[SYSTEMDS-3891] Add Generalized TSMMOOCPrimitive with Live Consumption

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

@@ -19,9 +19,6 @@
 
 package org.apache.sysds.runtime.instructions.ooc;
 
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-
 import org.apache.sysds.common.Opcodes;
 import org.apache.sysds.lops.MMTSJ;
 import org.apache.sysds.lops.MMTSJ.MMTSJType;
@@ -33,7 +30,6 @@
 import org.apache.sysds.runtime.instructions.InstructionUtils;
 import org.apache.sysds.runtime.instructions.cp.CPOperand;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
-import org.apache.sysds.runtime.matrix.data.LibMatrixReorg;
 import org.apache.sysds.runtime.matrix.data.MatrixBlock;
 import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
 import org.apache.sysds.runtime.matrix.operators.AggregateBinaryOperator;
@@ -67,45 +63,21 @@ public static TSMMOOCInstruction parseInstruction(String str) {
 	@Override
 	public void processInstruction(ExecutionContext ec) {
 		MatrixObject min = ec.getMatrixObject(input1);
+		if(!min.getDataCharacteristics().dimsKnown() || min.getBlocksize() <= 0)
+			throw new DMLRuntimeException("OOC TSMM requires known dimensions and a positive block size: "
+				+ min.getNumRows() + "x" + min.getNumColumns() + " (blocksize " + min.getBlocksize() + ")");
+
 		int numRowBlocks = Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks());
 		int numColBlocks = Math.toIntExact(min.getDataCharacteristics().getNumColBlocks());
 		if((_type.isLeft() && numColBlocks == 1) || (_type.isRight() && numRowBlocks == 1)) {
 			processSingleOutputTileInstruction(ec, min);
 			return;
 		}
 
-		int blocksPerJoinGroup = _type.isLeft() ? numColBlocks : numRowBlocks;
-		int partialsPerOutput = _type.isLeft() ? numRowBlocks : numColBlocks;
-
-		OOCStreamable<IndexedMatrixValue> inputStreamable = min.getStreamable();
-		final boolean createdCache = !inputStreamable.hasStreamCache();
-		final CachingStream inputCache = createdCache ? new CachingStream(min.getStreamHandle()) : inputStreamable
-			.getStreamCache();
-
-		OOCStream<List<IndexedMatrixValue>> groupedPartials = createWritableStream();
-		OOCStream<IndexedMatrixValue> partials = createWritableStream();
 		OOCStream<IndexedMatrixValue> out = createWritableStream();
-		addOutStream(out);
 		ec.getMatrixObject(output).setStreamHandle(out);
-
-		CompletableFuture<Void> joinFuture = joinManyOOC(inputCache.getReadStream(), inputCache.getReadStream(),
-			groupedPartials, this::createPartialOutputTiles, this::getJoinIndex, this::getJoinIndex, blocksPerJoinGroup,
-			blocksPerJoinGroup);
-		CompletableFuture<Void> expandFuture = expandOOC(groupedPartials, partials, values -> values);
-
 		BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
-		CompletableFuture<Void> outFuture = groupedReduceOOC(partials, out, (left, right) -> {
-			MatrixBlock result = ((MatrixBlock) left.getValue()).binaryOperations(plus, right.getValue());
-			left.setValue(result);
-			return left;
-		}, partialsPerOutput);
-
-		propagateFailuresToOutput(out, List.of(joinFuture, expandFuture, outFuture));
-
-		outFuture.whenComplete((result, error) -> {
-			if(createdCache)
-				inputCache.scheduleDeletion();
-		});
+		OOCInstructionUtils.tsmm(min.getStreamable(), out, _type, (AggregateBinaryOperator) _optr, plus, getContext());
 	}
 
 	private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObject min) {
@@ -119,52 +91,4 @@ private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObjec
 				((MatrixBlock) left.getValue()).binaryOperationsInPlace(plus, right.getValue())),
 			value -> ((MatrixBlock) value.getValue()).getExactSerializedSize(), getContext());
 	}
