Skip to content

[OOC] Wire OOC Aggregate Unary and Bugfix Race Condition #2592

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 18, 2026
Merged

[OOC] Wire OOC Aggregate Unary and Bugfix Race Condition #2592

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
177 changes: 100 additions & 77 deletions src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@
import org.apache.sysds.runtime.matrix.operators.AggregateUnaryOperator;
import org.apache.sysds.runtime.matrix.operators.Operator;
import org.apache.sysds.runtime.meta.DataCharacteristics;
import org.apache.sysds.runtime.ooc.primitives.GroupedReduceOOCPrimitive;
import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
import org.apache.sysds.runtime.ooc.util.OOCUtils;

import java.util.HashMap;

Expand Down Expand Up @@ -72,88 +75,108 @@ public static AggregateUnaryOOCInstruction parseInstruction(String str) {
public void processInstruction( ExecutionContext ec ) {
//TODO support all types of aggregations, currently only full aggregation, row aggregation and column aggregation

//setup operators and input queue
AggregateUnaryOperator aggun = (AggregateUnaryOperator) getOperator();
MatrixObject min = ec.getMatrixObject(input1);
DataCharacteristics chars = ec.getDataCharacteristics(input1.getName());
int blen = chars != null && chars.getBlocksize() > 0 ? chars.getBlocksize() : ConfigurationManager
.getBlocksize();

if(!aggun.isRowAggregate() && !aggun.isColAggregate()) {
processScalarAggregate(ec, min, aggun, blen);
return;
}
if(OOCUtils.getNumBlocks(chars) > 0) {
processPlannerMatrixAggregate(ec, min, aggun, blen);
return;
}

OOCStream<IndexedMatrixValue> qIn = min.getStreamHandle();
int blen = ConfigurationManager.getBlocksize();

if (aggun.isRowAggregate() || aggun.isColAggregate()) {
DataCharacteristics chars = ec.getDataCharacteristics(input1.getName());
// number of blocks to process per aggregation idx (row or column dim)
long emitThreshold = aggun.isRowAggregate()? chars.getNumColBlocks() : chars.getNumRowBlocks();
OOCMatrixBlockTracker aggTracker = new OOCMatrixBlockTracker(emitThreshold);
HashMap<Long, MatrixBlock> corrs = new HashMap<>(); // correction blocks

OOCStream<IndexedMatrixValue> qOut = createWritableStream();
OOCStream<IndexedMatrixValue> qLocal = createWritableStream();

ec.getMatrixObject(output).setStreamHandle(qOut);

// per-block aggregation (parallel map)
mapOOC(qIn, qLocal, tmp -> {
MatrixIndexes midx = aggun.isRowAggregate() ?
new MatrixIndexes(tmp.getIndexes().getRowIndex(), 1) :
new MatrixIndexes(1, tmp.getIndexes().getColumnIndex());

MatrixBlock ltmp = (MatrixBlock) ((MatrixBlock) tmp.getValue())
.aggregateUnaryOperations(aggun, new MatrixBlock(), blen, tmp.getIndexes());
return new IndexedMatrixValue(midx, ltmp);
});

// global reduce
addOutStream(qOut);
submitOOCTasks(qLocal, callback -> {
IndexedMatrixValue partial = callback.get();
synchronized(aggTracker) {
long idx = aggun.isRowAggregate() ? partial.getIndexes().getRowIndex() : partial.getIndexes()
.getColumnIndex();

MatrixBlock ret = aggTracker.get(idx);
boolean ready;
if(ret != null) {
MatrixBlock corr = corrs.get(idx);
OperationsOnMatrixValues.incrementalAggregation(ret,
_aop.existsCorrection() ? corr : null, (MatrixBlock) partial.getValue(), _aop,
true);
ready = aggTracker.incrementCount(idx);
}
else {
ret = (MatrixBlock) partial.getValue();
MatrixBlock corr = _aop.existsCorrection() ? new MatrixBlock(ret.getNumRows(),
ret.getNumColumns(), false) : null;
ready = aggTracker.putAndIncrementCount(idx, ret);
if(!ready && _aop.existsCorrection())
corrs.put(idx, corr);
}

if(ready) {
ret.dropLastRowsOrColumns(_aop.correction);
qOut.enqueue(new IndexedMatrixValue(partial.getIndexes(), ret));
aggTracker.remove(idx);
corrs.remove(idx);
}
long emitThreshold = aggun.isRowAggregate() ? chars.getNumColBlocks() : chars.getNumRowBlocks();
OOCMatrixBlockTracker aggTracker = new OOCMatrixBlockTracker(emitThreshold);
HashMap<Long, MatrixBlock> corrs = new HashMap<>();
OOCStream<IndexedMatrixValue> qOut = createWritableStream();
OOCStream<IndexedMatrixValue> qLocal = createWritableStream();
ec.getMatrixObject(output).setStreamHandle(qOut);

mapOOC(qIn, qLocal, tmp -> {
MatrixIndexes midx = aggun.isRowAggregate() ? new MatrixIndexes(tmp.getIndexes().getRowIndex(),
1) : new MatrixIndexes(1, tmp.getIndexes().getColumnIndex());
MatrixBlock ltmp = (MatrixBlock) ((MatrixBlock) tmp.getValue()).aggregateUnaryOperations(aggun,
new MatrixBlock(), blen, tmp.getIndexes());
return new IndexedMatrixValue(midx, ltmp);
});

addOutStream(qOut);
submitOOCTasks(qLocal, callback -> {
IndexedMatrixValue partial = callback.get();
synchronized(aggTracker) {
long idx = aggun.isRowAggregate() ? partial.getIndexes().getRowIndex() : partial.getIndexes()
.getColumnIndex();
MatrixBlock ret = aggTracker.get(idx);
boolean ready;
if(ret != null) {
MatrixBlock corr = corrs.get(idx);
OperationsOnMatrixValues.incrementalAggregation(ret, _aop.existsCorrection() ? corr : null,
(MatrixBlock) partial.getValue(), _aop, true);
ready = aggTracker.incrementCount(idx);
}
else {
ret = (MatrixBlock) partial.getValue();
MatrixBlock corr = _aop.existsCorrection() ? new MatrixBlock(ret.getNumRows(), ret.getNumColumns(),
false) : null;
ready = aggTracker.putAndIncrementCount(idx, ret);
if(!ready && _aop.existsCorrection())
corrs.put(idx, corr);
}
if(ready) {
ret.dropLastRowsOrColumns(_aop.correction);
qOut.enqueue(new IndexedMatrixValue(partial.getIndexes(), ret));
aggTracker.remove(idx);
corrs.remove(idx);
}
}).thenRun(qOut::closeInput);
}
// full aggregation
else {
OOCStream<MatrixBlock> qLocal = createWritableStream();

mapOOC(qIn, qLocal, tmp -> (MatrixBlock) tmp.getValue()
.aggregateUnaryOperations(aggun, new MatrixBlock(), blen, tmp.getIndexes()));

MatrixBlock ltmp;
int extra = _aop.correction.getNumRemovedRowsColumns();
MatrixBlock ret = new MatrixBlock(1, 1 + extra, _aop.initialValue);
MatrixBlock corr = new MatrixBlock(1,1+extra,false);
while((ltmp = qLocal.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) {
OperationsOnMatrixValues.incrementalAggregation(
ret, _aop.existsCorrection() ? corr : null, ltmp, _aop, true);
}
}).thenRun(qOut::closeInput);
}

//create scalar output
ec.setScalarOutput(output.getName(), new DoubleObject(ret.get(0, 0)));
}
private void processPlannerMatrixAggregate(ExecutionContext ec, MatrixObject input, AggregateUnaryOperator operator,
int blocksize) {
OOCStream<IndexedMatrixValue> outputStream = createWritableStream();
ec.getMatrixObject(output).setStreamHandle(outputStream);
GroupedReduceOOCPrimitive.Grouping grouping = operator
.isRowAggregate() ? GroupedReduceOOCPrimitive.Grouping.ROW_BLOCKS : GroupedReduceOOCPrimitive.Grouping.COL_BLOCKS;
OOCInstructionUtils.groupedReduceIndexed(input.getStreamable(), outputStream, grouping,
value -> aggregatePartial(value, operator, blocksize), this::mergeAggregate, this::finalizeAggregate,
getContext());
}

private void processScalarAggregate(ExecutionContext ec, MatrixObject input, AggregateUnaryOperator operator,
int blocksize) {
OOCStream<MatrixBlock> partials = createWritableStream();
mapOOC(input.getStreamHandle(), partials, value -> aggregatePartial(value, operator, blocksize));

int extra = _aop.correction.getNumRemovedRowsColumns();
MatrixBlock result = new MatrixBlock(1, 1 + extra, _aop.initialValue);
MatrixBlock correction = new MatrixBlock(1, 1 + extra, false);
MatrixBlock partial;
while((partial = partials.dequeue()) != LocalTaskQueue.NO_MORE_TASKS)
OperationsOnMatrixValues.incrementalAggregation(result, _aop.existsCorrection() ? correction : null,
partial, _aop, true);
ec.setScalarOutput(output.getName(), new DoubleObject(result.get(0, 0)));
}

private static MatrixBlock aggregatePartial(IndexedMatrixValue value, AggregateUnaryOperator operator,
int blocksize) {
return (MatrixBlock) value.getValue().aggregateUnaryOperations(operator, new MatrixBlock(), blocksize,
value.getIndexes());
}

private MatrixBlock mergeAggregate(MatrixBlock left, MatrixBlock right) {
OperationsOnMatrixValues.incrementalAggregation(left, null, right, _aop, true);
return left;
}

private MatrixBlock finalizeAggregate(MatrixBlock block) {
block.dropLastRowsOrColumns(_aop.correction);
return block;
}
}
33 changes: 27 additions & 6 deletions src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,7 @@ else if(meta.entry.getDataUnsafe() != null) {
readFuture = null;
}
else if(meta.readFuture == null) {
meta.entry.setState(BlockState.READING);
awaitRead(meta);
OOCFuture<BlockEntry> scheduled = _ioHandler.scheduleRead(meta.entry);
meta.readFuture = scheduled;
readFuture = scheduled;
Expand All @@ -354,8 +354,10 @@ else if(meta.readFuture == null) {
}
});
}
else
else {
awaitRead(meta);
readFuture = meta.readFuture;
}
}
if(releaseReserved) {
allowance.release(reservedBytes);
Expand All @@ -372,19 +374,24 @@ else if(meta.readFuture == null) {
try {
if(ex != null) {
release = true;
synchronized(OOCCacheImpl.this) {
finishRead(meta);
}
allowance.release(reservedBytes);
result.completeExceptionally(ex);
return;
}
BlockEntry pinned;
synchronized(OOCCacheImpl.this) {
if(getMeta(meta.entry) != meta || meta.entry.getDataUnsafe() == null) {
if(getMeta(meta.entry) != meta) {
release = true;
if(meta.entry.getState() == BlockState.READING)
meta.entry.setState(BlockState.COLD);
pinned = null;
}
else {
finishRead(meta);
if(meta.entry.getDataUnsafe() == null)
throw new IllegalStateException(
"Backing read left no data for entry: " + meta.entry.getKey());
completion = pinResident(meta);
Statistics.incrementOOCEvictionGet();
pinned = meta.entry;
Expand All @@ -404,6 +411,18 @@ else if(meta.readFuture == null) {
return result;
}

private void awaitRead(EntryMeta meta) {
meta.readWaiters++;
clearLive(meta.entry);
meta.entry.setState(BlockState.READING);
}

private void finishRead(EntryMeta meta) {
meta.readWaiters = Math.max(0, meta.readWaiters - 1);
if(meta.readWaiters == 0 && meta.entry.getState() == BlockState.READING)
meta.entry.setState(BlockState.COLD);
}

private DeferredCompletion pinResident(EntryMeta meta) {
BlockEntry entry = meta.entry;
if(isCacheOwned(entry)) {
Expand Down Expand Up @@ -624,7 +643,8 @@ private EvictController getOrCreateEvictController(long streamId) {
}

private void removeIfUnused(EntryMeta meta) {
if(meta.entry.getReferenceCount() > 0 || meta.entry.getPinCount() > 0 || meta.deferredUnpin != null)
if(meta.entry.getReferenceCount() > 0 || meta.entry.getPinCount() > 0 || meta.deferredUnpin != null ||
meta.readWaiters > 0)
return;
BlockEntry entry = meta.entry;
if(isCacheOwned(entry))
Expand Down Expand Up @@ -705,6 +725,7 @@ private static class EntryMeta {
private final BlockEntry entry;
private boolean backed;
private OOCFuture<BlockEntry> readFuture;
private int readWaiters;
private CacheUnpinHandle deferredUnpin;

private EntryMeta(BlockEntry entry) {
Expand Down
Loading
Loading