[SYSTEMDS-3891] Wire OOC Materialized Stores

src/main/java/org/apache/sysds/runtime/controlprogram/ParForProgramBlock.java

@@ -91,6 +91,7 @@
 import org.apache.sysds.runtime.instructions.cp.ScalarObject;
 import org.apache.sysds.runtime.instructions.cp.StringObject;
 import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction;
+import org.apache.sysds.runtime.instructions.ooc.TeeOOCInstruction;
 import org.apache.sysds.runtime.lineage.Lineage;
 import org.apache.sysds.runtime.lineage.LineageCacheConfig;
 import org.apache.sysds.runtime.lineage.LineageItem;
@@ -730,6 +731,7 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject
 		final LocalTaskQueue<Task> queue = new LocalTaskQueue<>();
 		final Thread[] threads         = new Thread[_numThreads];
 		final LocalParWorker[] workers = new LocalParWorker[_numThreads];
+		final Set<String> resultVarNames = _resultVars.stream().map(v -> v._name).collect(Collectors.toSet());
 		@SuppressWarnings("unchecked")
 		final HashMap<String, Data>[] workerBaselines = DMLScript.USE_OOC ? new HashMap[_numThreads] : null;
 		try
@@ -740,8 +742,11 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject
 				workers[i] = createParallelWorker( _pwIDs[i], queue, ec, i);
 				if(DMLScript.USE_OOC) {
 					workerBaselines[i] = new HashMap<>();
-					for(Map.Entry<String, Data> e : workers[i].getVariables().entrySet())
+					for(Map.Entry<String, Data> e : workers[i].getVariables().entrySet()) {
 						workerBaselines[i].put(e.getKey(), e.getValue());
+						if(!resultVarNames.contains(e.getKey()) && e.getValue() instanceof MatrixObject matrix)
+							TeeOOCInstruction.incrRef(matrix.getStreamable(), 1);
+					}
 				}
 				threads[i] = new Thread( workers[i] , "PARFOR");
 				threads[i].setPriority(Thread.MAX_PRIORITY);
@@ -785,8 +790,6 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject
 
 			// Step 4) collecting results from each parallel worker
 			//obtain results and cleanup other intermediates before result merge
-			Set<String> resultVarNames = _resultVars.stream()
-				.map(v -> v._name).collect(Collectors.toSet());
 			LocalVariableMap [] localVariables = new LocalVariableMap [_numThreads];
 			for( int i=0; i<_numThreads; i++ ) {
 				localVariables[i] = workers[i].getVariables();
@@ -796,6 +799,8 @@ private void executeLocalParFor( ExecutionContext ec, IntObject from, IntObject
 							Data current = localVariables[i].get(var);
 							if(current != null && current != workerBaselines[i].get(var))
 								VariableCPInstruction.processRmvarInstruction(workers[i].getExecutionContext(), var);
+							else if(current instanceof MatrixObject matrix)
+								TeeOOCInstruction.incrRef(matrix.getStreamable(), -1);
 						}
 					}
 				}

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

@@ -491,7 +491,7 @@ public synchronized OOCStream<IndexedMatrixValue> getStreamHandle() {
 		}
 
 		OOCStream<IndexedMatrixValue> stream = _streamHandle.getReadStream();
-		if(!stream.hasStreamCache())
+		if(!_streamHandle.hasStreamCache() && !_streamHandle.hasMaterializedStore())
 			_streamHandle = null; // To ensure read once
 		return stream;
 	}

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