-
-	private long getJoinIndex(IndexedMatrixValue value) {
-		return _type.isLeft() ? value.getIndexes().getRowIndex() : value.getIndexes().getColumnIndex();
-	}
-
-	private long getOutputIndex(IndexedMatrixValue value) {
-		return _type.isLeft() ? value.getIndexes().getColumnIndex() : value.getIndexes().getRowIndex();
-	}
-
-	private List<IndexedMatrixValue> createPartialOutputTiles(IndexedMatrixValue left, IndexedMatrixValue right) {
-		long leftIndex = getOutputIndex(left);
-		long rightIndex = getOutputIndex(right);
-		if(leftIndex > rightIndex)
-			return List.of();
-
-		MatrixBlock leftBlock = (MatrixBlock) left.getValue();
-		MatrixBlock rightBlock = (MatrixBlock) right.getValue();
-		if(leftIndex == rightIndex) {
-			MatrixBlock diagonal = leftBlock.transposeSelfMatrixMultOperations(new MatrixBlock(), _type);
-			return List.of(new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), diagonal));
-		}
-
-		MatrixBlock partial = multiplyOffDiagonal(leftBlock, rightBlock);
-		MatrixBlock mirror = LibMatrixReorg.transpose(partial);
-		return List.of(new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), partial),
-			new IndexedMatrixValue(new MatrixIndexes(rightIndex, leftIndex), mirror));
-	}
-
-	private MatrixBlock multiplyOffDiagonal(MatrixBlock leftBlock, MatrixBlock rightBlock) {
-		if(_type.isLeft()) {
-			MatrixBlock leftTranspose = LibMatrixReorg.transpose(leftBlock);
-			return leftTranspose.aggregateBinaryOperations(leftTranspose, rightBlock, new MatrixBlock(),
-				(AggregateBinaryOperator) _optr);
-		}
-
-		MatrixBlock rightTranspose = LibMatrixReorg.transpose(rightBlock);
-		return leftBlock.aggregateBinaryOperations(leftBlock, rightTranspose, new MatrixBlock(),
-			(AggregateBinaryOperator) _optr);
-	}
-
-	private static void propagateFailuresToOutput(OOCStream<?> out, List<CompletableFuture<Void>> futures) {
-		for(CompletableFuture<Void> future : futures) {
-			future.exceptionally(error -> {
-				out.propagateFailure(DMLRuntimeException.of(error));
-				return null;
-			});
-		}
-	}
 }

src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java

@@ -62,7 +62,9 @@ private static void injectMaterializations(OOCPrimitive primitive, Set<OOCPrimit
 				boundaries.put(input, boundary);
 			}
 			primitive.discardInputHandle(request.inputIndex());
-			boundary.registerRequest(request.expectedReaders());
+			boolean live = boundary.registerRequest(request.expectedReaders(), request.liveConsumer());
+			if(request.liveRegistration() != null)
+				request.liveRegistration().accept(live);
 			primitive.installMaterializedInput(request.inputIndex(), boundary);
 		}
 		for(OOCPrimitive child : primitive.getChildren())

src/main/java/org/apache/sysds/runtime/ooc/planning/OOCStoreLayout.java

@@ -26,11 +26,18 @@ public enum OOCStoreLayout {
 	ROW_MAJOR, COL_MAJOR;
 
 	public int linearize(MatrixIndexes indexes, DataCharacteristics characteristics) {
+		return linearize(indexes.getRowIndex(), indexes.getColumnIndex(), characteristics);
+	}
+
+	public int linearize(long row, long col, DataCharacteristics characteristics) {
+		checkCharacteristics(characteristics);
+		long index = this == ROW_MAJOR ? (row - 1) * characteristics.getNumColBlocks() + col -
+			1 : (col - 1) * characteristics.getNumRowBlocks() + row - 1;
+		return Math.toIntExact(index);
+	}
+
+	private static void checkCharacteristics(DataCharacteristics characteristics) {
 		if(characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0)
 			throw new IllegalArgumentException("Materialized store layout requires known dimensions and block size.");
-		long index = this == ROW_MAJOR ? (indexes.getRowIndex() - 1) * characteristics.getNumColBlocks() +
-			indexes.getColumnIndex() -
-			1 : (indexes.getColumnIndex() - 1) * characteristics.getNumRowBlocks() + indexes.getRowIndex() - 1;
-		return Math.toIntExact(index);
 	}
 }

src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java

@@ -19,8 +19,11 @@
 
 package org.apache.sysds.runtime.ooc.primitives;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import java.util.function.ToIntFunction;
 
 import org.apache.sysds.runtime.instructions.ooc.CachingStream;
