Skip to content

[SYSTEMDS-3891] OOC Primitive Simplification and Injection Robustness #2578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
janniklinde merged 1 commit into from
Aug 4, 2026
Merged

[SYSTEMDS-3891] OOC Primitive Simplification and Injection Robustness #2578

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,16 +56,12 @@ private static void injectMaterializations(OOCPrimitive primitive, Set<OOCPrimit
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;
});
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);
}
Expand Down
15 changes: 5 additions & 10 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -46,7 +45,7 @@
import org.apache.sysds.runtime.ooc.util.OOCUtils;

public final class GroupedReduceOOCPrimitive extends OOCPrimitive {
private final OOCStream<IndexedMatrixValue> _input;
private final OOCStreamable<IndexedMatrixValue> _input;
private final OOCStreamable<IndexedMatrixValue> _output;
private final BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> _merge;
private final AtomicBoolean _cleaned;
Expand All @@ -62,12 +61,7 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive {

public GroupedReduceOOCPrimitive(OOCStreamable<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> merge, StreamContext context) {
this(input.getReadStream(), output, merge, context);
}

private GroupedReduceOOCPrimitive(OOCStream<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> merge, StreamContext context) {
super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive()));
super(context, input);
_input = input;
_output = output;
_merge = merge;
Expand Down Expand Up @@ -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<IndexedMatrixValue> 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())
Expand All @@ -122,7 +117,7 @@ protected void startExecution() {
OOCUtils.estimateFullTileBytes(_output.getDataCharacteristics()));
long pinBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(logicalBytes);
long taskBytes = pinBytes + logicalBytes * 2;
AllocatedOOCStream<IndexedMatrixValue> admitted = new AllocatedOOCStream<>(_input, _allowance,
AllocatedOOCStream<IndexedMatrixValue> admitted = new AllocatedOOCStream<>(input, _allowance,
ignored -> taskBytes);
getContext().addInStream(admitted);
admitted.setSubscriber(this::accept);
Expand Down
36 changes: 17 additions & 19 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -42,22 +40,16 @@
import org.apache.sysds.runtime.ooc.util.StateTableUtils;

public class JoinOOCPrimitive extends OOCPrimitive {
private final OOCStream<IndexedMatrixValue> _left;
private final OOCStream<IndexedMatrixValue> _right;
private final OOCStreamable<IndexedMatrixValue> _left;
private final OOCStreamable<IndexedMatrixValue> _right;
private final OOCStreamable<IndexedMatrixValue> _output;
private final BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> _operation;
private StateTable<IndexedMatrixValue> _table;

public JoinOOCPrimitive(OOCStreamable<IndexedMatrixValue> left, OOCStreamable<IndexedMatrixValue> right,
OOCStreamable<IndexedMatrixValue> output, BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> operation,
StreamContext context) {
this(left.getReadStream(), right.getReadStream(), output, operation, context);
}

private JoinOOCPrimitive(OOCStream<IndexedMatrixValue> left, OOCStream<IndexedMatrixValue> right,
OOCStreamable<IndexedMatrixValue> output, BiFunction<MatrixBlock, MatrixBlock, MatrixBlock> 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;
Expand All @@ -84,6 +76,8 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) {

@Override
protected void startExecution() {
OOCStream<IndexedMatrixValue> left = getInputReadStream(0);
OOCStream<IndexedMatrixValue> right = getInputReadStream(1);
_table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID());
OOCStream<IndexedMatrixValue> output = _output.getWriteStream();
OOCStream<JoinWork> matches = new SubscribableTaskQueue<>();
Expand All @@ -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 {
Expand All @@ -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<JoinWork> matches, long taskBytes) {
private void drive(OOCStream<IndexedMatrixValue> leftInput, OOCStream<IndexedMatrixValue> rightInput,
OOCStream<JoinWork> matches, long taskBytes) {
long cols = _right.getDataCharacteristics().getNumColBlocks();
int unmatched = 0;
try {
while(true) {
OOCStream.QueueCallback<IndexedMatrixValue> left = _left.dequeueCB();
OOCStream.QueueCallback<IndexedMatrixValue> right = _right.dequeueCB();
OOCStream.QueueCallback<IndexedMatrixValue> left = leftInput.dequeueCB();
OOCStream.QueueCallback<IndexedMatrixValue> right = rightInput.dequeueCB();
boolean leftEos = left == null || left.isEos();
boolean rightEos = right == null || right.isEos();
if(leftEos || rightEos) {
Expand Down
22 changes: 8 additions & 14 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -31,43 +30,38 @@
import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;

public class MappingOOCPrimitive extends OOCPrimitive {
private final OOCStream<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.getReadStream(), output, operation, context);
}

private MappingOOCPrimitive(OOCStream<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
Function<IndexedMatrixValue, MatrixBlock> 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();
}

@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<IndexedMatrixValue> input = getInputReadStream(0);
OOCStream<IndexedMatrixValue> output = _output.getWriteStream();
OOCInstructionUtils
.submitAdmittedOOCTasks(_input, output,
.submitAdmittedOOCTasks(input, output,
value -> new IndexedMatrixValue(value.getIndexes(), _operation.apply(value)), _allowance, getContext())
.thenRun(this::onComplete);
}
Expand Down
5 changes: 2 additions & 3 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/MaterializeOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -55,7 +54,7 @@ public MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> source, OOCStor

private MaterializeOOCPrimitive(OOCStreamable<IndexedMatrixValue> 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<>();
Expand Down Expand Up @@ -100,7 +99,7 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) {
@Override
protected void startExecution() {
try {
OOCStream<IndexedMatrixValue> source = _source.getReservedReadStream();
OOCStream<IndexedMatrixValue> source = getInputReadStream(0);
MaterializedStore<IndexedMatrixValue> store = _reusable ? new MaterializedStore<>(
OOCCacheManager.getGlobalCache(),
CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(),
Expand Down
22 changes: 11 additions & 11 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<>();
Expand Down Expand Up @@ -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.");
Expand All @@ -141,13 +141,13 @@ public final void discardInputHandle(int index) {

@SuppressWarnings("unchecked")
protected final <T> OOCStream<T> getInputReadStream(int index) {
transferInputHandle(index);
consumeInputHandle(index);
return (OOCStream<T>) _inputs.get(index)._source.getReservedReadStream();
}

protected final OOCFuture<MaterializedStore<IndexedMatrixValue>> getMaterializedInput(int index) {
OOCFuture<MaterializedStore<IndexedMatrixValue>> 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;
Expand Down Expand Up @@ -184,8 +184,8 @@ public final void requestPattern(OOCAccessPattern accessPattern) {
private void rebuildInputChildren() {
List<OOCPrimitive> 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);
Expand All @@ -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();
}
Expand Down
3 changes: 1 addition & 2 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -42,7 +41,7 @@ public class PlannableDataGenOOCPrimitive extends OOCPrimitive {

public PlannableDataGenOOCPrimitive(OOCStreamable<IndexedMatrixValue> output,
Function<MatrixIndexes, MatrixBlock> operation, StreamContext context) {
super(context, List.of());
super(context);
_output = output;
_operation = operation;
}
Expand Down
22 changes: 8 additions & 14 deletions src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -32,41 +31,36 @@
import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;

public class TransposeOOCPrimitive extends OOCPrimitive {
private final OOCStream<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.getReadStream(), output, operation, context);
}

private TransposeOOCPrimitive(OOCStream<IndexedMatrixValue> input, OOCStreamable<IndexedMatrixValue> output,
Function<MatrixBlock, MatrixBlock> 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<IndexedMatrixValue> input = getInputReadStream(0);
OOCStream<IndexedMatrixValue> 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()));
Expand Down
Loading
Loading