From 0d0e9c26b2607c646fd6b7def768195067fb19f4 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Wed, 22 Jul 2026 12:23:51 +0200 Subject: [PATCH] [OOC] OOC Primitive Simplification and Injection Robustness --- .../runtime/ooc/planning/OOCPlanner.java | 16 ++++----- .../primitives/GroupedReduceOOCPrimitive.java | 15 +++----- .../ooc/primitives/JoinOOCPrimitive.java | 36 +++++++++---------- .../ooc/primitives/MappingOOCPrimitive.java | 22 +++++------- .../primitives/MaterializeOOCPrimitive.java | 5 ++- .../runtime/ooc/primitives/OOCPrimitive.java | 22 ++++++------ .../PlannableDataGenOOCPrimitive.java | 3 +- .../ooc/primitives/TransposeOOCPrimitive.java | 22 +++++------- .../store/MaterializedStoreStreamable.java | 4 +-- .../test/component/ooc/OOCPrimitiveTest.java | 3 +- 10 files changed, 62 insertions(+), 86 deletions(-) diff --git a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java index 38a5c609525..03eea8dcf31 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java @@ -56,16 +56,12 @@ private static void injectMaterializations(OOCPrimitive primitive, Set input = (OOCStreamable) 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; - }); + MaterializeOOCPrimitive boundary = boundaries.get(input); + if(boundary == null) { + boundary = new MaterializeOOCPrimitive(input, request.layout(), primitive.getContext()); + boundaries.put(input, boundary); + } + primitive.discardInputHandle(request.inputIndex()); boundary.registerRequest(request.expectedReaders()); primitive.installMaterializedInput(request.inputIndex(), boundary); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java index ad126af0c71..aa67497206e 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java @@ -19,7 +19,6 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; @@ -46,7 +45,7 @@ import org.apache.sysds.runtime.ooc.util.OOCUtils; public final class GroupedReduceOOCPrimitive extends OOCPrimitive { - private final OOCStream _input; + private final OOCStreamable _input; private final OOCStreamable _output; private final BiFunction _merge; private final AtomicBoolean _cleaned; @@ -62,12 +61,7 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive { public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStreamable output, BiFunction merge, StreamContext context) { - this(input.getReadStream(), output, merge, context); - } - - private GroupedReduceOOCPrimitive(OOCStream input, OOCStreamable output, - BiFunction merge, StreamContext context) { - super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); + super(context, input); _input = input; _output = output; _merge = merge; @@ -98,11 +92,12 @@ protected void startExecution() { DataCharacteristics inputDc = _input.getDataCharacteristics(); if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) throw new DMLRuntimeException("Grouped OOC reduction requires known input dimensions and block size."); + OOCStream input = getInputReadStream(0); _numGroups = Math.toIntExact(inputDc.getNumRowBlocks()); _groupSize = Math.toIntExact(inputDc.getNumColBlocks()); _outputStream = _output.getWriteStream(); _ready = new SubscribableTaskQueue<>(); - getContext().addInStream(_input).addOutStream(_outputStream, _ready); + getContext().addInStream(input).addOutStream(_outputStream, _ready); _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext()) @@ -122,7 +117,7 @@ protected void startExecution() { OOCUtils.estimateFullTileBytes(_output.getDataCharacteristics())); long pinBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(logicalBytes); long taskBytes = pinBytes + logicalBytes * 2; - AllocatedOOCStream admitted = new AllocatedOOCStream<>(_input, _allowance, + AllocatedOOCStream admitted = new AllocatedOOCStream<>(input, _allowance, ignored -> taskBytes); getContext().addInStream(admitted); admitted.setSubscriber(this::accept); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java index 99996eb7dac..a4ef2c30ced 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java @@ -19,10 +19,8 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.Objects; import java.util.concurrent.ExecutionException; import java.util.function.BiFunction; -import java.util.stream.Stream; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.CachingStream; @@ -42,8 +40,8 @@ import org.apache.sysds.runtime.ooc.util.StateTableUtils; public class JoinOOCPrimitive extends OOCPrimitive { - private final OOCStream _left; - private final OOCStream _right; + private final OOCStreamable _left; + private final OOCStreamable _right; private final OOCStreamable _output; private final BiFunction _operation; private StateTable _table; @@ -51,13 +49,7 @@ public class JoinOOCPrimitive extends OOCPrimitive { public JoinOOCPrimitive(OOCStreamable left, OOCStreamable right, OOCStreamable output, BiFunction operation, StreamContext context) { - this(left.getReadStream(), right.getReadStream(), output, operation, context); - } - - private JoinOOCPrimitive(OOCStream left, OOCStream right, - OOCStreamable output, BiFunction operation, - StreamContext context) { - super(context, Stream.of(left.getPrimitive(), right.getPrimitive()).filter(Objects::nonNull).toList()); + super(context, left, right); _left = left; _right = right; _output = output; @@ -84,6 +76,8 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { @Override protected void startExecution() { + OOCStream left = getInputReadStream(0); + OOCStream right = getInputReadStream(1); _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); OOCStream output = _output.getWriteStream(); OOCStream matches = new SubscribableTaskQueue<>(); @@ -95,10 +89,12 @@ protected void startExecution() { getContext().addOutStream(output); OOCInstructionUtils.submitOOCTasks(matches, callback -> { try(JoinWork work = callback.get()) { - IndexedMatrixValue left = work._left.get(); - IndexedMatrixValue right = work._right.get(); - OOCUtils.enqueueExact(output, new IndexedMatrixValue(left.getIndexes(), - _operation.apply((MatrixBlock) left.getValue(), (MatrixBlock) right.getValue())), work._budget); + IndexedMatrixValue mleft = work._left.get(); + IndexedMatrixValue mright = work._right.get(); + OOCUtils.enqueueExact(output, + new IndexedMatrixValue(mleft.getIndexes(), + _operation.apply((MatrixBlock) mleft.getValue(), (MatrixBlock) mright.getValue())), + work._budget); } }, callback -> true, (index, callback) -> callback.get().close(), getContext()).thenRun(() -> { try { @@ -110,16 +106,18 @@ protected void startExecution() { } }); - OOCInstructionUtils.submitOOCTask(() -> drive(matches, taskBytes), new StreamContext().addOutStream(output)); + OOCInstructionUtils.submitOOCTask(() -> drive(left, right, matches, taskBytes), + new StreamContext().addOutStream(output)); } - private void drive(OOCStream matches, long taskBytes) { + private void drive(OOCStream leftInput, OOCStream rightInput, + OOCStream matches, long taskBytes) { long cols = _right.getDataCharacteristics().getNumColBlocks(); int unmatched = 0; try { while(true) { - OOCStream.QueueCallback left = _left.dequeueCB(); - OOCStream.QueueCallback right = _right.dequeueCB(); + OOCStream.QueueCallback left = leftInput.dequeueCB(); + OOCStream.QueueCallback right = rightInput.dequeueCB(); boolean leftEos = left == null || left.isEos(); boolean rightEos = right == null || right.isEos(); if(leftEos || rightEos) { diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java index c4a96bc01dc..318c6c84cc4 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java @@ -19,7 +19,6 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.List; import java.util.function.Function; import org.apache.sysds.runtime.instructions.ooc.OOCStream; @@ -31,27 +30,20 @@ import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class MappingOOCPrimitive extends OOCPrimitive { - private final OOCStream _input; private final OOCStreamable _output; private final Function _operation; public MappingOOCPrimitive(OOCStreamable input, OOCStreamable output, Function operation, StreamContext context) { - this(input.getReadStream(), output, operation, context); - } - - private MappingOOCPrimitive(OOCStream input, OOCStreamable output, - Function operation, StreamContext context) { - super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); - _input = input; + super(context, input); _output = output; _operation = operation; } @Override protected void inferPatternsInternal() { - OOCAccessPattern inputPattern = getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().stream() - .findFirst().get().getAccessPattern(); + OOCPrimitive dependency = getInputDependency(0); + OOCAccessPattern inputPattern = dependency == null ? OOCAccessPattern.ANY : dependency.getAccessPattern(); _pattern = _pattern.preferred(inputPattern); inferParentPatterns(); } @@ -59,15 +51,17 @@ protected void inferPatternsInternal() { @Override protected void requestPatternInternal(OOCAccessPattern accessPattern) { _pattern = _pattern.preferred(accessPattern); - if(!getChildren().isEmpty()) - getChildren().forEach(c -> c.requestPattern(accessPattern)); + OOCPrimitive dependency = getInputDependency(0); + if(dependency != null) + dependency.requestPattern(accessPattern); } @Override protected void startExecution() { + OOCStream input = getInputReadStream(0); OOCStream output = _output.getWriteStream(); OOCInstructionUtils - .submitAdmittedOOCTasks(_input, output, + .submitAdmittedOOCTasks(input, output, value -> new IndexedMatrixValue(value.getIndexes(), _operation.apply(value)), _allowance, getContext()) .thenRun(this::onComplete); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java index e1002e98005..92bc28f8537 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java @@ -19,7 +19,6 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.ToIntFunction; @@ -55,7 +54,7 @@ public MaterializeOOCPrimitive(OOCStreamable source, OOCStor private MaterializeOOCPrimitive(OOCStreamable source, OOCStoreLayout layout, StreamContext context, boolean reusable) { - super(context, source.getPrimitive() == null ? List.of() : List.of(source.getPrimitive())); + super(context, source); _source = source; _layout = layout; _store = new OOCFuture<>(); @@ -100,7 +99,7 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { @Override protected void startExecution() { try { - OOCStream source = _source.getReservedReadStream(); + OOCStream source = getInputReadStream(0); MaterializedStore store = _reusable ? new MaterializedStore<>( OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(), diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java index 0bc20d54bd4..12d2d3a6d1f 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java @@ -64,7 +64,7 @@ protected OOCPrimitive(StreamContext context, OOCStreamable... inputs) { rebuildInputChildren(); } - private OOCPrimitive(StreamContext context) { + protected OOCPrimitive(StreamContext context) { _context = context; _children = new HashSet<>(); _parents = new HashSet<>(); @@ -108,19 +108,19 @@ public final OOCStreamable getInput(int index) { return _inputs.get(index)._source; } - public final OOCPrimitive getChildPrimitiveAt(int index) { - return _inputs.get(index)._primitive; + public final OOCPrimitive getInputDependency(int index) { + return _inputs.get(index)._dependency; } public final void installMaterializedInput(int index, MaterializeOOCPrimitive boundary) { if(hasStartedExecution()) throw new IllegalStateException("Cannot replace an input after primitive execution started."); InputSlot input = _inputs.get(index); - input._primitive = boundary; + input._dependency = boundary; rebuildInputChildren(); } - public final synchronized void transferInputHandle(int index) { + private synchronized void consumeInputHandle(int index) { InputSlot input = _inputs.get(index); if(!input._handleReserved) throw new IllegalStateException("Input " + index + " no longer owns a lazy handle."); @@ -141,13 +141,13 @@ public final void discardInputHandle(int index) { @SuppressWarnings("unchecked") protected final OOCStream getInputReadStream(int index) { - transferInputHandle(index); + consumeInputHandle(index); return (OOCStream) _inputs.get(index)._source.getReservedReadStream(); } protected final OOCFuture> getMaterializedInput(int index) { OOCFuture> materialized = ((MaterializeOOCPrimitive) _inputs - .get(index)._primitive).store(); + .get(index)._dependency).store(); if(materialized == null) throw new IllegalStateException("Input " + index + " was not materialized by the planner."); return materialized; @@ -184,8 +184,8 @@ public final void requestPattern(OOCAccessPattern accessPattern) { private void rebuildInputChildren() { List next = new ArrayList<>(); for(InputSlot input : _inputs) - if(input._primitive != null) - next.add(input._primitive); + if(input._dependency != null) + next.add(input._dependency); for(OOCPrimitive child : _children) if(!next.contains(child)) child._parents.remove(this); @@ -203,12 +203,12 @@ private void rebuildInputChildren() { private static final class InputSlot { private final OOCStreamable _source; - private OOCPrimitive _primitive; + private OOCPrimitive _dependency; private boolean _handleReserved; private InputSlot(OOCStreamable source) { _source = source; - _primitive = source.getPrimitive(); + _dependency = source.getPrimitive(); _handleReserved = true; source.reserveLazyHandle(); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java index 83ffcd2acf4..b604e66061a 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java @@ -19,7 +19,6 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.List; import java.util.function.Function; import org.apache.sysds.runtime.DMLRuntimeException; @@ -42,7 +41,7 @@ public class PlannableDataGenOOCPrimitive extends OOCPrimitive { public PlannableDataGenOOCPrimitive(OOCStreamable output, Function operation, StreamContext context) { - super(context, List.of()); + super(context); _output = output; _operation = operation; } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java index 18b127bfd09..c2aad428f4d 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java @@ -19,7 +19,6 @@ package org.apache.sysds.runtime.ooc.primitives; -import java.util.List; import java.util.function.Function; import org.apache.sysds.runtime.instructions.ooc.OOCStream; @@ -32,41 +31,36 @@ import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class TransposeOOCPrimitive extends OOCPrimitive { - private final OOCStream _input; private final OOCStreamable _output; private final Function _operation; public TransposeOOCPrimitive(OOCStreamable input, OOCStreamable output, Function operation, StreamContext context) { - this(input.getReadStream(), output, operation, context); - } - - private TransposeOOCPrimitive(OOCStream input, OOCStreamable output, - Function operation, StreamContext context) { - super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); - _input = input; + super(context, input); _output = output; _operation = operation; } @Override protected void inferPatternsInternal() { - _pattern = (getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().iterator().next().getAccessPattern()) - .transposed(); + OOCPrimitive dependency = getInputDependency(0); + _pattern = (dependency == null ? OOCAccessPattern.ANY : dependency.getAccessPattern()).transposed(); inferParentPatterns(); } @Override protected void requestPatternInternal(OOCAccessPattern accessPattern) { _pattern = accessPattern; - for(OOCPrimitive child : getChildren()) - child.requestPattern(accessPattern.transposed()); + OOCPrimitive dependency = getInputDependency(0); + if(dependency != null) + dependency.requestPattern(accessPattern.transposed()); } @Override protected void startExecution() { + OOCStream input = getInputReadStream(0); OOCStream output = _output.getWriteStream(); - OOCInstructionUtils.submitAdmittedOOCTasks(_input, output, value -> { + OOCInstructionUtils.submitAdmittedOOCTasks(input, output, value -> { MatrixIndexes indexes = value.getIndexes(); return new IndexedMatrixValue(new MatrixIndexes(indexes.getColumnIndex(), indexes.getRowIndex()), _operation.apply((MatrixBlock) value.getValue())); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java index 5d3879de72e..42f584ee2f8 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStoreStreamable.java @@ -208,8 +208,8 @@ private void markMaterializationDone() { @Override public synchronized void reserveLazyHandle() { - if(_deleteScheduled) - throw new DMLRuntimeException("Cannot reserve a reader on a materialized stream scheduled for deletion."); + if(_closed || (_deleteScheduled && _reservedReaders == 0)) + throw new DMLRuntimeException("Cannot reserve a reader on a closed materialized stream."); _reservedReaders++; } diff --git a/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java b/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java index 0fe0f2e01c3..ab83b474233 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/OOCPrimitiveTest.java @@ -110,7 +110,8 @@ public void testPlannerDoubleMaterialize() { sink.start(); Assert.assertEquals(1, sink.getChildren().size()); - Assert.assertTrue(sink.getChildPrimitiveAt(0) instanceof MaterializeOOCPrimitive); + Assert.assertTrue(sink.getInputDependency(0) instanceof MaterializeOOCPrimitive); + Assert.assertSame(sink.getInputDependency(0), sink.getInputDependency(1)); Assert.assertEquals(1, sink._executions); } finally {