@@ -43,6 +46,8 @@ public final class MaterializeOOCPrimitive extends OOCPrimitive {
 	private final OOCFuture<MaterializedStore<IndexedMatrixValue>> _store;
 	private final AtomicBoolean _finished;
 	private final boolean _reusable;
+	private final List<Consumer<OOCStream.QueueCallback<IndexedMatrixValue>>> _liveConsumers;
+	private MaterializedStore<IndexedMatrixValue> _materializedStore;
 	private int _expectedReaders;
 	private int _consumers;
 
@@ -59,21 +64,32 @@ private MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> source, OOCSto
 		_store = new OOCFuture<>();
 		_finished = new AtomicBoolean();
 		_reusable = reusable;
+		_liveConsumers = new ArrayList<>();
 	}
 
 	public static MaterializeOOCPrimitive reusable(OOCStreamable<IndexedMatrixValue> source) {
-		return new MaterializeOOCPrimitive(source, OOCStoreLayout.ROW_MAJOR, null, true);
+		return reusable(source, OOCStoreLayout.ROW_MAJOR);
 	}
 
-	public synchronized void registerRequest(int expectedReaders) {
-		if(_reusable)
-			throw new IllegalStateException("Reusable materialization registers readers dynamically.");
+	public static MaterializeOOCPrimitive reusable(OOCStreamable<IndexedMatrixValue> source, OOCStoreLayout layout) {
+		return new MaterializeOOCPrimitive(source, layout, null, true);
+	}
+
+	public synchronized boolean registerRequest(int expectedReaders,
+		Consumer<OOCStream.QueueCallback<IndexedMatrixValue>> liveConsumer) {
 		if(expectedReaders <= 0)
 			throw new IllegalArgumentException("Materialization request requires at least one reader.");
-		if(hasStartedExecution())
-			throw new IllegalStateException("Cannot register a consumer after materialization started.");
-		_expectedReaders = Math.addExact(_expectedReaders, expectedReaders);
-		_consumers = Math.addExact(_consumers, 1);
+		boolean live = !hasStartedExecution();
+		if(_materializedStore == null) {
+			if(!_reusable)
+				_expectedReaders = Math.addExact(_expectedReaders, expectedReaders);
+			_consumers = Math.addExact(_consumers, 1);
+		}
+		else
+			_materializedStore.registerConsumer(expectedReaders);
+		if(live && liveConsumer != null)
+			_liveConsumers.add(liveConsumer);
+		return live;
 	}
 
 	public OOCFuture<MaterializedStore<IndexedMatrixValue>> store() {
@@ -99,19 +115,24 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) {
 	protected void startExecution() {
 		try {
 			OOCStream<IndexedMatrixValue> source = getInputReadStream(0);
-			MaterializedStore<IndexedMatrixValue> store = _reusable ? new MaterializedStore<>(
-				OOCCacheManager.getGlobalCache(),
-				CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(),
-					CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers);
 			DataCharacteristics characteristics = _source.getDataCharacteristics();
+			boolean logicalLayout = characteristics != null && characteristics.dimsKnown() &&
+				characteristics.getBlocksize() > 0;
+			ToIntFunction<MatrixIndexes> linearize = logicalLayout ? indexes -> _layout.linearize(indexes,
+				characteristics) : null;
+			MaterializedStore<IndexedMatrixValue> store;
+			synchronized(this) {
+				int consumers = _reusable ? 1 + _consumers : _consumers;
+				store = new MaterializedStore<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID(),
+					_reusable ? -1 : _expectedReaders, consumers, logicalLayout ? _layout : null,
+					logicalLayout ? characteristics : null);
+				_materializedStore = store;
+			}
 			AtomicInteger nextIndex = new AtomicInteger();
-			ToIntFunction<MatrixIndexes> linearize;
-			if(_reusable &&
-				(characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0))
-				linearize = ignored -> nextIndex.getAndIncrement();
-			else
-				linearize = indexes -> _layout.linearize(indexes, characteristics);
-			OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, linearize, _allowance);
+			ToIntFunction<MatrixIndexes> publicationIndex = linearize != null ? linearize : ignored -> nextIndex
+				.getAndIncrement();
+			OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, publicationIndex, _allowance,
+				_liveConsumers);
 			materializer.completion().whenComplete((ignored, error) -> {
 				if(error != null)
 					fail(error);

src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java

@@ -25,6 +25,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
@@ -229,6 +230,15 @@ private InputSlot(OOCStreamable<?> source) {
 		}
 	}
 
-	public record OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders) {
+	public record OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders,
+		Consumer<OOCStream.QueueCallback<IndexedMatrixValue>> liveConsumer, Consumer<Boolean> liveRegistration) {
+		public OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders) {
+			this(inputIndex, layout, expectedReaders, null, null);
+		}
+
+		public OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders,
+			Consumer<OOCStream.QueueCallback<IndexedMatrixValue>> liveConsumer) {
+			this(inputIndex, layout, expectedReaders, liveConsumer, null);
+		}
 	}
 }