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 37b2ba93a77..0707601b12f 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.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.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 void processInstruction(ExecutionContext ec) { } private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObject min) { - OOCStream qIn = min.getStreamHandle(); + OOCStream out = createWritableStream(); + ec.getMatrixObject(output).setStreamHandle(out); BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString()); - MatrixBlock resultBlock = null; - - OOCStream 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 5cec3ae981d..4bde5c9df4d 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 static boolean canClaimMemory() { return getCache().isWithinLimits() && OOCInstruction.getComputeInFlight() <= OOCInstruction.getComputeBackpressureThreshold(); } - public static OOCCacheScheduler.HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback) { + public static OOCCacheScheduler.HandoverHandle handover(BlockKey key, + InMemoryQueueCallback 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 ad161e95303..d8ff79dd797 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 @@ interface HandoverHandle { OOCStream.QueueCallback reclaim(); } - HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback); + HandoverHandle handover(BlockKey key, InMemoryQueueCallback 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 3f5601adbae..96de368ccc9 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 BlockEntry putAndPin(BlockKey key, Object data, long size) { } @Override - public HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback) { + public HandoverHandle handover(BlockKey key, InMemoryQueueCallback callback) { if(!this._running) throw new IllegalStateException("Cache scheduler has been shut down."); PendingHandover handover = new PendingHandover(key, callback); @@ -1085,7 +1085,7 @@ private void registerWaiter(BlockKey key, DeferredReadRequest request, int index } private boolean commitHandover(PendingHandover pending) { - InMemoryQueueCallback callback = pending.takeForCommit(); + InMemoryQueueCallback callback = pending.takeForCommit(); if(callback == null) return false; try { @@ -1135,12 +1135,12 @@ private DeferredReadWaiter(DeferredReadRequest request, int index) { private static class PendingHandover implements HandoverHandle { private final BlockKey _key; private final CompletableFuture _completionFuture; - private InMemoryQueueCallback _callback; + private InMemoryQueueCallback _callback; private boolean _committed; private boolean _cancelled; private boolean _committing; - private PendingHandover(BlockKey key, InMemoryQueueCallback callback) { + private PendingHandover(BlockKey key, InMemoryQueueCallback callback) { _key = key; _completionFuture = new CompletableFuture<>(); _callback = callback; @@ -1180,11 +1180,11 @@ private synchronized boolean isCancelled() { return _cancelled; } - private synchronized InMemoryQueueCallback takeForCommit() { + private synchronized InMemoryQueueCallback takeForCommit() { if(_committed || _cancelled || _committing) return null; _committing = true; - InMemoryQueueCallback callback = _callback; + InMemoryQueueCallback 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 ffba3910b26..d2375b80c1e 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 CachedAllowance(MemoryBroker broker) { _handoverSchedulingRequested = false; } - public void handover(InMemoryQueueCallback callback, int index) { + public void handover(InMemoryQueueCallback callback, int index) { if(callback == null) throw new IllegalArgumentException("Cannot hand over null callback."); callback.transferOwnershipBlocking(this); - InMemoryQueueCallback root = (InMemoryQueueCallback) callback.keepOpen(); + InMemoryQueueCallback root = callback.keepOpen(); callback.close(); root.getHandle().attachCachedAllowance(this, index); @@ -88,7 +88,7 @@ public OOCStream.QueueCallback tryGet(int index) { while(true) { BlockKey cacheKey = null; OOCCacheScheduler.HandoverHandle handover = null; - InMemoryQueueCallback local = null; + InMemoryQueueCallback local = null; synchronized(entry) { if(entry._local != null && entry._handover == null) @@ -124,7 +124,7 @@ else if(entry._cacheKey != null) if(!future.isDone()) return null; boolean committed = future.join(); - InMemoryQueueCallback localToClose = null; + InMemoryQueueCallback localToClose = null; synchronized(entry) { if(entry._handover != handover) continue; @@ -171,8 +171,8 @@ public CompletableFuture> get(int in throw DMLRuntimeException.of(ex.getCause() == null ? ex : ex.getCause()); return committed == true; }).thenCompose(committed -> { - InMemoryQueueCallback localToClose = null; - InMemoryQueueCallback local = null; + InMemoryQueueCallback localToClose = null; + InMemoryQueueCallback local = null; BlockKey key; synchronized(entry) { @@ -215,7 +215,7 @@ public void clear(int index) { while(true) { OOCCacheScheduler.HandoverHandle handover = null; BlockKey forgetKey = null; - InMemoryQueueCallback localToClose = null; + InMemoryQueueCallback localToClose = null; synchronized(entry) { if(entry._local != null && entry._handover == null) { @@ -410,7 +410,7 @@ private long tryStartCacheHandover(SlotEntry entry) { if(bytes <= 0) return 0; - InMemoryQueueCallback retained = (InMemoryQueueCallback) entry._local.keepOpen(); + InMemoryQueueCallback retained = entry._local.keepOpen(); try { entry._cacheKey = new BlockKey(_streamId, _nextBlockId.getAndIncrement()); entry._handover = OOCCacheManager.handover(entry._cacheKey, retained); @@ -448,7 +448,7 @@ private void finishPendingHandover(SlotEntry entry) { onFinishedHandover(bytes); } - private void closeRoot(InMemoryQueueCallback local) { + private void closeRoot(InMemoryQueueCallback local) { local.getHandle().detachCachedAllowance(); local.close(); } @@ -486,12 +486,12 @@ private void ensureCapacity(int index) { } private static final class SlotEntry { - private InMemoryQueueCallback _local; + private InMemoryQueueCallback _local; private BlockKey _cacheKey; private OOCCacheScheduler.HandoverHandle _handover; private long _pendingBytes; - private SlotEntry(InMemoryQueueCallback local) { + private SlotEntry(InMemoryQueueCallback 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 7fafc042e41..ae712305cb7 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 @@ 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 { - private CallbackHandle _handle; +public class InMemoryQueueCallback implements OOCStream.QueueCallback { + private CallbackHandle _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 payload) { + this(payload.value(), null, payload.owner(), payload.bytes()); + payload.transfer(); + } + + private InMemoryQueueCallback(CallbackHandle handle) { _handle = handle; _closed = false; } @Override - public IndexedMatrixValue get() { + public T get() { return _handle.get(); } @Override - public synchronized OOCStream.QueueCallback keepOpen() { + public synchronized InMemoryQueueCallback 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 boolean isFailure() { return _handle._failure != null; } - CallbackHandle getHandle() { + CallbackHandle getHandle() { return _handle; } - static final class CallbackHandle { - private volatile IndexedMatrixValue _result; + static final class CallbackHandle { + 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 @@ private CallbackHandle(IndexedMatrixValue result, DMLRuntimeException failure, M _cacheIdx = -1; } - private IndexedMatrixValue get() { + private T get() { if(_failure != null) throw _failure; return _result; @@ -174,8 +175,8 @@ boolean isExclusiveToRoot() { 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 @@ private void closeFinal() { } } - public IndexedMatrixValue takeManagedResultForHandover() { + public T takeManagedResultForHandover() { return _handle.takeManagedResultForHandover(); } - public synchronized ManagedPayload extractManagedPayload() { + public synchronized ManagedPayload extractManagedPayload() { if(_closed) throw new IllegalStateException("Cannot extract a managed payload from a closed callback."); - CallbackHandle handle = _handle; + CallbackHandle handle = _handle; synchronized(handle) { if(handle._failure != null) throw handle._failure; @@ -203,7 +204,7 @@ public synchronized ManagedPayload extractManagedPayload() { 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 00000000000..9f5ee85feda --- /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 extends OOCPrimitive { + private final OOCStreamable _input; + private final OOCStreamable _output; + private final Function _partial; + private final BiFunction _merge; + private final ToLongFunction _size; + private ManagedPayload _accumulator; + + public ReduceOOCPrimitive(OOCStreamable input, OOCStreamable output, Function partial, + BiFunction merge, ToLongFunction 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 input = getInputReadStream(0); + OOCStream output = _output.getWriteStream(); + long inputBytes = OOCUtils.estimateOutputTileBytes(_input.getDataCharacteristics()); + long outputBytes = OOCUtils.estimateOutputTileBytes(_output.getDataCharacteristics()); + long taskBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(inputBytes) + 2 * outputBytes; + AllocatedOOCStream admitted = new AllocatedOOCStream<>(input, _allowance, ignored -> taskBytes); + getContext().addOutStream(output); + OOCInstructionUtils.submitOOCTasks(admitted, callback -> { + ReservationBudget budget = null; + ManagedPayload 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 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 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 4f2bcc34853..3eedb32b8e6 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.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.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 static void groupedReduceIndexed(OOCStreamable input, output.assignPrimitive(new GroupedReduceOOCPrimitive(input, output, grouping, partial, merge, finish, context)); } + public static void reduce(OOCStreamable input, OOCStream output, Function partial, + BiFunction merge, ToLongFunction 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 7981736753f..c1a0e6439c1 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 static void enqueueExact(OOCStream out, IndexedMatrix OOCStream.QueueCallback 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 c94e0ac5f1d..336a7aecfd0 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 void testDataGenMapTransposePipeline() { values); } + @Test + public void testReduce() { + 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(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 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 5ecbb070b48..7e75a4d52ec 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 void testCallbackPutOrTake() throws Exception { _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 left = referenced.left(); @@ -95,7 +95,7 @@ public void testCallbackPutOrTake() throws Exception { _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 8660252e634..5986f6521db 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 long testNew(boolean optimal) { OOCStream leftStream = new SubscribableTaskQueue<>(); OOCStream rightStream = new SubscribableTaskQueue<>(); - OOCStream outStream = new SubscribableTaskQueue<>(); + OOCStream> outStream = new SubscribableTaskQueue<>(); long startMillis = System.currentTimeMillis(); @@ -248,31 +248,31 @@ public long testNew(boolean optimal) { rightStream.closeInput(); }).start(); - OOCStream leftStreamOut = new SubscribableTaskQueue<>(); - OOCStream leftStreamOutOut = new SubscribableTaskQueue<>(); - OOCStream rightStreamOut = new SubscribableTaskQueue<>(); + OOCStream> leftStreamOut = new SubscribableTaskQueue<>(); + OOCStream> leftStreamOutOut = new SubscribableTaskQueue<>(); + OOCStream> 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 future = new CompletableFuture<>(); @@ -283,7 +283,7 @@ public long testNew(boolean optimal) { future.complete(null); return; } - InMemoryQueueCallback inner = cb.get(); + InMemoryQueueCallback inner = cb.get(); try(cb; inner) { ctr.incrementAndGet(); double checksum =((MatrixBlock)inner.get().getValue()).sum(); @@ -394,14 +394,15 @@ public CompletableFuture joinOOC(OOCStream l, OOCStrea return super.joinOOC(l, r, out, joinFn, IndexedMatrixValue::getIndexes); } - public CompletableFuture join(OOCStream l, OOCStream r, - OOCStream out, Runnable memoryReserver, CachedAllowance cache, - BiFunction joinFn) { + public CompletableFuture join(OOCStream> l, + OOCStream> r, + OOCStream> out, Runnable memoryReserver, CachedAllowance cache, + BiFunction> joinFn) { OOCStream, OOCStream.QueueCallback, Integer>> intermediate = createWritableStream(); new Thread(() -> { - InMemoryQueueCallback next; + InMemoryQueueCallback next; IndexedMatrixValue nextValue; boolean nextLeft = true; AtomicInteger pendingRequests = new AtomicInteger(1);