src/main/java/org/apache/sysds/runtime/compress/lib/CLALibStack.java - systemds - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 package org.apache.sysds.runtime.compress.lib;

 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.stream.Collectors;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.compress.CompressedMatrixBlock;
 import org.apache.sysds.runtime.compress.DMLCompressionException;
 import org.apache.sysds.runtime.compress.colgroup.AColGroup;
 import org.apache.sysds.runtime.compress.colgroup.AColGroup.CompressionType;
 import org.apache.sysds.runtime.compress.colgroup.indexes.IIterate;
 import org.apache.sysds.runtime.matrix.data.MatrixBlock;
 import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
 import org.apache.sysds.runtime.util.CommonThreadPool;

 public final class CLALibStack {
 	protected static final Log LOG = LogFactory.getLog(CLALibStack.class.getName());

 	private CLALibStack() {
 		// private constructor
 	}

 	/**
 	 * Combine the map of index matrix blocks into a single MatrixBlock.
 	 *
 	 * The intension is that the combining is able to resolve differences in the different MatrixBlocks allocation.
 	 *
 	 * @param m The map of Index to MatrixBLocks
 	 * @param k The parallelization degree allowed for this operation
 	 * @return The combined matrix.
 	 */
 	public static MatrixBlock combine(Map<MatrixIndexes, MatrixBlock> m, int k) {
 		final MatrixIndexes lookup = new MatrixIndexes(1, 1);
 		MatrixBlock b = m.get(lookup);
 		if(b == null)
 			throw new DMLCompressionException("Invalid map to combine does not contain the top left map MatrixBlock");
 		final int blen = Math.max(b.getNumColumns(), b.getNumRows());

 		// Dynamically find rlen, clen and blen;
 		final long rows = findRLength(m, b);
 		// TODO utilize the known size of m to extrapolate how many row blocks there are
 		// and only use the last rowblock to calculate the total number of rows.
 		final long cols = findCLength(m, b);

 		return combine(m, lookup, (int) rows, (int) cols, blen, k);
 	}

 	public static MatrixBlock combine(Map<MatrixIndexes, MatrixBlock> m, int rlen, int clen, int blen, int k) {
 		final MatrixIndexes lookup = new MatrixIndexes();
 		return combine(m, lookup, rlen, clen, blen, k);
 	}

 	private static long findRLength(Map<MatrixIndexes, MatrixBlock> m, MatrixBlock b) {
 		final MatrixIndexes lookup = new MatrixIndexes(1, 1);
 		long rows = 0;
 		while((b = m.get(lookup)) != null) {
 			rows += b.getNumRows();
 			lookup.setIndexes(lookup.getRowIndex() + 1, 1);
 		}
 		return rows;
 	}

 	private static long findCLength(Map<MatrixIndexes, MatrixBlock> m, MatrixBlock b) {
 		final MatrixIndexes lookup = new MatrixIndexes(1, 1);
 		long cols = 0;
 		while((b = m.get(lookup)) != null) {
 			cols += b.getNumColumns();
 			lookup.setIndexes(1, lookup.getColumnIndex() + 1);
 		}
 		return cols;
 	}

 	private static MatrixBlock combine(final Map<MatrixIndexes, MatrixBlock> m, final MatrixIndexes lookup,
 		final int rlen, final int clen, final int blen, final int k) {

 		if(rlen < blen) // Shortcut, in case file only contains one block in r length.
 			return combineColumnGroups(m, lookup, rlen, clen, blen, k);

 		final CompressionType[] colTypes = new CompressionType[clen];
 		// Look through the first blocks in to the top.
 		for(int bc = 0; bc * blen < clen; bc++) {
 			lookup.setIndexes(1, bc + 1); // get first blocks
 			final MatrixBlock b = m.get(lookup);
 			if(!(b instanceof CompressedMatrixBlock)) {
 				LOG.warn("Found uncompressed matrix in Map of matrices, this is not"
 					+ " supported in combine therefore falling back to decompression");
 				return combineViaDecompression(m, rlen, clen, blen, k);
 			}
 			final CompressedMatrixBlock cmb = (CompressedMatrixBlock) b;
 			if(cmb.isOverlapping()) {
 				LOG.warn("Not supporting overlapping combine yet falling back to decompression");
 				return combineViaDecompression(m, rlen, clen, blen, k);
 			}
 			final List<AColGroup> gs = cmb.getColGroups();
 			final int off = bc * blen;
 			for(AColGroup g : gs) {
 				try {
 					final IIterate cols = g.getColIndices().iterator();
 					final CompressionType t = g.getCompType();
 					while(cols.hasNext())
 						colTypes[cols.next() + off] = t;
 				}
 				catch(Exception e) {
 					throw new DMLCompressionException("Failed combining: " + g.toString());
 				}
 			}
 		}

 		// Look through the Remaining blocks down in the rows.
 		for(int br = 1; br * blen < rlen; br++) {
 			for(int bc = 0; bc * blen < clen; bc++) {
 				lookup.setIndexes(br + 1, bc + 1); // get first blocks
 				final MatrixBlock b = m.get(lookup);
 				if(!(b instanceof CompressedMatrixBlock)) {
 					LOG.warn("Found uncompressed matrix in Map of matrices, this is not"
 						+ " supported in combine therefore falling back to decompression");
 					return combineViaDecompression(m, rlen, clen, blen, k);
 				}
 				final CompressedMatrixBlock cmb = (CompressedMatrixBlock) b;
 				if(cmb.isOverlapping()) {
 					LOG.warn("Not supporting overlapping combine yet falling back to decompression");
 					return combineViaDecompression(m, rlen, clen, blen, k);
 				}
 				final List<AColGroup> gs = cmb.getColGroups();
 				final int off = bc * blen;
 				for(AColGroup g : gs) {
 					final IIterate cols = g.getColIndices().iterator();
 					final CompressionType t = g.getCompType();
 					while(cols.hasNext()) {
 						final int c = cols.next();
 						if(colTypes[c + off] != t) {
 							LOG.warn("Not supported different types of column groups to combine."
 								+ "Falling back to decompression of all blocks");
 							return combineViaDecompression(m, rlen, clen, blen, k);
 						}
 					}
 				}
 			}
 		}

 		return combineColumnGroups(m, lookup, rlen, clen, blen, k);
 	}

 	private static MatrixBlock combineViaDecompression(final Map<MatrixIndexes, MatrixBlock> m, final int rlen,
 		final int clen, final int blen, int k) {
 		final MatrixBlock out = new MatrixBlock(rlen, clen, false);
 		out.allocateDenseBlock();
 		for(Entry<MatrixIndexes, MatrixBlock> e : m.entrySet()) {
 			final MatrixIndexes ix = e.getKey();
 			final MatrixBlock block = e.getValue();
 			if(block != null) {
 				final int row_offset = (int) (ix.getRowIndex() - 1) * blen;
 				final int col_offset = (int) (ix.getColumnIndex() - 1) * blen;
 				block.putInto(out, row_offset, col_offset, false);
 			}
 		}
 		out.setNonZeros(-1);
 		out.examSparsity(true);
 		return out;
 	}

 	// It is known all of the matrices are Compressed and they are non overlapping.
 	private static MatrixBlock combineColumnGroups(final Map<MatrixIndexes, MatrixBlock> m, final MatrixIndexes lookup,
 		final int rlen, final int clen, final int blen, int k) {

 		final AColGroup[][] finalCols = new AColGroup[clen][]; // temp array for combining
 		final int blocksInColumn = (rlen - 1) / blen + 1;


 		// Add all the blocks into linear structure.
 		for(int br = 0; br * blen < rlen; br++) {
 			for(int bc = 0; bc * blen < clen; bc++) {
 				lookup.setIndexes(br + 1, bc + 1);
 				final CompressedMatrixBlock cmb = (CompressedMatrixBlock) m.get(lookup);
 				for(AColGroup g : cmb.getColGroups()) {
 					final AColGroup gc = bc > 0 ? g.shiftColIndices(bc * blen) : g;
 					final int c = gc.getColIndices().get(0);
 					if(br == 0)
 						finalCols[c] = new AColGroup[blocksInColumn];
 					else if(finalCols[c] == null) {
 						LOG.warn("Combining via decompression. There was an column"
 							+ " assigned not assigned in block 1 indicating spark compression");
 						return combineViaDecompression(m, rlen, clen, blen, k);
 					}
 					finalCols[c][br] = gc;
 					if(br != 0 && (finalCols[c][0] == null || !finalCols[c][br].getColIndices().equals(finalCols[c][0].getColIndices()))){
 						LOG.warn("Combining via decompression. There was an column with different index");
 						return combineViaDecompression(m, rlen, clen, blen, k);
 					}

 				}
 			}
 		}


 		final ExecutorService pool = CommonThreadPool.get(Math.max(Math.min(clen / 500, k), 1));
 		try {

 			List<AColGroup> finalGroups = pool.submit(() -> {
 				return Arrays//
 					.stream(finalCols)//
 					.filter(x -> x != null)// filter all columns that are contained in other groups.
 					.parallel()//
 					.map(x -> {
 						return combineN(x);
 					}).collect(Collectors.toList());
 			}).get();
 			if(finalGroups.contains(null)) {
 				LOG.warn("Combining via decompression. There was a column group that did not append ");
 				return combineViaDecompression(m, rlen, clen, blen, k);
 			}
 			return new CompressedMatrixBlock(rlen, clen, -1, false, finalGroups);
 		}
 		catch(InterruptedException | ExecutionException e) {
 			throw new DMLRuntimeException("Failed to combine column groups", e);
 		}
 		finally {
 			pool.shutdown();
 		}
 	}

 	private static AColGroup combineN(AColGroup[] groups) {
 		try {
 			return AColGroup.appendN(groups);

 		}
 		catch(Exception e) {
 			throw new DMLCompressionException("Failed to combine groups:\n" + Arrays.toString(groups), e);
 		}
 	}
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	package org.apache.sysds.runtime.compress.lib;

	import java.util.Arrays;
	import java.util.List;
	import java.util.Map;
	import java.util.Map.Entry;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.ExecutorService;
	import java.util.stream.Collectors;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.sysds.runtime.DMLRuntimeException;
	import org.apache.sysds.runtime.compress.CompressedMatrixBlock;
	import org.apache.sysds.runtime.compress.DMLCompressionException;
	import org.apache.sysds.runtime.compress.colgroup.AColGroup;
	import org.apache.sysds.runtime.compress.colgroup.AColGroup.CompressionType;
	import org.apache.sysds.runtime.compress.colgroup.indexes.IIterate;
	import org.apache.sysds.runtime.matrix.data.MatrixBlock;
	import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
	import org.apache.sysds.runtime.util.CommonThreadPool;

	public final class CLALibStack {
	protected static final Log LOG = LogFactory.getLog(CLALibStack.class.getName());

	private CLALibStack() {
	// private constructor
	}

	/**
	* Combine the map of index matrix blocks into a single MatrixBlock.
	*
	* The intension is that the combining is able to resolve differences in the different MatrixBlocks allocation.
	*
	* @param m The map of Index to MatrixBLocks
	* @param k The parallelization degree allowed for this operation
	* @return The combined matrix.
	*/
	public static MatrixBlock combine(Map<MatrixIndexes, MatrixBlock> m, int k) {
	final MatrixIndexes lookup = new MatrixIndexes(1, 1);
	MatrixBlock b = m.get(lookup);
	if(b == null)
	throw new DMLCompressionException("Invalid map to combine does not contain the top left map MatrixBlock");
	final int blen = Math.max(b.getNumColumns(), b.getNumRows());

	// Dynamically find rlen, clen and blen;
	final long rows = findRLength(m, b);
	// TODO utilize the known size of m to extrapolate how many row blocks there are
	// and only use the last rowblock to calculate the total number of rows.
	final long cols = findCLength(m, b);

	return combine(m, lookup, (int) rows, (int) cols, blen, k);
	}

	public static MatrixBlock combine(Map<MatrixIndexes, MatrixBlock> m, int rlen, int clen, int blen, int k) {
	final MatrixIndexes lookup = new MatrixIndexes();
	return combine(m, lookup, rlen, clen, blen, k);
	}

	private static long findRLength(Map<MatrixIndexes, MatrixBlock> m, MatrixBlock b) {
	final MatrixIndexes lookup = new MatrixIndexes(1, 1);
	long rows = 0;
	while((b = m.get(lookup)) != null) {
	rows += b.getNumRows();
	lookup.setIndexes(lookup.getRowIndex() + 1, 1);
	}
	return rows;
	}

	private static long findCLength(Map<MatrixIndexes, MatrixBlock> m, MatrixBlock b) {
	final MatrixIndexes lookup = new MatrixIndexes(1, 1);
	long cols = 0;
	while((b = m.get(lookup)) != null) {
	cols += b.getNumColumns();
	lookup.setIndexes(1, lookup.getColumnIndex() + 1);
	}
	return cols;
	}

	private static MatrixBlock combine(final Map<MatrixIndexes, MatrixBlock> m, final MatrixIndexes lookup,
	final int rlen, final int clen, final int blen, final int k) {

	if(rlen < blen) // Shortcut, in case file only contains one block in r length.
	return combineColumnGroups(m, lookup, rlen, clen, blen, k);

	final CompressionType[] colTypes = new CompressionType[clen];
	// Look through the first blocks in to the top.
	for(int bc = 0; bc * blen < clen; bc++) {
	lookup.setIndexes(1, bc + 1); // get first blocks
	final MatrixBlock b = m.get(lookup);
	if(!(b instanceof CompressedMatrixBlock)) {
	LOG.warn("Found uncompressed matrix in Map of matrices, this is not"
	+ " supported in combine therefore falling back to decompression");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	final CompressedMatrixBlock cmb = (CompressedMatrixBlock) b;
	if(cmb.isOverlapping()) {
	LOG.warn("Not supporting overlapping combine yet falling back to decompression");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	final List<AColGroup> gs = cmb.getColGroups();
	final int off = bc * blen;
	for(AColGroup g : gs) {
	try {
	final IIterate cols = g.getColIndices().iterator();
	final CompressionType t = g.getCompType();
	while(cols.hasNext())
	colTypes[cols.next() + off] = t;
	}
	catch(Exception e) {
	throw new DMLCompressionException("Failed combining: " + g.toString());
	}
	}
	}

	// Look through the Remaining blocks down in the rows.
	for(int br = 1; br * blen < rlen; br++) {
	for(int bc = 0; bc * blen < clen; bc++) {
	lookup.setIndexes(br + 1, bc + 1); // get first blocks
	final MatrixBlock b = m.get(lookup);
	if(!(b instanceof CompressedMatrixBlock)) {
	LOG.warn("Found uncompressed matrix in Map of matrices, this is not"
	+ " supported in combine therefore falling back to decompression");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	final CompressedMatrixBlock cmb = (CompressedMatrixBlock) b;
	if(cmb.isOverlapping()) {
	LOG.warn("Not supporting overlapping combine yet falling back to decompression");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	final List<AColGroup> gs = cmb.getColGroups();
	final int off = bc * blen;
	for(AColGroup g : gs) {
	final IIterate cols = g.getColIndices().iterator();
	final CompressionType t = g.getCompType();
	while(cols.hasNext()) {
	final int c = cols.next();
	if(colTypes[c + off] != t) {
	LOG.warn("Not supported different types of column groups to combine."
	+ "Falling back to decompression of all blocks");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	}
	}
	}
	}

	return combineColumnGroups(m, lookup, rlen, clen, blen, k);
	}

	private static MatrixBlock combineViaDecompression(final Map<MatrixIndexes, MatrixBlock> m, final int rlen,
	final int clen, final int blen, int k) {
	final MatrixBlock out = new MatrixBlock(rlen, clen, false);
	out.allocateDenseBlock();
	for(Entry<MatrixIndexes, MatrixBlock> e : m.entrySet()) {
	final MatrixIndexes ix = e.getKey();
	final MatrixBlock block = e.getValue();
	if(block != null) {
	final int row_offset = (int) (ix.getRowIndex() - 1) * blen;
	final int col_offset = (int) (ix.getColumnIndex() - 1) * blen;
	block.putInto(out, row_offset, col_offset, false);
	}
	}
	out.setNonZeros(-1);
	out.examSparsity(true);
	return out;
	}

	// It is known all of the matrices are Compressed and they are non overlapping.
	private static MatrixBlock combineColumnGroups(final Map<MatrixIndexes, MatrixBlock> m, final MatrixIndexes lookup,
	final int rlen, final int clen, final int blen, int k) {

	final AColGroup[][] finalCols = new AColGroup[clen][]; // temp array for combining
	final int blocksInColumn = (rlen - 1) / blen + 1;


	// Add all the blocks into linear structure.
	for(int br = 0; br * blen < rlen; br++) {
	for(int bc = 0; bc * blen < clen; bc++) {
	lookup.setIndexes(br + 1, bc + 1);
	final CompressedMatrixBlock cmb = (CompressedMatrixBlock) m.get(lookup);
	for(AColGroup g : cmb.getColGroups()) {
	final AColGroup gc = bc > 0 ? g.shiftColIndices(bc * blen) : g;
	final int c = gc.getColIndices().get(0);
	if(br == 0)
	finalCols[c] = new AColGroup[blocksInColumn];
	else if(finalCols[c] == null) {
	LOG.warn("Combining via decompression. There was an column"
	+ " assigned not assigned in block 1 indicating spark compression");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	finalCols[c][br] = gc;
	if(br != 0 && (finalCols[c][0] == null \|\| !finalCols[c][br].getColIndices().equals(finalCols[c][0].getColIndices()))){
	LOG.warn("Combining via decompression. There was an column with different index");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}

	}
	}
	}


	final ExecutorService pool = CommonThreadPool.get(Math.max(Math.min(clen / 500, k), 1));
	try {

	List<AColGroup> finalGroups = pool.submit(() -> {
	return Arrays//
	.stream(finalCols)//
	.filter(x -> x != null)// filter all columns that are contained in other groups.
	.parallel()//
	.map(x -> {
	return combineN(x);
	}).collect(Collectors.toList());
	}).get();
	if(finalGroups.contains(null)) {
	LOG.warn("Combining via decompression. There was a column group that did not append ");
	return combineViaDecompression(m, rlen, clen, blen, k);
	}
	return new CompressedMatrixBlock(rlen, clen, -1, false, finalGroups);
	}
	catch(InterruptedException \| ExecutionException e) {
	throw new DMLRuntimeException("Failed to combine column groups", e);
	}
	finally {
	pool.shutdown();
	}
	}

	private static AColGroup combineN(AColGroup[] groups) {
	try {
	return AColGroup.appendN(groups);

	}
	catch(Exception e) {
	throw new DMLCompressionException("Failed to combine groups:\n" + Arrays.toString(groups), e);
	}
	}
	}