[SYSTEMDS-3891] Add OOC Primitives

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

@@ -20,6 +20,7 @@
 package org.apache.sysds.runtime.instructions.ooc;
 
 import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 import java.util.function.Consumer;
 
@@ -48,6 +49,12 @@ static <T> QueueCallback<T> eos(DMLRuntimeException e) {
 	 */
 	void setSubscriber(Consumer<QueueCallback<T>> subscriber);
 
+	default void start() {
+		OOCPrimitive primitive = getPrimitive();
+		if(primitive != null)
+			primitive.tryStartExecution();
+	}
+
 	interface QueueCallback<T> extends AutoCloseable {
 		T get();
 

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

@@ -21,6 +21,7 @@
 
 import org.apache.sysds.runtime.controlprogram.caching.CacheableData;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 public interface OOCStreamable<T> {
 	OOCStream<T> getReadStream();
@@ -38,4 +39,12 @@ public interface OOCStreamable<T> {
 	CacheableData<?> getData();
 
 	void setData(CacheableData<?> data);
+
+	default OOCPrimitive getPrimitive() {
+		return null;
+	}
+
+	default void assignPrimitive(OOCPrimitive primitive) {
+		throw new UnsupportedOperationException("Stream does not support primitive assignment");
+	}
 }

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

@@ -23,6 +23,7 @@
 import org.apache.sysds.runtime.controlprogram.caching.CacheableData;
 import org.apache.sysds.runtime.controlprogram.parfor.LocalTaskQueue;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 import org.apache.sysds.runtime.ooc.util.OOCUtils;
 
 import java.util.LinkedList;
@@ -37,6 +38,7 @@ public class SubscribableTaskQueue<T> extends LocalTaskQueue<OOCStream.QueueCall
 	private final AtomicInteger _blockCount = new AtomicInteger(0);
 	private QueueCallback<T> _lastDequeued = null;
 	private CacheableData<?> _cdata;
+	private volatile OOCPrimitive _primitive;
 	private volatile Consumer<QueueCallback<T>> _subscriber = null;
 	private String _watchdogId;
 
@@ -276,6 +278,18 @@ public CachingStream getStreamCache() {
 		return null;
 	}
 
+	@Override
+	public OOCPrimitive getPrimitive() {
+		return _primitive;
+	}
+
+	@Override
+	public void assignPrimitive(OOCPrimitive primitive) {
+		if(_primitive != null)
+			throw new IllegalStateException("Primitive already assigned");
+		_primitive = primitive;
+	}
+
 	@Override
 	public DataCharacteristics getDataCharacteristics() {
 		return _cdata == null ? null : _cdata.getDataCharacteristics();

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

@@ -0,0 +1,53 @@
+/*
+ * 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;
+
+public enum OOCAccessPattern {
+	ROW_MAJOR, COL_MAJOR, ANY, UNKNOWN, UNSET;
+
+	public OOCAccessPattern fused(OOCAccessPattern other) {
+		return switch(this) {
+			case ANY -> other;
+			case UNKNOWN, UNSET -> this;
+			case ROW_MAJOR -> other == ROW_MAJOR || other == ANY ? this : UNKNOWN;
+			case COL_MAJOR -> other == COL_MAJOR || other == ANY ? this : UNKNOWN;
+		};
+	}
+
+	public OOCAccessPattern transposed() {
+		return switch(this) {
+			case ROW_MAJOR -> COL_MAJOR;
+			case COL_MAJOR -> ROW_MAJOR;
+			default -> this;
+		};
+	}
+
+	public OOCAccessPattern preferred(OOCAccessPattern preferred) {
+		return this == ANY || this == UNSET ? preferred : this;
+	}
+
+	public boolean isPlannable() {
+		return this == ROW_MAJOR || this == COL_MAJOR || this == ANY;
+	}
+
+	public boolean isUnset() {
+		return this == UNSET;
+	}
+}

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

@@ -0,0 +1,92 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+
+public abstract class OOCPrimitive {
+	private final List<OOCPrimitive> _children;
+	private final List<OOCPrimitive> _parents;
+	private final AtomicBoolean _executionStarted;
+	protected OOCAccessPattern _pattern;
+
+	protected OOCPrimitive(List<OOCPrimitive> children) {
+		_parents = new ArrayList<>();
+		List<OOCPrimitive> uniqueChildren = new ArrayList<>(children.size());
+		for(OOCPrimitive child : children) {
+			if(containsIdentity(uniqueChildren, child))
+				continue;
+			uniqueChildren.add(child);
+			child.addParent(this);
+		}
+		_children = List.copyOf(uniqueChildren);
+		_executionStarted = new AtomicBoolean();
+		_pattern = OOCAccessPattern.UNSET;
+	}
+
+	public final List<OOCPrimitive> getChildren() {
+		return _children;
+	}
+
+	public final List<OOCPrimitive> getParents() {
+		return List.copyOf(_parents);
+	}
+
+	private void addParent(OOCPrimitive parent) {
+		if(!containsIdentity(_parents, parent))
+			_parents.add(parent);
+	}
+
+	protected final void inferParentPatterns() {
+		for(OOCPrimitive parent : _parents)
+			if(parent._pattern.isUnset())
+				parent.inferPatterns();
+	}
+
+	public final OOCAccessPattern getAccessPattern() {
+		return _pattern;
+	}
+
+	public final boolean hasStartedExecution() {
+		return _executionStarted.get();
+	}
+
+	public final void tryStartExecution() {
+		if(_executionStarted.compareAndSet(false, true))
+			startExecution();
+	}
+
+	private static boolean containsIdentity(List<OOCPrimitive> primitives, OOCPrimitive primitive) {
+		for(OOCPrimitive current : primitives)
+			if(current == primitive)
+				return true;
+		return false;
+	}
+
+	protected abstract void startExecution();
+
+	public abstract void inferPatterns();
+
+	public abstract void requestPattern(OOCAccessPattern accessPattern);
+}

src/main/java/org/apache/sysds/runtime/ooc/stream/AllocatedOOCStream.java

@@ -27,6 +27,7 @@
 import org.apache.sysds.runtime.ooc.cache.OOCFuture;
 import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
 import org.apache.sysds.runtime.ooc.memory.ReservationBudget;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 public final class AllocatedOOCStream<T> extends SubscribableTaskQueue<T> {
 	private final OOCStream<T> _source;
@@ -49,6 +50,11 @@ public static ReservationBudget detachBudget(OOCStream.QueueCallback<?> callback
 		return callback instanceof BudgetedQueueCallback<?> budgeted ? budgeted.detachBudget() : null;
 	}
 
+	@Override
+	public OOCPrimitive getPrimitive() {
+		return _source.getPrimitive();
+	}
+
 	private void admit(OOCStream.QueueCallback<T> callback) {
 		if(callback.isFailure()) {
 			try(callback) {

src/main/java/org/apache/sysds/runtime/ooc/stream/FilteredOOCStream.java

@@ -24,6 +24,7 @@
 import org.apache.sysds.runtime.instructions.ooc.CachingStream;
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -88,6 +89,11 @@ public CachingStream getStreamCache() {
 		return _sourceStream.getStreamCache();
 	}
 
+	@Override
+	public OOCPrimitive getPrimitive() {
+		return _sourceStream.getPrimitive();
+	}
+
 	@Override
 	public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 		_sourceStream.setSubscriber(cb -> {

src/main/java/org/apache/sysds/runtime/ooc/stream/MergedOOCStream.java

@@ -197,6 +197,12 @@ public OOCStream<T> getReadStream() {
 		return this;
 	}
 
+	@Override
+	public void start() {
+		for(OOCStream<T> source : _sources)
+			source.start();
+	}
+
 	@Override
 	public OOCStream<T> getWriteStream() {
 		throw new UnsupportedOperationException();

src/main/java/org/apache/sysds/runtime/ooc/stream/SplittingOOCStream.java

@@ -24,6 +24,7 @@
 import org.apache.sysds.runtime.instructions.ooc.CachingStream;
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -144,6 +145,11 @@ public CachingStream getStreamCache() {
 		return _sourceStream.getStreamCache();
 	}
 
+	@Override
+	public OOCPrimitive getPrimitive() {
+		return _sourceStream.getPrimitive();
+	}
+
 	@Override
 	public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 		throw new UnsupportedOperationException();

src/main/java/org/apache/sysds/runtime/ooc/stream/SubOOCStream.java

@@ -25,6 +25,7 @@
 import org.apache.sysds.runtime.instructions.ooc.OOCStream;
 import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue;
 import org.apache.sysds.runtime.meta.DataCharacteristics;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
 
 import java.util.function.Consumer;
 
@@ -86,6 +87,11 @@ public CachingStream getStreamCache() {
 		return _sourceStream.getStreamCache();
 	}
 
+	@Override
+	public OOCPrimitive getPrimitive() {
+		return _sourceStream.getPrimitive();
+	}
+
 	@Override
 	public void setSubscriber(Consumer<QueueCallback<T>> subscriber) {
 		_taskQueue.setSubscriber(cb -> {

src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java

@@ -0,0 +1,78 @@
+/*
+ * 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.test.component.ooc;
+
+import java.util.List;
+
+import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive;
+import org.apache.sysds.runtime.ooc.stream.FilteredOOCStream;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class OOCPrimitiveTest {
+	@Test
+	public void testGraphPatternsAndExecution() {
+		TestPrimitive source = new TestPrimitive(List.of());
+		TestPrimitive sink = new TestPrimitive(List.of(source, source));
+
+		Assert.assertEquals(List.of(source), sink.getChildren());
+		Assert.assertEquals(List.of(sink), source.getParents());
+		source.inferPatterns();
+		Assert.assertEquals(OOCAccessPattern.ANY, source.getAccessPattern());
+		Assert.assertEquals(OOCAccessPattern.ANY, sink.getAccessPattern());
+		Assert.assertEquals(OOCAccessPattern.COL_MAJOR, OOCAccessPattern.ROW_MAJOR.transposed());
+		Assert.assertEquals(OOCAccessPattern.UNKNOWN, OOCAccessPattern.ROW_MAJOR.fused(OOCAccessPattern.COL_MAJOR));
+
+		SubscribableTaskQueue<Integer> stream = new SubscribableTaskQueue<>();
+		stream.assignPrimitive(sink);
+		FilteredOOCStream<Integer> filtered = new FilteredOOCStream<>(stream, ignored -> true);
+		Assert.assertSame(sink, filtered.getPrimitive());
+		stream.start();
+		filtered.start();
+		Assert.assertTrue(sink.hasStartedExecution());
+		Assert.assertEquals(1, sink._executions);
+	}
+
+	private static final class TestPrimitive extends OOCPrimitive {
+		private int _executions;
+
+		private TestPrimitive(List<OOCPrimitive> children) {
+			super(children);
+		}
+
+		@Override
+		protected void startExecution() {
+			_executions++;
+		}
+
+		@Override
+		public void inferPatterns() {
+			_pattern = OOCAccessPattern.ANY;
+			inferParentPatterns();
+		}
+
+		@Override
+		public void requestPattern(OOCAccessPattern accessPattern) {
+			_pattern = accessPattern;
+		}
+	}
+}