[SYSTEMDS-3924] New primitive for OOC stream creation / reset

This patch improves the OOC backend by a new primitive that
automatically creates or resets existing OOC streams on
getStreamHandle such that OOC instruction don't need to handle this
issue individually, but can still probe the existence of active streams
via hasStreamHandle. As a result if there are OOC instructions that
consume materialized intermediates, we automatically create new
streams from these intermediates.

Author: mboehm7

src/main/java/org/apache/sysds/hops/rewrite/RewriteInjectOOCTee.java

@@ -51,6 +51,8 @@
  */
 public class RewriteInjectOOCTee extends HopRewriteRule {
 
+	public static boolean APPLY_ONLY_XtX_PATTERN = false;
+	
 	private static final Set<Long> rewrittenHops = new HashSet<>();
 	private static final Map<Long, Hop> handledHop = new HashMap<>();
 
@@ -140,7 +142,7 @@ private void findRewriteCandidates(Hop hop) {
 			&& hop.getDataType().isMatrix()
 			&& !HopRewriteUtils.isData(hop, OpOpData.TEE)
 			&& hop.getParent().size() > 1
-			&& isSelfTranposePattern(hop)) //FIXME remove
+			&& (!APPLY_ONLY_XtX_PATTERN || isSelfTranposePattern(hop))) //FIXME remove
 		{
 			rewriteCandidates.add(hop);
 		}

src/main/java/org/apache/sysds/runtime/controlprogram/caching/CacheableData.java

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.LongStream;
 
 import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.apache.commons.logging.Log;
@@ -48,19 +49,22 @@
 import org.apache.sysds.runtime.instructions.fed.InitFEDInstruction;
 import org.apache.sysds.runtime.instructions.gpu.context.GPUContext;
 import org.apache.sysds.runtime.instructions.gpu.context.GPUObject;
+import org.apache.sysds.runtime.instructions.ooc.ResettableStream;
 import org.apache.sysds.runtime.instructions.spark.data.BroadcastObject;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
 import org.apache.sysds.runtime.instructions.spark.data.RDDObject;
 import org.apache.sysds.runtime.io.FileFormatProperties;
 import org.apache.sysds.runtime.io.IOUtilFunctions;
 import org.apache.sysds.runtime.io.ReaderWriterFederated;
 import org.apache.sysds.runtime.lineage.LineageItem;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
 import org.apache.sysds.runtime.meta.MatrixCharacteristics;
 import org.apache.sysds.runtime.meta.MetaData;
 import org.apache.sysds.runtime.meta.MetaDataFormat;
 import org.apache.sysds.runtime.util.HDFSTool;
 import org.apache.sysds.runtime.util.LocalFileUtils;
+import org.apache.sysds.runtime.util.UtilFunctions;
 import org.apache.sysds.utils.Statistics;
 import org.apache.sysds.utils.stats.InfrastructureAnalyzer;
 
@@ -466,8 +470,44 @@ public boolean hasBroadcastHandle() {
 	}
 
 	public LocalTaskQueue<IndexedMatrixValue> getStreamHandle() {
+		if( !hasStreamHandle() ) {
+			_streamHandle = new LocalTaskQueue<>();
+			DataCharacteristics dc = getDataCharacteristics();
+			MatrixBlock src = (MatrixBlock)acquireReadAndRelease();
+			LongStream.range(0, dc.getNumBlocks())
+				.mapToObj(i -> UtilFunctions.createIndexedMatrixBlock(src, dc, i))
+				.forEach( blk -> {
+					try{ 
+						_streamHandle.enqueueTask(blk); 
+					}
+					catch(Exception ex) {
+						throw new DMLRuntimeException(ex);
+				}});
+			_streamHandle.closeInput();
+		}
+		else if(_streamHandle != null && _streamHandle.isProcessed() 
+			&& _streamHandle instanceof ResettableStream) 
+		{
+			try {
+				((ResettableStream)_streamHandle).reset();
+			}
+			catch(Exception ex) {
+				throw new DMLRuntimeException(ex);
+			}
+		}
+		
 		return _streamHandle;
 	}
+	
+	/**
+	 * Probes if stream handle is existing, because <code>getStreamHandle<code>
+	 * creates a new stream if not existing.
+	 * 
+	 * @return true if existing, false otherwise
+	 */
+	public boolean hasStreamHandle() {
+		return _streamHandle != null && !_streamHandle.isProcessed();
+	} 
 
 	@SuppressWarnings({ "rawtypes", "unchecked" })
 	public void setBroadcastHandle( BroadcastObject bc ) {
@@ -592,7 +632,7 @@ && getRDDHandle() == null) ) {
 					//mark for initial local write despite read operation
 					_requiresLocalWrite = false;
 				}
