You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by ja...@apache.org on 2026/08/18 15:53:20 UTC
(systemds) branch main updated: [SYSTEMDS-3891] Add Generic OOC Reduce
This is an automated email from the ASF dual-hosted git repository.
janniklinde pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new beb1fe3194 [SYSTEMDS-3891] Add Generic OOC Reduce
beb1fe3194 is described below
commit beb1fe3194e9cff2e9992a097d0490ae3b2a3fd0
Author: Jannik Lindemann <52...@users.noreply.github.com>
AuthorDate: Tue Aug 18 17:53:13 2026 +0200
[SYSTEMDS-3891] Add Generic OOC Reduce
Assisted-by: AI
---
.../instructions/ooc/TSMMOOCInstruction.java | 27 ++--
.../sysds/runtime/ooc/cache/OOCCacheManager.java | 3 +-
.../ooc/cache/legacy/OOCCacheScheduler.java | 2 +-
.../ooc/cache/legacy/OOCLRUCacheScheduler.java | 12 +-
.../sysds/runtime/ooc/memory/CachedAllowance.java | 22 ++--
.../runtime/ooc/memory/InMemoryQueueCallback.java | 47 +++----
.../runtime/ooc/primitives/ReduceOOCPrimitive.java | 146 +++++++++++++++++++++
.../runtime/ooc/util/OOCInstructionUtils.java | 7 +
.../apache/sysds/runtime/ooc/util/OOCUtils.java | 2 +-
.../sysds/test/component/ooc/OOCPrimitiveTest.java | 24 ++++
.../test/component/ooc/StateTableUtilsTest.java | 4 +-
.../ooc/memory/OOCMemoryAllowanceTest.java | 27 ++--
12 files changed, 247 insertions(+), 76 deletions(-)
diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java
index 37b2ba93a7..0707601b12 100644
--- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java
+++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java
@@ -28,7 +28,6 @@ import org.apache.sysds.lops.MMTSJ.MMTSJType;
import org.apache.sysds.runtime.DMLRuntimeException;
import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
-import org.apache.sysds.runtime.controlprogram.parfor.LocalTaskQueue;
import org.apache.sysds.runtime.functionobjects.Multiply;
import org.apache.sysds.runtime.functionobjects.Plus;
import org.apache.sysds.runtime.instructions.InstructionUtils;
@@ -41,6 +40,7 @@ import org.apache.sysds.runtime.matrix.operators.AggregateBinaryOperator;
import org.apache.sysds.runtime.matrix.operators.AggregateOperator;
import org.apache.sysds.runtime.matrix.operators.BinaryOperator;
import org.apache.sysds.runtime.matrix.operators.Operator;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
public class TSMMOOCInstruction extends ComputationOOCInstruction {
private final MMTSJType _type;
@@ -109,24 +109,15 @@ public class TSMMOOCInstruction extends ComputationOOCInstruction {
}
private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObject min) {
- OOCStream<IndexedMatrixValue> qIn = min.getStreamHandle();
+ OOCStream<IndexedMatrixValue> out = createWritableStream();
+ ec.getMatrixObject(output).setStreamHandle(out);
BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
- MatrixBlock resultBlock = null;
-
- OOCStream<MatrixBlock> tmpStream = createWritableStream();
- mapOOC(qIn, tmpStream,
- tmp -> ((MatrixBlock) tmp.getValue())
- .transposeSelfMatrixMultOperations(new MatrixBlock(), _type));
-
- MatrixBlock tmp;
- while((tmp = tmpStream.dequeue()) != LocalTaskQueue.NO_MORE_TASKS) {
- if(resultBlock == null)
- resultBlock = tmp;
- else
- resultBlock.binaryOperationsInPlace(plus, tmp);
- }
-
- ec.setMatrixOutput(output.getName(), resultBlock);
+ OOCInstructionUtils.reduce(min.getStreamable(), out,
+ value -> new IndexedMatrixValue(new MatrixIndexes(1, 1),
+ ((MatrixBlock) value.getValue()).transposeSelfMatrixMultOperations(new MatrixBlock(), _type)),
+ (left, right) -> new IndexedMatrixValue(new MatrixIndexes(1, 1),
+ ((MatrixBlock) left.getValue()).binaryOperationsInPlace(plus, right.getValue())),
+ value -> ((MatrixBlock) value.getValue()).getExactSerializedSize(), getContext());
}
private long getJoinIndex(IndexedMatrixValue value) {
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java
index 5cec3ae981..4bde5c9df4 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheManager.java
@@ -281,7 +281,8 @@ public class OOCCacheManager {
return getCache().isWithinLimits() && OOCInstruction.getComputeInFlight() <= OOCInstruction.getComputeBackpressureThreshold();
}
- public static OOCCacheScheduler.HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback) {
+ public static OOCCacheScheduler.HandoverHandle handover(BlockKey key,
+ InMemoryQueueCallback<IndexedMatrixValue> callback) {
return getCache().handover(key, callback);
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCCacheScheduler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCCacheScheduler.java
index ad161e9530..d8ff79dd79 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCCacheScheduler.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCCacheScheduler.java
@@ -106,7 +106,7 @@ public interface OOCCacheScheduler {
OOCStream.QueueCallback<IndexedMatrixValue> reclaim();
}
- HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback);
+ HandoverHandle handover(BlockKey key, InMemoryQueueCallback<IndexedMatrixValue> callback);
/**
* Places a new source-backed block in the cache and registers the location with the IO handler. The entry is
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java
index 3f5601adba..96de368ccc 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/legacy/OOCLRUCacheScheduler.java
@@ -293,7 +293,7 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
}
@Override
- public HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback) {
+ public HandoverHandle handover(BlockKey key, InMemoryQueueCallback<IndexedMatrixValue> callback) {
if(!this._running)
throw new IllegalStateException("Cache scheduler has been shut down.");
PendingHandover handover = new PendingHandover(key, callback);
@@ -1085,7 +1085,7 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
}
private boolean commitHandover(PendingHandover pending) {
- InMemoryQueueCallback callback = pending.takeForCommit();
+ InMemoryQueueCallback<IndexedMatrixValue> callback = pending.takeForCommit();
if(callback == null)
return false;
try {
@@ -1135,12 +1135,12 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
private static class PendingHandover implements HandoverHandle {
private final BlockKey _key;
private final CompletableFuture<Boolean> _completionFuture;
- private InMemoryQueueCallback _callback;
+ private InMemoryQueueCallback<IndexedMatrixValue> _callback;
private boolean _committed;
private boolean _cancelled;
private boolean _committing;
- private PendingHandover(BlockKey key, InMemoryQueueCallback callback) {
+ private PendingHandover(BlockKey key, InMemoryQueueCallback<IndexedMatrixValue> callback) {
_key = key;
_completionFuture = new CompletableFuture<>();
_callback = callback;
@@ -1180,11 +1180,11 @@ public class OOCLRUCacheScheduler implements OOCCacheScheduler {
return _cancelled;
}
- private synchronized InMemoryQueueCallback takeForCommit() {
+ private synchronized InMemoryQueueCallback<IndexedMatrixValue> takeForCommit() {
if(_committed || _cancelled || _committing)
return null;
_committing = true;
- InMemoryQueueCallback callback = _callback;
+ InMemoryQueueCallback<IndexedMatrixValue> callback = _callback;
_callback = null;
return callback;
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/CachedAllowance.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/CachedAllowance.java
index ffba3910b2..d2375b80c1 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/memory/CachedAllowance.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/CachedAllowance.java
@@ -56,12 +56,12 @@ public class CachedAllowance extends SyncMemoryAllowance {
_handoverSchedulingRequested = false;
}
- public void handover(InMemoryQueueCallback callback, int index) {
+ public void handover(InMemoryQueueCallback<IndexedMatrixValue> callback, int index) {
if(callback == null)
throw new IllegalArgumentException("Cannot hand over null callback.");
callback.transferOwnershipBlocking(this);
- InMemoryQueueCallback root = (InMemoryQueueCallback) callback.keepOpen();
+ InMemoryQueueCallback<IndexedMatrixValue> root = callback.keepOpen();
callback.close();
root.getHandle().attachCachedAllowance(this, index);
@@ -88,7 +88,7 @@ public class CachedAllowance extends SyncMemoryAllowance {
while(true) {
BlockKey cacheKey = null;
OOCCacheScheduler.HandoverHandle handover = null;
- InMemoryQueueCallback local = null;
+ InMemoryQueueCallback<IndexedMatrixValue> local = null;
synchronized(entry) {
if(entry._local != null && entry._handover == null)
@@ -124,7 +124,7 @@ public class CachedAllowance extends SyncMemoryAllowance {
if(!future.isDone())
return null;
boolean committed = future.join();
- InMemoryQueueCallback localToClose = null;
+ InMemoryQueueCallback<IndexedMatrixValue> localToClose = null;
synchronized(entry) {
if(entry._handover != handover)
continue;
@@ -171,8 +171,8 @@ public class CachedAllowance extends SyncMemoryAllowance {
throw DMLRuntimeException.of(ex.getCause() == null ? ex : ex.getCause());
return committed == true;
}).thenCompose(committed -> {
- InMemoryQueueCallback localToClose = null;
- InMemoryQueueCallback local = null;
+ InMemoryQueueCallback<IndexedMatrixValue> localToClose = null;
+ InMemoryQueueCallback<IndexedMatrixValue> local = null;
BlockKey key;
synchronized(entry) {
@@ -215,7 +215,7 @@ public class CachedAllowance extends SyncMemoryAllowance {
while(true) {
OOCCacheScheduler.HandoverHandle handover = null;
BlockKey forgetKey = null;
- InMemoryQueueCallback localToClose = null;
+ InMemoryQueueCallback<IndexedMatrixValue> localToClose = null;
synchronized(entry) {
if(entry._local != null && entry._handover == null) {
@@ -410,7 +410,7 @@ public class CachedAllowance extends SyncMemoryAllowance {
if(bytes <= 0)
return 0;
- InMemoryQueueCallback retained = (InMemoryQueueCallback) entry._local.keepOpen();
+ InMemoryQueueCallback<IndexedMatrixValue> retained = entry._local.keepOpen();
try {
entry._cacheKey = new BlockKey(_streamId, _nextBlockId.getAndIncrement());
entry._handover = OOCCacheManager.handover(entry._cacheKey, retained);
@@ -448,7 +448,7 @@ public class CachedAllowance extends SyncMemoryAllowance {
onFinishedHandover(bytes);
}
- private void closeRoot(InMemoryQueueCallback local) {
+ private void closeRoot(InMemoryQueueCallback<IndexedMatrixValue> local) {
local.getHandle().detachCachedAllowance();
local.close();
}
@@ -486,12 +486,12 @@ public class CachedAllowance extends SyncMemoryAllowance {
}
private static final class SlotEntry {
- private InMemoryQueueCallback _local;
+ private InMemoryQueueCallback<IndexedMatrixValue> _local;
private BlockKey _cacheKey;
private OOCCacheScheduler.HandoverHandle _handover;
private long _pendingBytes;
- private SlotEntry(InMemoryQueueCallback local) {
+ private SlotEntry(InMemoryQueueCallback<IndexedMatrixValue> local) {
_local = local;
}
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java
index 7fafc042e4..ae712305cb 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java
@@ -21,36 +21,38 @@ package org.apache.sysds.runtime.ooc.memory;
import org.apache.sysds.runtime.DMLRuntimeException;
import org.apache.sysds.runtime.instructions.ooc.OOCStream;
-import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
-
import java.util.concurrent.atomic.AtomicInteger;
-public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMatrixValue> {
- private CallbackHandle _handle;
+public class InMemoryQueueCallback<T> implements OOCStream.QueueCallback<T> {
+ private CallbackHandle<T> _handle;
private boolean _closed;
- public InMemoryQueueCallback(IndexedMatrixValue result, DMLRuntimeException failure, MemoryAllowance allow,
- long reservedBytes) {
- _handle = new CallbackHandle(result, failure, allow, reservedBytes);
+ public InMemoryQueueCallback(T result, DMLRuntimeException failure, MemoryAllowance allow, long reservedBytes) {
+ _handle = new CallbackHandle<>(result, failure, allow, reservedBytes);
_closed = false;
}
- private InMemoryQueueCallback(CallbackHandle handle) {
+ public InMemoryQueueCallback(ManagedPayload<T> payload) {
+ this(payload.value(), null, payload.owner(), payload.bytes());
+ payload.transfer();
+ }
+
+ private InMemoryQueueCallback(CallbackHandle<T> handle) {
_handle = handle;
_closed = false;
}
@Override
- public IndexedMatrixValue get() {
+ public T get() {
return _handle.get();
}
@Override
- public synchronized OOCStream.QueueCallback<IndexedMatrixValue> keepOpen() {
+ public synchronized InMemoryQueueCallback<T> keepOpen() {
if(_closed)
throw new IllegalStateException("Cannot keep open a closed callback");
_handle._refCtr.incrementAndGet();
- return new InMemoryQueueCallback(_handle);
+ return new InMemoryQueueCallback<>(_handle);
}
@Override
@@ -126,20 +128,19 @@ public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMat
return _handle._failure != null;
}
- CallbackHandle getHandle() {
+ CallbackHandle<T> getHandle() {
return _handle;
}
- static final class CallbackHandle {
- private volatile IndexedMatrixValue _result;
+ static final class CallbackHandle<T> {
+ private volatile T _result;
private final AtomicInteger _refCtr;
private MemoryAllowance _allow;
private long _reservedBytes;
private volatile DMLRuntimeException _failure;
private int _cacheIdx;
- private CallbackHandle(IndexedMatrixValue result, DMLRuntimeException failure, MemoryAllowance allow,
- long reservedBytes) {
+ private CallbackHandle(T result, DMLRuntimeException failure, MemoryAllowance allow, long reservedBytes) {
_result = result;
_failure = failure;
_refCtr = new AtomicInteger(1);
@@ -148,7 +149,7 @@ public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMat
_cacheIdx = -1;
}
- private IndexedMatrixValue get() {
+ private T get() {
if(_failure != null)
throw _failure;
return _result;
@@ -174,8 +175,8 @@ public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMat
return _refCtr.get() == 1;
}
- private synchronized IndexedMatrixValue takeManagedResultForHandover() {
- IndexedMatrixValue result = _result;
+ private synchronized T takeManagedResultForHandover() {
+ T result = _result;
_result = null;
return result;
}
@@ -188,14 +189,14 @@ public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMat
}
}
- public IndexedMatrixValue takeManagedResultForHandover() {
+ public T takeManagedResultForHandover() {
return _handle.takeManagedResultForHandover();
}
- public synchronized ManagedPayload<IndexedMatrixValue> extractManagedPayload() {
+ public synchronized ManagedPayload<T> extractManagedPayload() {
if(_closed)
throw new IllegalStateException("Cannot extract a managed payload from a closed callback.");
- CallbackHandle handle = _handle;
+ CallbackHandle<T> handle = _handle;
synchronized(handle) {
if(handle._failure != null)
throw handle._failure;
@@ -203,7 +204,7 @@ public class InMemoryQueueCallback implements OOCStream.QueueCallback<IndexedMat
throw new IllegalStateException("Cannot extract a managed payload while callback aliases exist.");
if(handle._cacheIdx >= 0)
throw new IllegalStateException("Cannot extract a managed payload from a cached-slot callback.");
- IndexedMatrixValue result = handle._result;
+ T result = handle._result;
if(result == null)
throw new IllegalStateException("Cannot extract a managed payload from an empty callback.");
long bytes = handle._reservedBytes;
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReduceOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReduceOOCPrimitive.java
new file mode 100644
index 0000000000..9f5ee85fed
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReduceOOCPrimitive.java
@@ -0,0 +1,146 @@
+/*
+ * 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.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.ToLongFunction;
+
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.ooc.OOCStreamable;
+import org.apache.sysds.runtime.ooc.cache.OOCCacheManager;
+import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback;
+import org.apache.sysds.runtime.ooc.memory.ManagedPayload;
+import org.apache.sysds.runtime.ooc.memory.ReservationBudget;
+import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern;
+import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream;
+import org.apache.sysds.runtime.ooc.stream.StreamContext;
+import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils;
+import org.apache.sysds.runtime.ooc.util.OOCUtils;
+
+public final class ReduceOOCPrimitive<I, O> extends OOCPrimitive {
+ private final OOCStreamable<I> _input;
+ private final OOCStreamable<O> _output;
+ private final Function<I, O> _partial;
+ private final BiFunction<O, O, O> _merge;
+ private final ToLongFunction<O> _size;
+ private ManagedPayload<O> _accumulator;
+
+ public ReduceOOCPrimitive(OOCStreamable<I> input, OOCStreamable<O> output, Function<I, O> partial,
+ BiFunction<O, O, O> merge, ToLongFunction<O> size, StreamContext context) {
+ super(context, input);
+ _input = input;
+ _output = output;
+ _partial = partial;
+ _merge = merge;
+ _size = size;
+ }
+
+ @Override
+ protected void inferPatternsInternal() {
+ _pattern = OOCAccessPattern.ANY;
+ inferParentPatterns();
+ }
+
+ @Override
+ protected void requestPatternInternal(OOCAccessPattern accessPattern) {
+ _pattern = OOCAccessPattern.ANY;
+ }
+
+ @Override
+ protected void startExecution() {
+ OOCStream<I> input = getInputReadStream(0);
+ OOCStream<O> output = _output.getWriteStream();
+ long inputBytes = OOCUtils.estimateOutputTileBytes(_input.getDataCharacteristics());
+ long outputBytes = OOCUtils.estimateOutputTileBytes(_output.getDataCharacteristics());
+ long taskBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(inputBytes) + 2 * outputBytes;
+ AllocatedOOCStream<I> admitted = new AllocatedOOCStream<>(input, _allowance, ignored -> taskBytes);
+ getContext().addOutStream(output);
+ OOCInstructionUtils.submitOOCTasks(admitted, callback -> {
+ ReservationBudget budget = null;
+ ManagedPayload<O> partial = null;
+ try {
+ budget = AllocatedOOCStream.detachBudget(callback).enableReuse();
+ O value = _partial.apply(callback.get());
+ long bytes = _size.applyAsLong(value);
+ budget.reserveBlocking(bytes);
+ partial = new ManagedPayload<>(value, bytes, budget);
+ synchronized(this) {
+ if(_accumulator != null) {
+ value = _merge.apply(_accumulator.value(), partial.value());
+ bytes = _size.applyAsLong(value);
+ budget.reserveBlocking(bytes);
+ _accumulator.release();
+ partial.release();
+ partial = new ManagedPayload<>(value, bytes, budget);
+ }
+ _accumulator = partial;
+ partial = null;
+ budget.close();
+ budget = null;
+ }
+ }
+ catch(Throwable error) {
+ fail(error);
+ throw DMLRuntimeException.of(error);
+ }
+ finally {
+ if(partial != null)
+ partial.release();
+ if(budget != null)
+ budget.close();
+ }
+ }, getContext()).thenRun(() -> {
+ ManagedPayload<O> result;
+ synchronized(this) {
+ result = _accumulator;
+ _accumulator = null;
+ }
+ try {
+ if(hasFailed()) {
+ if(result != null)
+ result.release();
+ return;
+ }
+ if(result == null)
+ throw new DMLRuntimeException("Cannot reduce an empty OOC stream");
+ OOCStream.QueueCallback<O> callback = new InMemoryQueueCallback<>(result);
+ try {
+ output.enqueue(callback);
+ callback = null;
+ }
+ finally {
+ if(callback != null)
+ callback.close();
+ }
+ output.closeInput();
+ }
+ catch(Throwable error) {
+ if(result != null)
+ result.release();
+ fail(error);
+ }
+ finally {
+ onComplete();
+ }
+ });
+ }
+}
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 4f2bcc3485..3eedb32b8e 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
@@ -31,6 +31,7 @@ import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.function.ToIntFunction;
+import java.util.function.ToLongFunction;
import org.apache.sysds.api.DMLScript;
import org.apache.sysds.runtime.DMLRuntimeException;
@@ -47,6 +48,7 @@ import org.apache.sysds.runtime.ooc.primitives.GroupedReduceOOCPrimitive;
import org.apache.sysds.runtime.ooc.primitives.JoinOOCPrimitive;
import org.apache.sysds.runtime.ooc.primitives.MappingOOCPrimitive;
import org.apache.sysds.runtime.ooc.primitives.PlannableDataGenOOCPrimitive;
+import org.apache.sysds.runtime.ooc.primitives.ReduceOOCPrimitive;
import org.apache.sysds.runtime.ooc.primitives.TransposeOOCPrimitive;
import org.apache.sysds.runtime.ooc.stats.OOCEventLog;
import org.apache.sysds.runtime.ooc.store.MaterializedStore;
@@ -114,6 +116,11 @@ public final class OOCInstructionUtils {
output.assignPrimitive(new GroupedReduceOOCPrimitive(input, output, grouping, partial, merge, finish, context));
}
+ public static <I, O> void reduce(OOCStreamable<I> input, OOCStream<O> output, Function<I, O> partial,
+ BiFunction<O, O, O> merge, ToLongFunction<O> size, StreamContext context) {
+ output.assignPrimitive(new ReduceOOCPrimitive<>(input, output, partial, merge, size, context));
+ }
+
public static int getComputeInFlight() {
return COMPUTE_IN_FLIGHT.get();
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java
index 7981736753..c1a0e6439c 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java
@@ -163,7 +163,7 @@ public class OOCUtils {
OOCStream.QueueCallback<IndexedMatrixValue> callback = null;
try {
budget.reserveBlocking(bytes);
- callback = new InMemoryQueueCallback(value, null, budget, bytes);
+ callback = new InMemoryQueueCallback<>(value, null, budget, bytes);
budget.close();
out.enqueue(callback);
callback = null;
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 c94e0ac5f1..336a7aecfd 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
@@ -188,6 +188,30 @@ public class OOCPrimitiveTest {
values);
}
+ @Test
+ public void testReduce() {
+ SubscribableTaskQueue<IndexedMatrixValue> input = new SubscribableTaskQueue<>();
+ SubscribableTaskQueue<MatrixBlock> 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(1, 1, 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.reduce(input, output, value -> new MatrixBlock(1, 1, 2 * value.getValue().get(0, 0)),
+ (left, right) -> new MatrixBlock(1, 1, left.get(0, 0) + right.get(0, 0)),
+ MatrixBlock::getExactSerializedSize, new StreamContext());
+
+ output.start();
+ try(OOCStream.QueueCallback<MatrixBlock> callback = output.dequeueCB()) {
+ Assert.assertEquals(204, callback.get().get(0, 0), 0);
+ }
+ Assert.assertNull(output.dequeueCB());
+ }
+
@Test
public void testGroupedReduceModes() {
Assert.assertEquals(Map.of("1,1", 136d, "2,1", 166d),
diff --git a/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java
index 5ecbb070b4..7e75a4d52e 100644
--- a/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java
+++ b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java
@@ -84,7 +84,7 @@ public class StateTableUtilsTest {
_producer.reserveBlocking(TILE_BYTES);
StateTableUtils.Match referenced = StateTableUtils
- .putOrTake(_table, 0, new InMemoryQueueCallback(tile(2.0), null, _producer, TILE_BYTES), _reader)
+ .putOrTake(_table, 0, new InMemoryQueueCallback<>(tile(2.0), null, _producer, TILE_BYTES), _reader)
.get(WAIT_SECONDS, TimeUnit.SECONDS);
Assert.assertNotNull(referenced);
try(OOCStream.QueueCallback<IndexedMatrixValue> left = referenced.left();
@@ -95,7 +95,7 @@ public class StateTableUtilsTest {
_producer.reserveBlocking(TILE_BYTES);
Assert.assertNull(StateTableUtils
- .putOrTake(_table, 1, new InMemoryQueueCallback(tile(3.0), null, _producer, TILE_BYTES), _reader)
+ .putOrTake(_table, 1, new InMemoryQueueCallback<>(tile(3.0), null, _producer, TILE_BYTES), _reader)
.get(WAIT_SECONDS, TimeUnit.SECONDS));
StateTableUtils.Match copied = StateTableUtils
.putOrTake(_table, 1, new OOCStream.SimpleQueueCallback<>(tile(4.0), null), _reader)
diff --git a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
index 8660252e63..5986f6521d 100644
--- a/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
+++ b/src/test/java/org/apache/sysds/test/component/ooc/memory/OOCMemoryAllowanceTest.java
@@ -223,7 +223,7 @@ public class OOCMemoryAllowanceTest {
OOCStream<Integer> leftStream = new SubscribableTaskQueue<>();
OOCStream<Integer> rightStream = new SubscribableTaskQueue<>();
- OOCStream<InMemoryQueueCallback> outStream = new SubscribableTaskQueue<>();
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> outStream = new SubscribableTaskQueue<>();
long startMillis = System.currentTimeMillis();
@@ -248,31 +248,31 @@ public class OOCMemoryAllowanceTest {
rightStream.closeInput();
}).start();
- OOCStream<InMemoryQueueCallback> leftStreamOut = new SubscribableTaskQueue<>();
- OOCStream<InMemoryQueueCallback> leftStreamOutOut = new SubscribableTaskQueue<>();
- OOCStream<InMemoryQueueCallback> rightStreamOut = new SubscribableTaskQueue<>();
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> leftStreamOut = new SubscribableTaskQueue<>();
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> leftStreamOutOut = new SubscribableTaskQueue<>();
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> rightStreamOut = new SubscribableTaskQueue<>();
test.map(leftStream, leftStreamOut, i -> {
var imv = new IndexedMatrixValue(new MatrixIndexes(i.longValue(), 1L), new MatrixBlock(1000, 1, 5.0));
- return new InMemoryQueueCallback(imv, null, leftAllowance, 8 * 1000);
+ return new InMemoryQueueCallback<>(imv, null, leftAllowance, 8 * 1000);
});
test.map(leftStreamOut, leftStreamOutOut, cb -> {
try(cb) {
var imv = new IndexedMatrixValue(cb.get().getIndexes(), cb.get().getValue()
.scalarOperations(new RightScalarOperator(Plus.getPlusFnObject(), 2.0), new MatrixBlock()));
- return new InMemoryQueueCallback(imv, null, leftAllowance, 8 * 1000);
+ return new InMemoryQueueCallback<>(imv, null, leftAllowance, 8 * 1000);
}
});
test.map(rightStream, rightStreamOut, i -> {
var imv = new IndexedMatrixValue(new MatrixIndexes(i.longValue(), 1L), new MatrixBlock(1000, 1, 3.0));
- return new InMemoryQueueCallback(imv, null, rightAllowance, 8 * 1000);
+ return new InMemoryQueueCallback<>(imv, null, rightAllowance, 8 * 1000);
});
test.join(leftStreamOutOut, rightStreamOut, outStream, () -> joinAllowance.reserveBlocking(8 * 1000), cache,
(l, r) -> {
var imv = new IndexedMatrixValue(l.getIndexes(), ((MatrixBlock)l.getValue()).binaryOperations(new BinaryOperator(
Plus.getPlusFnObject()), r.getValue()));
- return new InMemoryQueueCallback(imv, null, joinAllowance, 8 * 1000);
+ return new InMemoryQueueCallback<>(imv, null, joinAllowance, 8 * 1000);
});
CompletableFuture<Void> future = new CompletableFuture<>();
@@ -283,7 +283,7 @@ public class OOCMemoryAllowanceTest {
future.complete(null);
return;
}
- InMemoryQueueCallback inner = cb.get();
+ InMemoryQueueCallback<IndexedMatrixValue> inner = cb.get();
try(cb; inner) {
ctr.incrementAndGet();
double checksum =((MatrixBlock)inner.get().getValue()).sum();
@@ -394,14 +394,15 @@ public class OOCMemoryAllowanceTest {
return super.joinOOC(l, r, out, joinFn, IndexedMatrixValue::getIndexes);
}
- public CompletableFuture<Void> join(OOCStream<InMemoryQueueCallback> l, OOCStream<InMemoryQueueCallback> r,
- OOCStream<InMemoryQueueCallback> out, Runnable memoryReserver, CachedAllowance cache,
- BiFunction<IndexedMatrixValue, IndexedMatrixValue, InMemoryQueueCallback> joinFn) {
+ public CompletableFuture<Void> join(OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> l,
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> r,
+ OOCStream<InMemoryQueueCallback<IndexedMatrixValue>> out, Runnable memoryReserver, CachedAllowance cache,
+ BiFunction<IndexedMatrixValue, IndexedMatrixValue, InMemoryQueueCallback<IndexedMatrixValue>> joinFn) {
OOCStream<Tuple3<OOCStream.QueueCallback<IndexedMatrixValue>, OOCStream.QueueCallback<IndexedMatrixValue>, Integer>> intermediate = createWritableStream();
new Thread(() -> {
- InMemoryQueueCallback next;
+ InMemoryQueueCallback<IndexedMatrixValue> next;
IndexedMatrixValue nextValue;
boolean nextLeft = true;
AtomicInteger pendingRequests = new AtomicInteger(1);