From 5668ac5de99cc14f348b2a0a692b6a7e5657662e Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Mon, 17 Aug 2026 13:20:38 +0200 Subject: [PATCH] [OOC] Wire OOC Aggregate Unary and Bugfix Race Condition Assisted-by: AI --- .../ooc/AggregateUnaryOOCInstruction.java | 177 ++++++++++-------- .../sysds/runtime/ooc/cache/OOCCacheImpl.java | 33 +++- .../primitives/GroupedReduceOOCPrimitive.java | 113 +++++++++-- .../runtime/ooc/util/OOCInstructionUtils.java | 10 +- .../test/component/ooc/OOCPrimitiveTest.java | 38 ++++ 5 files changed, 267 insertions(+), 104 deletions(-) diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java index ac4e9bac919..0de35832f32 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/AggregateUnaryOOCInstruction.java @@ -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; @@ -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 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 corrs = new HashMap<>(); // correction blocks - - OOCStream qOut = createWritableStream(); - OOCStream 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 corrs = new HashMap<>(); + OOCStream qOut = createWritableStream(); + OOCStream 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 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 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 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; } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java index 9b0008e84ab..90f81d43bed 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java @@ -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 scheduled = _ioHandler.scheduleRead(meta.entry); meta.readFuture = scheduled; readFuture = scheduled; @@ -354,8 +354,10 @@ else if(meta.readFuture == null) { } }); } - else + else { + awaitRead(meta); readFuture = meta.readFuture; + } } if(releaseReserved) { allowance.release(reservedBytes); @@ -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; @@ -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)) { @@ -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)) @@ -705,6 +725,7 @@ private static class EntryMeta { private final BlockEntry entry; private boolean backed; private OOCFuture readFuture; + private int readWaiters; private CacheUnpinHandle deferredUnpin; private EntryMeta(BlockEntry entry) { 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 b380848aa7d..c94c2108dcf 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 @@ -22,6 +22,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiFunction; +import java.util.function.Function; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.CachingStream; @@ -45,9 +46,16 @@ import org.apache.sysds.runtime.ooc.util.OOCUtils; public final class GroupedReduceOOCPrimitive extends OOCPrimitive { + public enum Grouping { + ROW_BLOCKS, COL_BLOCKS + } + private final OOCStreamable _input; private final OOCStreamable _output; + private final Grouping _grouping; + private final Function _partial; private final BiFunction _merge; + private final Function _finish; private final AtomicBoolean _cleaned; private final AtomicBoolean _sourceComplete; private final AtomicInteger _active; @@ -60,10 +68,21 @@ public final class GroupedReduceOOCPrimitive extends OOCPrimitive { public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStreamable output, BiFunction merge, StreamContext context) { + this(input, output, Grouping.ROW_BLOCKS, value -> (MatrixBlock) value.getValue(), merge, Function.identity(), + context); + } + + public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStreamable output, + Grouping grouping, Function partial, + BiFunction merge, Function finish, + StreamContext context) { super(context, input); _input = input; _output = output; + _grouping = grouping; + _partial = partial; _merge = merge; + _finish = finish; _cleaned = new AtomicBoolean(); _sourceComplete = new AtomicBoolean(); _active = new AtomicInteger(1); @@ -72,17 +91,21 @@ public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStr @Override protected void inferPatternsInternal() { - _pattern = OOCAccessPattern.ROW_MAJOR; + _pattern = groupingPattern(); for(OOCPrimitive child : getChildren()) - child.requestPattern(OOCAccessPattern.ROW_MAJOR); + child.requestPattern(_pattern); inferParentPatterns(); } @Override protected void requestPatternInternal(OOCAccessPattern accessPattern) { - _pattern = OOCAccessPattern.ROW_MAJOR; + _pattern = groupingPattern(); for(OOCPrimitive child : getChildren()) - child.requestPattern(OOCAccessPattern.ROW_MAJOR); + child.requestPattern(_pattern); + } + + private OOCAccessPattern groupingPattern() { + return _grouping == Grouping.COL_BLOCKS ? OOCAccessPattern.COL_MAJOR : OOCAccessPattern.ROW_MAJOR; } @Override @@ -91,8 +114,7 @@ protected void startExecution() { 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()); + configureGroups(inputDc); _outputStream = _output.getWriteStream(); _ready = new SubscribableTaskQueue<>(); getContext().addInStream(input).addOutStream(_outputStream, _ready); @@ -121,6 +143,25 @@ protected void startExecution() { admitted.setSubscriber(this::accept); } + private void configureGroups(DataCharacteristics inputDc) { + long rowBlocks = inputDc.getNumRowBlocks(); + long colBlocks = inputDc.getNumColBlocks(); + switch(_grouping) { + case ROW_BLOCKS: + _numGroups = Math.toIntExact(rowBlocks); + _groupSize = Math.toIntExact(colBlocks); + break; + case COL_BLOCKS: + _numGroups = Math.toIntExact(colBlocks); + _groupSize = Math.toIntExact(rowBlocks); + break; + default: + throw new IllegalStateException("Unsupported grouped-reduce grouping: " + _grouping); + } + if(_numGroups <= 0 || _groupSize <= 0) + throw new DMLRuntimeException("Grouped OOC reduction requires non-empty input block geometry."); + } + private void accept(OOCStream.QueueCallback callback) { if(callback.isEos() || callback.isFailure()) { try(callback) { @@ -140,11 +181,11 @@ private void accept(OOCStream.QueueCallback callback) { try(callback) { budget = AllocatedOOCStream.detachBudget(callback).enableReuse(); IndexedMatrixValue input = callback.get(); - int group = Math.toIntExact(input.getIndexes().getRowIndex() - 1); - if(group < 0 || group >= _numGroups) - throw new DMLRuntimeException("Invalid grouped-reduce row block: " + (group + 1)); - IndexedMatrixValue value = new IndexedMatrixValue(new MatrixIndexes(group + 1L, 1), input.getValue()); - payload = payload(value, budget); + int group = group(input.getIndexes()); + MatrixBlock partial = _partial.apply(input); + if(partial == null) + throw new DMLRuntimeException("Grouped OOC reduction produced a null partial block."); + payload = payload(partialValue(group, 1, partial), budget); reduce(group, payload, budget); payload = null; budget = null; @@ -211,9 +252,11 @@ private void process(MergeWork work) { IndexedMatrixValue right = work._incoming.value(); int count = Math.addExact(multiplicity(left), multiplicity(right)); if(count > _groupSize) - throw new DMLRuntimeException("Too many partial tiles for grouped-reduce row " + (work._group + 1)); + throw new DMLRuntimeException("Too many partial tiles for grouped-reduce group " + (work._group + 1)); MatrixBlock value = _merge.apply((MatrixBlock) left.getValue(), (MatrixBlock) right.getValue()); - merged = payload(new IndexedMatrixValue(new MatrixIndexes(work._group + 1L, count), value), budget); + if(value == null) + throw new DMLRuntimeException("Grouped OOC reduction produced a null merged block."); + merged = payload(partialValue(work._group, count, value), budget); work.releaseIncoming(); released = work.closeExistingAsync(); } @@ -240,14 +283,18 @@ private void process(MergeWork work) { } private void finalizeGroup(int group, ManagedPayload payload, ReservationBudget budget) { - IndexedMatrixValue accumulated = payload.value(); - IndexedMatrixValue output = new IndexedMatrixValue(new MatrixIndexes(group + 1L, 1), accumulated.getValue()); - payload.release(); try { - OOCUtils.enqueueExact(_outputStream, output, budget); + MatrixBlock outputBlock = _finish.apply((MatrixBlock) payload.value().getValue()); + if(outputBlock == null) + throw new DMLRuntimeException("Grouped OOC reduction produced a null final block."); + payload.release(); + payload = null; + OOCUtils.enqueueExact(_outputStream, new IndexedMatrixValue(outputIndexes(group), outputBlock), budget); _finalizedGroups.incrementAndGet(); } catch(Throwable failure) { + if(payload != null) + payload.release(); budget.close(); fail(failure); } @@ -260,8 +307,34 @@ private static ManagedPayload payload(IndexedMatrixValue val return new ManagedPayload<>(value, bytes, budget); } - private static int multiplicity(IndexedMatrixValue value) { - return Math.toIntExact(value.getIndexes().getColumnIndex()); + private int group(MatrixIndexes indexes) { + long group = switch(_grouping) { + case ROW_BLOCKS -> indexes.getRowIndex() - 1; + case COL_BLOCKS -> indexes.getColumnIndex() - 1; + }; + if(group < 0 || group >= _numGroups) + throw new DMLRuntimeException("Invalid grouped-reduce group index: " + group); + return Math.toIntExact(group); + } + + private IndexedMatrixValue partialValue(int group, int count, MatrixBlock value) { + MatrixIndexes indexes = switch(_grouping) { + case ROW_BLOCKS -> new MatrixIndexes(group + 1L, count); + case COL_BLOCKS -> new MatrixIndexes(count, group + 1L); + }; + return new IndexedMatrixValue(indexes, value); + } + + private int multiplicity(IndexedMatrixValue value) { + return Math.toIntExact( + _grouping == Grouping.COL_BLOCKS ? value.getIndexes().getRowIndex() : value.getIndexes().getColumnIndex()); + } + + private MatrixIndexes outputIndexes(int group) { + return switch(_grouping) { + case ROW_BLOCKS -> new MatrixIndexes(group + 1L, 1); + case COL_BLOCKS -> new MatrixIndexes(1, group + 1L); + }; } private void finishSource() { @@ -275,7 +348,7 @@ private void completeOne() { return; if(!hasFailed() && _finalizedGroups.get() != _numGroups) fail(new DMLRuntimeException( - "Grouped reduction completed " + _finalizedGroups.get() + " of " + _numGroups + " row groups.")); + "Grouped reduction completed " + _finalizedGroups.get() + " of " + _numGroups + " groups.")); try { _ready.closeInput(); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java index e93876dfc03..4f2bcc34853 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java @@ -103,7 +103,15 @@ public static void indexedBroadcastMap(OOCStreamable streame public static void rowGroupedReduce(OOCStreamable input, OOCStream output, BiFunction merge, StreamContext context) { - output.assignPrimitive(new GroupedReduceOOCPrimitive(input, output, merge, context)); + groupedReduceIndexed(input, output, GroupedReduceOOCPrimitive.Grouping.ROW_BLOCKS, + value -> (MatrixBlock) value.getValue(), merge, Function.identity(), context); + } + + public static void groupedReduceIndexed(OOCStreamable input, + OOCStream output, GroupedReduceOOCPrimitive.Grouping grouping, + Function partial, BiFunction merge, + Function finish, StreamContext context) { + output.assignPrimitive(new GroupedReduceOOCPrimitive(input, output, grouping, partial, merge, finish, context)); } public static int getComputeInFlight() { 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 ab83b474233..c94e0ac5f1d 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 @@ -39,6 +39,7 @@ import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.primitives.GroupedReduceOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.MaterializeOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; import org.apache.sysds.runtime.ooc.store.CountingLiveness; @@ -187,6 +188,43 @@ public void testDataGenMapTransposePipeline() { values); } + @Test + public void testGroupedReduceModes() { + Assert.assertEquals(Map.of("1,1", 136d, "2,1", 166d), + runGroupedReduce(GroupedReduceOOCPrimitive.Grouping.ROW_BLOCKS, 2, 1)); + Assert.assertEquals(Map.of("1,1", 132d, "1,2", 134d, "1,3", 136d), + runGroupedReduce(GroupedReduceOOCPrimitive.Grouping.COL_BLOCKS, 1, 3)); + } + + private static Map runGroupedReduce(GroupedReduceOOCPrimitive.Grouping grouping, long outputRows, + long outputCols) { + SubscribableTaskQueue input = new SubscribableTaskQueue<>(); + SubscribableTaskQueue output = new SubscribableTaskQueue<>(); + input.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 3, 1), FileFormat.BINARY))); + output.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(outputRows, outputCols, 1), FileFormat.BINARY))); + for(long[] indexes : List.of(new long[] {2, 3}, new long[] {1, 1}, new long[] {2, 1}, new long[] {1, 3}, + new long[] {1, 2}, new long[] {2, 2})) + input.enqueue(new IndexedMatrixValue(new MatrixIndexes(indexes[0], indexes[1]), + new MatrixBlock(1, 1, indexes[0] * 10d + indexes[1]))); + input.closeInput(); + OOCInstructionUtils.groupedReduceIndexed(input, output, grouping, value -> (MatrixBlock) value.getValue(), + (left, right) -> new MatrixBlock(1, 1, left.get(0, 0) + right.get(0, 0)), + value -> new MatrixBlock(1, 1, value.get(0, 0) + 100), new StreamContext()); + + output.start(); + Map values = new HashMap<>(); + OOCStream.QueueCallback callback; + while((callback = output.dequeueCB()) != null) + try(OOCStream.QueueCallback current = callback) { + IndexedMatrixValue value = current.get(); + values.put(value.getIndexes().getRowIndex() + "," + value.getIndexes().getColumnIndex(), + value.getValue().get(0, 0)); + } + return values; + } + @Test public void testJoinOutOfOrder() { SubscribableTaskQueue left = new SubscribableTaskQueue<>();