@@ -32,6 +32,13 @@ public interface OOCStreamable<T> {
 
 	CachingStream getStreamCache();
 
+	default boolean hasMaterializedStore() {
+		return false;
+	}
+
+	default void scheduleMaterializedStoreDeletion() {
+	}
+
 	boolean isProcessed();
 
 	DataCharacteristics getDataCharacteristics();

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

@@ -24,22 +24,24 @@
 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.ooc.store.MaterializedStoreStreamable;
 
 import java.util.concurrent.ConcurrentHashMap;
 
 public class TeeOOCInstruction extends ComputationOOCInstruction {
 
-	private static final ConcurrentHashMap<CachingStream, Integer> refCtr = new ConcurrentHashMap<>();
+	private static final ConcurrentHashMap<OOCStreamable<IndexedMatrixValue>, Integer> refCtr = new ConcurrentHashMap<>();
 
 	public static void reset() {
 		if (!refCtr.isEmpty()) {
 			System.err.println("There are some dangling streams still in the cache: " + refCtr);
-			for(CachingStream cache : refCtr.keySet()) {
+			for(OOCStreamable<IndexedMatrixValue> stream : refCtr.keySet()) {
 				try {
-					cache.scheduleDeletion();
+					scheduleDeletion(stream);
 				}
 				catch(Exception ex) {
-					System.err.println("Failed to schedule deletion for dangling stream " + cache + ": " + ex.getMessage());
+					System.err
+						.println("Failed to schedule deletion for dangling stream " + stream + ": " + ex.getMessage());
 				}
 			}
 			refCtr.clear();
@@ -50,19 +52,26 @@ public static void reset() {
 	 * Increments the reference counter of a stream by the set amount.
 	 */
 	public static void incrRef(OOCStreamable<IndexedMatrixValue> stream, int incr) {
-		if (!stream.hasStreamCache())
+		if(!stream.hasStreamCache() && !stream.hasMaterializedStore())
 			return;
-		CachingStream cache = stream.getStreamCache();
+		OOCStreamable<IndexedMatrixValue> handle = stream.hasStreamCache() ? stream.getStreamCache() : stream;
 
-		Integer ref = refCtr.compute(cache, (k, v) -> {
+		Integer ref = refCtr.compute(handle, (k, v) -> {
 			if (v == null)
 				v = 0;
 			v += incr;
 			return v <= 0 ? null : v;
 		});
 
-		if (ref == null)
-			cache.scheduleDeletion();
+		if(ref == null)
+			scheduleDeletion(handle);
+	}
+
+	private static void scheduleDeletion(OOCStreamable<IndexedMatrixValue> stream) {
+		if(stream.hasMaterializedStore())
+			stream.scheduleMaterializedStoreDeletion();
+		else
+			stream.getStreamCache().scheduleDeletion();
 	}
 
 	protected TeeOOCInstruction(OOCType type, CPOperand in1, CPOperand out, String opcode, String istr) {
@@ -82,15 +91,15 @@ public void processInstruction(ExecutionContext ec) {
 		//get input stream
 		MatrixObject min = ec.getMatrixObject(input1);
 		OOCStreamable<IndexedMatrixValue> streamable = min.getStreamable();
-		CachingStream handle;
+		OOCStreamable<IndexedMatrixValue> handle;
 
-		if(streamable.hasStreamCache()) {
-			handle = streamable.getStreamCache();
+		if(streamable.hasStreamCache() || streamable.hasMaterializedStore()) {
+			handle = streamable.hasStreamCache() ? streamable.getStreamCache() : streamable;
 			incrRef(handle, 1);
 		}
 		else {
-			// We also set the input stream handle
-			handle = new CachingStream(min.getStreamHandle());
+			// The input and output matrix objects both retain the new reusable handle.
+			handle = new MaterializedStoreStreamable(min.getStreamHandle(), min);
 			min.setStreamHandle(handle);
 			incrRef(handle, 2);
 		}

src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java

@@ -43,6 +43,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.LockSupport;
 import java.util.function.LongUnaryOperator;
+import java.util.function.Supplier;
 
 public final class OOCPackedCache implements OOCCache {
 	private static final long PACKED_STREAM_ID = CachingStream._streamSeq.getNextID();
@@ -231,11 +232,8 @@ public OOCFuture<BlockEntry> pin(long sId, long tId, MemoryAllowance allowance)
 		if(!(location instanceof SealedPackLocation packed))
 			return _physical.pin(sId, tId, allowance);
 
-		return packed.state().pin(_physical, allowance, false).map(physicalEntry -> {
-			if(physicalEntry == null)
-				return null;
-			return createLogicalPin(new BlockKey(sId, tId), packed);
-		});
+		BlockKey key = new BlockKey(sId, tId);
+		return pinLogical(key, packed, () -> packed.state().pin(_physical, allowance, false));
 	}
 
 	@Override
@@ -248,11 +246,8 @@ public OOCFuture<BlockEntry> pinAdmitted(long sId, long tId, MemoryAllowance all
 		if(!(location instanceof SealedPackLocation packed))
 			return _physical.pinAdmitted(sId, tId, allowance);
 
-		return packed.state().pinAdmitted(_physical, allowance).map(physicalEntry -> {
-			if(physicalEntry == null)
-				return null;
-			return createLogicalPin(new BlockKey(sId, tId), packed);
-		});
+		BlockKey key = new BlockKey(sId, tId);
+		return pinLogical(key, packed, () -> packed.state().pinAdmitted(_physical, allowance));
 	}
 
 	@Override
@@ -265,9 +260,19 @@ public BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance) {
 		if(!(location instanceof SealedPackLocation packed))
 			return _physical.pinIfLive(sId, tId, allowance);
 
-		if(packed.state().pinIfLive(_physical, allowance) == null)
-			return null;
-		return createLogicalPin(new BlockKey(sId, tId), packed);
+		BlockKey key = new BlockKey(sId, tId);
+		packed.retain();
+		try {
+			if(packed.state().pinIfLive(_physical, allowance) == null) {
+				releaseLocation(key, packed);
+				return null;
+			}
+			return createLogicalPin(key, packed);
+		}
+		catch(RuntimeException | Error error) {
+			releaseLocation(key, packed);
+			throw error;
+		}
 	}
 
 	@Override
@@ -410,6 +415,7 @@ private UnpinHandle unpinPending(BlockEntry entry, PendingLogicalPin pin, Memory
 			entry.unpin();
 			entry.setCacheMeta(null);
 			PackedUnpinHandle handle = pin.builder().unpinProducer(entry, pin.slot(), allowance);
+			releasePendingPin(entry.getKey(), pin.builder(), pin.slot());
 			if(pin.builder().sealed && pin.builder().activePins == 0)
 				pin.builder().transferProducerOwnership(_physical);
 			scheduleSeal(pin.builder());
@@ -426,7 +432,9 @@ private UnpinHandle unpinPacked(BlockEntry entry, PackedLogicalPin pin, MemoryAl
 		}
 		entry.unpin();
 		entry.setCacheMeta(null);
-		return pin.location().state().unpin(this, _packReleaseDelayMs, allowance);
+		UnpinHandle handle = pin.location().state().unpin(this, _packReleaseDelayMs, allowance);
+		releaseLocation(entry.getKey(), pin.location());
+		return handle;
 	}
 
 	void enqueueRelease(PackedPinState state) {
@@ -492,6 +500,34 @@ private SealedPackLocation forceSeal(PendingPackLocation pending) {
 		}
 	}
 
+	private OOCFuture<BlockEntry> pinLogical(BlockKey key, SealedPackLocation location,
+		Supplier<OOCFuture<BlockEntry>> pin) {
+		location.retain();
+		OOCFuture<BlockEntry> physical;
+		try {
+			physical = pin.get();
+		}
+		catch(RuntimeException | Error error) {
+			releaseLocation(key, location);
+			throw error;
+		}
+		physical.whenComplete((entry, error) -> {
+			if(entry == null || error != null)
+				releaseLocation(key, location);
+		});
+		return physical.map(entry -> entry == null ? null : createLogicalPin(key, location));
+	}
+
+	private void releasePendingPin(BlockKey key, PackBuilder builder, int slot) {
+		if(builder.sealed) {
+			PackedCacheLocation location = getLocation(key.getStreamId(), key.getSequenceNumber());
+			if(location instanceof SealedPackLocation packed)
+				releaseLocation(key, packed);
+		}
+		else if(builder.releaseSlot(slot) == 0)
+			clearLocation(key);
+	}
+
 	private static BlockEntry createLogicalPin(BlockKey logicalKey, SealedPackLocation location) {
 		PackedBlock block = (PackedBlock) location.state().physicalEntry.getDataUnsafe();
 		Object data = block.values[location.slot()];

src/main/java/org/apache/sysds/runtime/ooc/cache/packed/PackBuilder.java

@@ -59,7 +59,7 @@ int append(long streamId, long tileId, Object value, long size) {
 		tileIds[slot] = tileId;
 		values[slot] = value;
 		sizes[slot] = size;
-		refCounts[slot] = 1;
+		refCounts[slot] = 2;
 		bytes += size;
 		activePins++;
 		return slot;

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

@@ -21,12 +21,16 @@
 
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.ToIntFunction;
 
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.instructions.ooc.CachingStream;
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
 import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
 import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
 import org.apache.sysds.runtime.ooc.cache.OOCCacheManager;
 import org.apache.sysds.runtime.ooc.cache.OOCFuture;
 import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
@@ -40,19 +44,32 @@ public final class MaterializeOOCPrimitive extends OOCPrimitive {
 	private final OOCStoreLayout _layout;
 	private final OOCFuture<MaterializedStore<IndexedMatrixValue>> _store;
 	private final AtomicBoolean _finished;
+	private final boolean _reusable;
 	private int _expectedReaders;
 	private int _consumers;
 
 	public MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> source, OOCStoreLayout layout,
 		StreamContext context) {
+		this(source, layout, context, false);
+	}
+
+	private MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> source, OOCStoreLayout layout,
+		StreamContext context, boolean reusable) {
 		super(context, source.getPrimitive() == null ? List.of() : List.of(source.getPrimitive()));
 		_source = source;
 		_layout = layout;
 		_store = new OOCFuture<>();
 		_finished = new AtomicBoolean();
+		_reusable = reusable;
+	}
+
+	public static MaterializeOOCPrimitive reusable(OOCStreamable<IndexedMatrixValue> source) {
+		return new MaterializeOOCPrimitive(source, OOCStoreLayout.ROW_MAJOR, null, true);
 	}
 
 	public synchronized void registerRequest(int expectedReaders) {
+		if(_reusable)
+			throw new IllegalStateException("Reusable materialization registers readers dynamically.");
 		if(expectedReaders <= 0)
 			throw new IllegalArgumentException("Materialization request requires at least one reader.");
 		if(hasStartedExecution())
@@ -84,10 +101,19 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) {
 	protected void startExecution() {
 		try {
 			OOCStream<IndexedMatrixValue> source = _source.getReservedReadStream();
-			MaterializedStore<IndexedMatrixValue> store = new MaterializedStore<>(OOCCacheManager.getGlobalCache(),
-				CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers);
-			OOCStreamMaterializer materializer = new OOCStreamMaterializer(store,
-				indexes -> _layout.linearize(indexes, _source.getDataCharacteristics()), _allowance);
+			MaterializedStore<IndexedMatrixValue> store = _reusable ? new MaterializedStore<>(
+				OOCCacheManager.getGlobalCache(),
+				CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(),
+					CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers);
+			DataCharacteristics characteristics = _source.getDataCharacteristics();
+			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);
 			materializer.completion().whenComplete((ignored, error) -> {
 				if(error != null)
 					fail(error);