[SYSTEMDS-3891] Add Materialized Planning Support

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

@@ -79,6 +79,7 @@ public class CachingStream implements OOCStreamable<IndexedMatrixValue> {
 
 	private boolean _deletable = false;
 	private int _maxConsumptionCount = 0;
+	private int _lazyHandleReservations = 0;
 	private String _watchdogId = null;
 
 	public CachingStream(OOCStream<IndexedMatrixValue> source) {
@@ -356,7 +357,7 @@ public synchronized void scheduleDeletion() {
 		if (_deletable)
 			return; // Deletion already scheduled
 
-		if (_cacheInProgress && _maxConsumptionCount == 0)
+		if(_cacheInProgress && _maxConsumptionCount == 0 && _lazyHandleReservations == 0)
 			System.out.println("[WARN] Scheduling deletion for caching stream with no listeners: " + this);
 
 		_deletable = true;
@@ -370,6 +371,8 @@ public String toString() {
 	}
 
 	private synchronized void tryDeleteBlock(int i) {
+		if(_lazyHandleReservations > 0)
+			return;
 		int cnt = _consumptionCounts.getInt(i);
 		if (cnt > _maxConsumptionCount)
 			throw new DMLRuntimeException("Cannot have more than " + _maxConsumptionCount + " consumptions.");
@@ -582,6 +585,11 @@ public OOCStream<IndexedMatrixValue> getReadStream() {
 		return new PlaybackStream(this);
 	}
 
+	@Override
+	public OOCStream<IndexedMatrixValue> getReservedReadStream() {
+		return new PlaybackStream(this, consumeLazyHandleReservation());
+	}
+
 	@Override
 	public OOCStream<IndexedMatrixValue> getWriteStream() {
 		return _source.getWriteStream();
@@ -721,6 +729,29 @@ public synchronized void incrSubscriberCount(int count) {
 		_maxConsumptionCount += count;
 	}
 
+	@Override
+	public synchronized void reserveLazyHandle() {
+		_lazyHandleReservations++;
+	}
+
+	@Override
+	public synchronized void discardHandle() {
+		if(_lazyHandleReservations <= 0)
+			return;
+		_lazyHandleReservations--;
+		if(_deletable)
+			for(int i = 0; i < _consumptionCounts.size(); i++)
+				tryDeleteBlock(i);
+	}
+
+	private synchronized boolean consumeLazyHandleReservation() {
+		if(_lazyHandleReservations <= 0)
+			return false;
+		_lazyHandleReservations--;
+		_maxConsumptionCount++;
+		return true;
+	}
+
 	/**
 	 * Artificially increase the processing count of a block.
 	 */

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

@@ -47,4 +47,14 @@ default OOCPrimitive getPrimitive() {
 	default void assignPrimitive(OOCPrimitive primitive) {
 		throw new UnsupportedOperationException("Stream does not support primitive assignment");
 	}
+
+	default OOCStream<T> getReservedReadStream() {
+		return getReadStream();
+	}
+
+	default void reserveLazyHandle() {
+	}
+
+	default void discardHandle() {
+	}
 }

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

@@ -37,10 +37,15 @@ public class PlaybackStream implements OOCStream<IndexedMatrixValue> {
 	private QueueCallback<IndexedMatrixValue> _lastDequeue;
 
 	public PlaybackStream(CachingStream streamCache) {
+		this(streamCache, false);
+	}
+
+	public PlaybackStream(CachingStream streamCache, boolean reserved) {
 		this._streamCache = streamCache;
 		this._streamIdx = new AtomicInteger(0);
 		this._subscriberSet = new AtomicBoolean(false);
-		streamCache.incrSubscriberCount(1);
+		if(!reserved)
+			streamCache.incrSubscriberCount(1);
 	}
 
 	@Override

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

@@ -25,10 +25,14 @@
 import java.util.List;
 import java.util.Set;
 
+import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
+import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.ooc.primitives.MaterializeOOCPrimitive;
 import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 public final class OOCPlanner {
 	public static void compile(OOCPrimitive root) {
+		injectMaterializations(root, Collections.newSetFromMap(new IdentityHashMap<>()), new IdentityHashMap<>());
 		List<OOCPrimitive> primitives = new ArrayList<>();
 		collect(root, Collections.newSetFromMap(new IdentityHashMap<>()), primitives);
 		if(primitives.isEmpty())
@@ -44,6 +48,31 @@ public static void compile(OOCPrimitive root) {
 			primitive.tryStartExecution();
 	}
 
+	@SuppressWarnings("unchecked")
+	private static void injectMaterializations(OOCPrimitive primitive, Set<OOCPrimitive> visited,
+		IdentityHashMap<OOCStreamable<IndexedMatrixValue>, MaterializeOOCPrimitive> boundaries) {
+		if(primitive.hasStartedExecution() || !visited.add(primitive))
+			return;
+		for(OOCPrimitive.OOCMaterializedInputRequest request : primitive.requiredMaterializedInputs()) {
+			OOCStreamable<IndexedMatrixValue> input = (OOCStreamable<IndexedMatrixValue>) primitive
+				.getInput(request.inputIndex());
+			MaterializeOOCPrimitive boundary = boundaries.compute(input, (k, v) -> {
+				if(v == null) {
+					MaterializeOOCPrimitive p = new MaterializeOOCPrimitive(input, request.layout(),
+						primitive.getContext());
+					primitive.transferInputHandle(request.inputIndex());
+					return p;
+				}
+				primitive.discardInputHandle(request.inputIndex());
+				return v;
+			});
+			boundary.registerRequest(request.expectedReaders());
+			primitive.installMaterializedInput(request.inputIndex(), boundary);
+		}
+		for(OOCPrimitive child : primitive.getChildren())
+			injectMaterializations(child, visited, boundaries);
+	}
+
 	private static void collect(OOCPrimitive primitive, Set<OOCPrimitive> visited, List<OOCPrimitive> result) {
 		if(primitive.hasStartedExecution() || !visited.add(primitive))
 			return;

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

@@ -0,0 +1,36 @@
+/*
+ * 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.ooc.planning;
+
+import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
+import org.apache.sysds.runtime.meta.DataCharacteristics;
+
+public enum OOCStoreLayout {
+	ROW_MAJOR;
+
+	public int linearize(MatrixIndexes indexes, DataCharacteristics characteristics) {
+		if(characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0)
+			throw new IllegalArgumentException("Materialized store layout requires known dimensions and block size.");
+		long columns = characteristics.getNumColBlocks();
+		long index = Math.addExact(Math.multiplyExact(indexes.getRowIndex() - 1, columns),
+			indexes.getColumnIndex() - 1);
+		return Math.toIntExact(index);
+	}
+}

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

@@ -50,8 +50,8 @@ private MappingOOCPrimitive(OOCStream<IndexedMatrixValue> input, OOCStreamable<I
 
 	@Override
 	protected void inferPatternsInternal() {
-		OOCAccessPattern inputPattern = getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0)
-			.getAccessPattern();
+		OOCAccessPattern inputPattern = getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().stream()
+			.findFirst().get().getAccessPattern();
 		_pattern = _pattern.preferred(inputPattern);
 		inferParentPatterns();
 	}
@@ -60,7 +60,7 @@ protected void inferPatternsInternal() {
 	protected void requestPatternInternal(OOCAccessPattern accessPattern) {
 		_pattern = _pattern.preferred(accessPattern);
 		if(!getChildren().isEmpty())
-			getChildren().get(0).requestPattern(accessPattern);
+			getChildren().forEach(c -> c.requestPattern(accessPattern));
 	}
 
 	@Override

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

@@ -0,0 +1,117 @@
+/*
+ * 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.ooc.primitives;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+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.ooc.cache.OOCCacheManager;
+import org.apache.sysds.runtime.ooc.cache.OOCFuture;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout;
+import org.apache.sysds.runtime.ooc.store.MaterializedStore;
+import org.apache.sysds.runtime.ooc.store.OOCStreamMaterializer;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
+
+public final class MaterializeOOCPrimitive extends OOCPrimitive {
+	private final OOCStreamable<IndexedMatrixValue> _source;
+	private final OOCStoreLayout _layout;
+	private final OOCFuture<MaterializedStore<IndexedMatrixValue>> _store;
+	private final AtomicBoolean _finished;
+	private int _expectedReaders;
+	private int _consumers;
+
+	public MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> source, OOCStoreLayout layout,
+		StreamContext context) {
+		super(context, source.getPrimitive() == null ? List.of() : List.of(source.getPrimitive()));
+		_source = source;
+		_layout = layout;
+		_store = new OOCFuture<>();
+		_finished = new AtomicBoolean();
+	}
+
+	public synchronized void registerRequest(int expectedReaders) {
+		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);
+	}
+
+	public OOCFuture<MaterializedStore<IndexedMatrixValue>> store() {
+		return _store;
+	}
+
+	@Override
+	protected void inferPatternsInternal() {
+		_pattern = OOCAccessPattern.ROW_MAJOR;
+		for(OOCPrimitive child : getChildren())
+			child.requestPattern(OOCAccessPattern.ROW_MAJOR);
+		inferParentPatterns();
+	}
+
+	@Override
+	protected void requestPatternInternal(OOCAccessPattern accessPattern) {
+		_pattern = OOCAccessPattern.ROW_MAJOR;
+		for(OOCPrimitive child : getChildren())
+			child.requestPattern(OOCAccessPattern.ROW_MAJOR);
+	}
+
+	@Override
+	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);
+			materializer.completion().whenComplete((ignored, error) -> {
+				if(error != null)
+					fail(error);
+				finish();
+			});
+			if(getContext() != null)
+				getContext().addInStream(source);
+			_store.complete(store);
+			materializer.attach(source);
+		}
+		catch(Throwable failure) {
+			_store.completeExceptionally(failure);
+			fail(failure);
+			finish();
+		}
+	}
+
+	private void fail(Throwable error) {
+		if(getContext() != null)
+			getContext().failAll(DMLRuntimeException.of(error));
+	}
+
+	private void finish() {
+		if(_finished.compareAndSet(false, true))
+			onComplete();
+	}
+}