-				else if( getStreamHandle() != null ) {
+				else if( hasStreamHandle() ) {
 					_data = readBlobFromStream( getStreamHandle() );
 				}
 				else if( getRDDHandle()==null || getRDDHandle().allowsShortCircuitRead() ) {
@@ -909,7 +949,7 @@ public synchronized void exportData (String fName, String outputFormat, int repl
 			// a) get the matrix
 			boolean federatedWrite = (outputFormat != null ) &&  outputFormat.contains("federated");
 
-			if(getStreamHandle()!=null) {
+			if(hasStreamHandle()) {
 				try {
 					long totalNnz = writeStreamToHDFS(fName, outputFormat, replication, formatProperties);
 					updateDataCharacteristics(new MatrixCharacteristics(

src/main/java/org/apache/sysds/runtime/controlprogram/parfor/LocalTaskQueue.java

@@ -106,6 +106,10 @@ public synchronized void closeInput()
 		_closedInput = true;
 		notifyAll(); //notify all waiting readers
 	}
+	
+	public synchronized boolean isProcessed() {
+		return _closedInput && _data.isEmpty();
+	}
 
 	@Override
 	public synchronized String toString() 

src/main/java/org/apache/sysds/runtime/instructions/ooc/ResettableStream.java

@@ -83,7 +83,7 @@ public synchronized IndexedMatrixValue dequeueTask()
 	 * This can only be called once the stream is fully consumed once.
 	 */
 	public synchronized void reset() throws InterruptedException {
-		if (_cacheInProgress) {
+		while (_cacheInProgress) {
 			// Attempted to reset a stream that's not been fully cached yet.
 			wait();
 		}
@@ -94,4 +94,9 @@ public synchronized void reset() throws InterruptedException {
 	public synchronized void closeInput() {
 		_source.closeInput();
 	}
+	
+	@Override
+	public synchronized boolean isProcessed() {
+		return false;
+	}
 }

src/main/java/org/apache/sysds/runtime/util/UtilFunctions.java

@@ -53,8 +53,10 @@
 import org.apache.sysds.runtime.frame.data.columns.HashIntegerArray;
 import org.apache.sysds.runtime.frame.data.columns.HashLongArray;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
 import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
 import org.apache.sysds.runtime.matrix.data.Pair;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
 import org.apache.sysds.runtime.meta.TensorCharacteristics;
 import org.apache.sysds.runtime.transform.encode.ColumnEncoderRecode;
 
@@ -1471,4 +1473,26 @@ public static String[] cleanAndTokenizeRow(String[] row) {
 
 		return joined.split("\\s+");
 	}
+	
+	public static IndexedMatrixValue createIndexedMatrixBlock(MatrixBlock mb, DataCharacteristics mc, long ix) {
+		try {
+			//compute block indexes
+			long blockRow = ix / mc.getNumColBlocks();
+			long blockCol = ix % mc.getNumColBlocks();
+			//compute block sizes
+			int maxRow = UtilFunctions.computeBlockSize(mc.getRows(), blockRow+1, mc.getBlocksize());
+			int maxCol = UtilFunctions.computeBlockSize(mc.getCols(), blockCol+1, mc.getBlocksize());
+			//copy sub-matrix to block
+			MatrixBlock block = new MatrixBlock(maxRow, maxCol, mb.isInSparseFormat());
+			int row_offset = (int)blockRow*mc.getBlocksize();
+			int col_offset = (int)blockCol*mc.getBlocksize();
+			block = mb.slice( row_offset, row_offset+maxRow-1,
+				col_offset, col_offset+maxCol-1, false, block );
+			//create key-value pair
+			return new IndexedMatrixValue(new MatrixIndexes(blockRow+1, blockCol+1), block);
+		}
+		catch(DMLRuntimeException ex) {
+			throw new RuntimeException(ex);
+		}
+	}
 }

src/test/java/org/apache/sysds/test/functions/ooc/lmDSTest.java

@@ -20,6 +20,7 @@
 package org.apache.sysds.test.functions.ooc;
 
 import org.apache.sysds.common.Types;
+import org.apache.sysds.hops.rewrite.RewriteInjectOOCTee;
 import org.apache.sysds.runtime.io.MatrixWriter;
 import org.apache.sysds.runtime.io.MatrixWriterFactory;
 import org.apache.sysds.runtime.matrix.data.LibCommonsMath;
@@ -70,9 +71,12 @@ public void testlmDS2() {
 	private void runMatrixVectorMultiplicationTest(int cols)
 	{
 		Types.ExecMode platformOld = setExecMode(Types.ExecMode.SINGLE_NODE);
-
+		boolean oldFlag = RewriteInjectOOCTee.APPLY_ONLY_XtX_PATTERN;
+		
 		try
 		{
+			RewriteInjectOOCTee.APPLY_ONLY_XtX_PATTERN = true;
+			
 			getAndLoadTestConfiguration(TEST_NAME1);
 			String HOME = SCRIPT_DIR + TEST_DIR;
 			fullDMLScriptName = HOME + TEST_NAME1 + ".dml";
@@ -117,6 +121,7 @@ private void runMatrixVectorMultiplicationTest(int cols)
 		}
 		finally {
 			resetExecMode(platformOld);
+			RewriteInjectOOCTee.APPLY_ONLY_XtX_PATTERN = oldFlag;
 		}
 	}
 }

src/test/scripts/functions/ooc/lmDS.dml

@@ -25,4 +25,5 @@ y = read($2)
 XtX = t(X) %*% X; # 500 x 500
 Xty = t(X) %*% y; # 500 x 1
 R = solve(XtX, Xty)
+print(sum(R!=0))
 write(R, $3, format="binary")