[SYSTEMDS-3891] Add OOC Planner and First Primitive Implementations

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

@@ -51,8 +51,8 @@ static <T> QueueCallback<T> eos(DMLRuntimeException e) {
 
 	default void start() {
 		OOCPrimitive primitive = getPrimitive();
-		if(primitive != null)
-			primitive.tryStartExecution();
+		if(primitive != null && !primitive.hasStartedExecution())
+			primitive.start();
 	}
 
 	interface QueueCallback<T> extends AutoCloseable {

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

@@ -0,0 +1,54 @@
+/*
+ * 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 java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Set;
+
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
+
+public final class OOCPlanner {
+	public static void compile(OOCPrimitive root) {
+		List<OOCPrimitive> primitives = new ArrayList<>();
+		collect(root, Collections.newSetFromMap(new IdentityHashMap<>()), primitives);
+		if(primitives.isEmpty())
+			return;
+
+		for(int i = primitives.size() - 1; i >= 0; i--)
+			if(primitives.get(i).getAccessPattern().isUnset())
+				primitives.get(i).inferPatterns();
+		if(root.getAccessPattern() == OOCAccessPattern.ANY || root.getAccessPattern().isUnset())
+			root.requestPattern(OOCAccessPattern.ROW_MAJOR);
+
+		for(OOCPrimitive primitive : primitives)
+			primitive.tryStartExecution();
+	}
+
+	private static void collect(OOCPrimitive primitive, Set<OOCPrimitive> visited, List<OOCPrimitive> result) {
+		if(primitive.hasStartedExecution() || !visited.add(primitive))
+			return;
+		result.add(primitive);
+		for(OOCPrimitive child : primitive.getChildren())
+			collect(child, visited, result);
+	}
+}

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

@@ -0,0 +1,76 @@
+/*
+ * 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.function.Function;
+
+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.MatrixBlock;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
+
+public class MappingOOCPrimitive extends OOCPrimitive {
+	private final OOCStreamable<IndexedMatrixValue> _input;
+	private final OOCStreamable<IndexedMatrixValue> _output;
+	private final Function<IndexedMatrixValue, MatrixBlock> _operation;
+
+	public MappingOOCPrimitive(OOCStreamable<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
+		Function<IndexedMatrixValue, MatrixBlock> operation, StreamContext context) {
+		this(input.getPrimitive(), input, output, operation, context);
+	}
+
+	private MappingOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable<IndexedMatrixValue> input,
+		OOCStreamable<IndexedMatrixValue> output, Function<IndexedMatrixValue, MatrixBlock> operation,
+		StreamContext context) {
+		super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive));
+		_input = input;
+		_output = output;
+		_operation = operation;
+	}
+
+	@Override
+	protected void inferPatternsInternal() {
+		OOCAccessPattern inputPattern = getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0)
+			.getAccessPattern();
+		_pattern = _pattern.preferred(inputPattern);
+		inferParentPatterns();
+	}
+
+	@Override
+	protected void requestPatternInternal(OOCAccessPattern accessPattern) {
+		_pattern = _pattern.preferred(accessPattern);
+		if(!getChildren().isEmpty())
+			getChildren().get(0).requestPattern(accessPattern);
+	}
+
+	@Override
+	protected void startExecution() {
+		OOCStream<IndexedMatrixValue> input = _input.getReadStream();
+		OOCStream<IndexedMatrixValue> output = _output.getWriteStream();
+		OOCInstructionUtils
+			.submitAdmittedOOCTasks(input, output,
+				value -> new IndexedMatrixValue(value.getIndexes(), _operation.apply(value)), _allowance, getContext())
+			.thenRun(this::onComplete);
+	}
+}

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

@@ -23,15 +23,24 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker;
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance;
 import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.planning.OOCPlanner;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
 
 public abstract class OOCPrimitive {
+	private final StreamContext _context;
 	private final List<OOCPrimitive> _children;
 	private final List<OOCPrimitive> _parents;
+	private final AtomicBoolean _started;
 	private final AtomicBoolean _executionStarted;
 	protected OOCAccessPattern _pattern;
+	protected MemoryAllowance _allowance;
 
-	protected OOCPrimitive(List<OOCPrimitive> children) {
+	protected OOCPrimitive(StreamContext context, List<OOCPrimitive> children) {
+		_context = context;
 		_parents = new ArrayList<>();
 		List<OOCPrimitive> uniqueChildren = new ArrayList<>(children.size());
 		for(OOCPrimitive child : children) {
@@ -41,10 +50,15 @@ protected OOCPrimitive(List<OOCPrimitive> children) {
 			child.addParent(this);
 		}
 		_children = List.copyOf(uniqueChildren);
+		_started = new AtomicBoolean();
 		_executionStarted = new AtomicBoolean();
 		_pattern = OOCAccessPattern.UNSET;
 	}
 
+	public final StreamContext getContext() {
+		return _context;
+	}
+
 	public final List<OOCPrimitive> getChildren() {
 		return _children;
 	}
@@ -72,9 +86,30 @@ public final boolean hasStartedExecution() {
 		return _executionStarted.get();
 	}
 
+	public final void start() {
+		if(_started.compareAndSet(false, true))
+			OOCPlanner.compile(this);
+	}
+
 	public final void tryStartExecution() {
-		if(_executionStarted.compareAndSet(false, true))
+		if(_executionStarted.compareAndSet(false, true)) {
+			_allowance = new SyncMemoryAllowance(GlobalMemoryBroker.get());
 			startExecution();
+		}
+	}
+
+	public final void onComplete() {
+		_allowance.shutdown();
+	}
+
+	public final void inferPatterns() {
+		if(!hasStartedExecution())
+			inferPatternsInternal();
+	}
+
+	public final void requestPattern(OOCAccessPattern accessPattern) {
+		if(!hasStartedExecution() && _pattern != accessPattern)
+			requestPatternInternal(accessPattern);
 	}
 
 	private static boolean containsIdentity(List<OOCPrimitive> primitives, OOCPrimitive primitive) {
@@ -86,7 +121,7 @@ private static boolean containsIdentity(List<OOCPrimitive> primitives, OOCPrimit
 
 	protected abstract void startExecution();
 
-	public abstract void inferPatterns();
+	protected abstract void inferPatternsInternal();
 
-	public abstract void requestPattern(OOCAccessPattern accessPattern);
+	protected abstract void requestPatternInternal(OOCAccessPattern accessPattern);
 }

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

@@ -0,0 +1,91 @@
+/*
+ * 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.function.Function;
+
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
+import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue;
+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.ooc.memory.ReservationBudget;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
+import org.apache.sysds.runtime.ooc.util.OOCUtils;
+
+public class PlannableDataGenOOCPrimitive extends OOCPrimitive {
+	private final OOCStreamable<IndexedMatrixValue> _output;
+	private final Function<MatrixIndexes, MatrixBlock> _operation;
+
+	public PlannableDataGenOOCPrimitive(OOCStreamable<IndexedMatrixValue> output,
+		Function<MatrixIndexes, MatrixBlock> operation, StreamContext context) {
+		super(context, List.of());
+		_output = output;
+		_operation = operation;
+	}
+
+	@Override
+	protected void inferPatternsInternal() {
+		if(_pattern.isUnset())
+			_pattern = OOCAccessPattern.ANY;
+		inferParentPatterns();
+	}
+
+	@Override
+	protected void requestPatternInternal(OOCAccessPattern accessPattern) {
+		_pattern = accessPattern;
+	}
+
+	@Override
+	protected void startExecution() {
+		OOCStream<MatrixIndexes> work = new SubscribableTaskQueue<>();
+		OOCStream<IndexedMatrixValue> output = _output.getWriteStream();
+		long outputBytes = OOCUtils.estimateOutputTileBytes(_output.getDataCharacteristics());
+		AllocatedOOCStream<MatrixIndexes> admitted = new AllocatedOOCStream<>(work, _allowance, ignored -> outputBytes);
+		getContext().addOutStream(output);
+		OOCInstructionUtils.submitOOCTasks(admitted, callback -> {
+			ReservationBudget budget = AllocatedOOCStream.detachBudget(callback);
+			try {
+				MatrixIndexes indexes = callback.get();
+				OOCUtils.enqueueExact(output, new IndexedMatrixValue(indexes, _operation.apply(indexes)), budget);
+				budget = null;
+			}
+			finally {
+				if(budget != null)
+					budget.close();
+			}
+		}, getContext()).thenRun(output::closeInput).exceptionally(error -> {
+			output.propagateFailure(DMLRuntimeException.of(error));
+			return null;
+		}).thenRun(this::onComplete);
+
+		OOCInstructionUtils.submitOOCTask(() -> {
+			for(MatrixIndexes indexes : OOCUtils.getAccessPattern(_output.getDataCharacteristics(), _pattern))
+				work.enqueue(indexes);
+			work.closeInput();
+		}, new StreamContext().addOutStream(work));
+	}
+}

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

@@ -0,0 +1,76 @@
+/*
+ * 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.function.Function;
+
+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.MatrixBlock;
+import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
+
+public class TransposeOOCPrimitive extends OOCPrimitive {
+	private final OOCStreamable<IndexedMatrixValue> _input;
+	private final OOCStreamable<IndexedMatrixValue> _output;
+	private final Function<MatrixBlock, MatrixBlock> _operation;
+
+	public TransposeOOCPrimitive(OOCStreamable<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
+		Function<MatrixBlock, MatrixBlock> operation, StreamContext context) {
+		this(input.getPrimitive(), input, output, operation, context);
+	}
+
+	private TransposeOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable<IndexedMatrixValue> input,
+		OOCStreamable<IndexedMatrixValue> output, Function<MatrixBlock, MatrixBlock> operation, StreamContext context) {
+		super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive));
+		_input = input;
+		_output = output;
+		_operation = operation;
+	}
+
+	@Override
+	protected void inferPatternsInternal() {
+		_pattern = (getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0).getAccessPattern())
+			.transposed();
+		inferParentPatterns();
+	}
+
+	@Override
+	protected void requestPatternInternal(OOCAccessPattern accessPattern) {
+		_pattern = accessPattern;
+		if(!getChildren().isEmpty())
+			getChildren().get(0).requestPattern(accessPattern.transposed());
+	}
+
+	@Override
+	protected void startExecution() {
+		OOCStream<IndexedMatrixValue> input = _input.getReadStream();
+		OOCStream<IndexedMatrixValue> output = _output.getWriteStream();
+		OOCInstructionUtils.submitAdmittedOOCTasks(input, output, value -> {
+			MatrixIndexes indexes = value.getIndexes();
+			return new IndexedMatrixValue(new MatrixIndexes(indexes.getColumnIndex(), indexes.getRowIndex()),
+				_operation.apply((MatrixBlock) value.getValue()));
+		}, _allowance, getContext()).thenRun(this::onComplete);
+	}
+}