diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/MMultOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/MMultOOCInstruction.java index 81f1102811d..d176a0c4184 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/MMultOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/MMultOOCInstruction.java @@ -33,6 +33,9 @@ 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.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.store.CountingLiveness; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class MMultOOCInstruction extends ComputationOOCInstruction { @@ -60,6 +63,33 @@ public void processInstruction( ExecutionContext ec ) { // 1. Identify the inputs MatrixObject min = ec.getMatrixObject(input1); // big matrix MatrixObject vin = ec.getMatrixObject(input2); // streamed vector + DataCharacteristics mdc = min.getDataCharacteristics(); + DataCharacteristics vdc = vin.getDataCharacteristics(); + + if(min != vin && mdc.getRows() > 0 && mdc.getCols() > 0 && vdc.getCols() > 0 && + mdc.getCols() == vdc.getRows() && vdc.getNumColBlocks() == 1) { + OOCStream partials = createWritableStream(); + OOCStream out = createWritableStream(); + partials.setData(min); + ec.getMatrixObject(output).setStreamHandle(out); + OOCInstructionUtils.indexedBroadcastMap(min.getStreamable(), vin.getStreamable(), partials, + left -> Math.toIntExact(left.getIndexes().getColumnIndex() - 1), + () -> new CountingLiveness(Math.toIntExact(vin.getDataCharacteristics().getNumRowBlocks()), + Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks())), + (left, right) -> { + MatrixBlock leftBlock = (MatrixBlock) left.getValue(); + MatrixBlock rightBlock = (MatrixBlock) right.getValue(); + MatrixBlock partial = leftBlock.aggregateBinaryOperations(leftBlock, rightBlock, new MatrixBlock(), + (AggregateBinaryOperator) _optr); + MatrixIndexes indexes = left.getIndexes(); + return new IndexedMatrixValue(new MatrixIndexes(indexes.getRowIndex(), indexes.getColumnIndex()), + partial); + }, getContext()); + BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString()); + OOCInstructionUtils.rowGroupedReduce(partials, out, + (left, right) -> left.binaryOperations(plus, right, new MatrixBlock()), getContext()); + return; + } int emitLeftThreshold = (int)vin.getDataCharacteristics().getNumColBlocks(); int emitRightThreshold = (int)min.getDataCharacteristics().getNumRowBlocks(); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java index aec88891595..8bbdb0e24de 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java @@ -24,6 +24,14 @@ import java.util.function.LongUnaryOperator; public interface OOCCache { + /** + * Maximum bytes charged given the logical bytes of the requested entry. Use this method for reservation budget + * planning as logical byte size and pinned entry bytes may differ. + */ + default long maxPhysicalPinBytes(long logicalBytes) { + return logicalBytes; + } + /** * Pins an item backed by an allowance. A successful pin transfers memory ownership from the cache to the owner of * the allowance and guarantees data availability. While pinned, the bytes of the entry are not counted as diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java index 9bbdcbeae5a..ad08c637bcf 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java @@ -129,6 +129,14 @@ public OOCPackedCache(OOCCacheImpl physical, long packThresholdBytes, long packT }); } + @Override + public long maxPhysicalPinBytes(long logicalBytes) { + if(logicalBytes >= _packThresholdBytes) + return logicalBytes; + return _packTargetBytes > Long.MAX_VALUE - _packThresholdBytes ? Long.MAX_VALUE : _packTargetBytes + + _packThresholdBytes; + } + @Override public BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance) { if(size >= _packThresholdBytes) diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java index 47e08202dfb..26bf4477f6e 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/ReservationBudget.java @@ -26,6 +26,7 @@ public final class ReservationBudget implements MemoryAllowance, AutoCloseable { private long _outstanding; private long _available; private boolean _closed; + private boolean _reusable; public ReservationBudget(MemoryAllowance parent, long bytes) { if(parent == null) @@ -37,6 +38,13 @@ public ReservationBudget(MemoryAllowance parent, long bytes) { _available = bytes; } + public synchronized ReservationBudget enableReuse() { + if(_closed || getUsedMemory() != 0) + throw new IllegalStateException("Budget reuse must be enabled before reserving memory."); + _reusable = true; + return this; + } + @Override public synchronized boolean tryReserve(long bytes) { checkNonNegative(bytes); @@ -64,13 +72,19 @@ public void release(long bytes) { checkNonNegative(bytes); if(bytes == 0) return; + boolean releaseParent; synchronized(this) { long used = _outstanding - _available; if(bytes > used) throw new IllegalStateException("Cannot release " + bytes + " bytes from a budget using " + used); - _outstanding -= bytes; + releaseParent = _closed || !_reusable; + if(releaseParent) + _outstanding -= bytes; + else + _available += bytes; } - _parent.release(bytes); + if(releaseParent) + _parent.release(bytes); } @Override diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java new file mode 100644 index 00000000000..e2a9be63c25 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/BroadcastOOCPrimitive.java @@ -0,0 +1,318 @@ +/* + * 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.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import java.util.function.Supplier; +import java.util.function.ToIntFunction; + +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.instructions.ooc.SubscribableTaskQueue; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.memory.ReservationBudget; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.store.IndexedMaterializedStoreReader; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; +import org.apache.sysds.runtime.ooc.store.StoreLease; +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 BroadcastOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _broadcast; + private final OOCStreamable _output; + private final ToIntFunction _lookup; + private final Supplier _liveness; + private final BiFunction _operation; + private final AtomicBoolean _cleaned; + private final AtomicBoolean _failed; + private final AtomicBoolean _sourceComplete; + private final AtomicInteger _active; + private MaterializedStore _store; + private IndexedMaterializedStoreReader _reader; + private OOCStream _ready; + private OOCStream _outputStream; + + public BroadcastOOCPrimitive(OOCStreamable streamed, + OOCStreamable broadcast, OOCStreamable output, + ToIntFunction lookup, Supplier liveness, + BiFunction operation, StreamContext context) { + super(context, streamed, broadcast); + _broadcast = broadcast; + _output = output; + _lookup = lookup; + _liveness = liveness; + _operation = operation; + _cleaned = new AtomicBoolean(); + _failed = new AtomicBoolean(); + _sourceComplete = new AtomicBoolean(); + _active = new AtomicInteger(1); + } + + @Override + public List requiredMaterializedInputs() { + return List.of(new OOCMaterializedInputRequest(1, OOCStoreLayout.ROW_MAJOR, 1)); + } + + @Override + protected void inferPatternsInternal() { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + } + + @Override + protected void startExecution() { + _outputStream = _output.getWriteStream(); + _ready = new SubscribableTaskQueue<>(); + getContext().addOutStream(_outputStream, _ready); + OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext()) + .whenComplete((ignored, error) -> { + try { + if(error != null) + fail(error); + _outputStream.closeInput(); + } + catch(Throwable failure) { + fail(failure); + } + finally { + cleanup(); + } + }); + + getMaterializedInput(1).whenComplete((store, error) -> { + if(error != null) { + fail(error); + finishSource(); + return; + } + _store = store; + store.completion().whenComplete((ignored, completionError) -> { + if(completionError != null) { + fail(completionError); + finishSource(); + return; + } + try { + _reader = store.openIndexedReader(_liveness.get()); + startBroadcast(); + } + catch(Throwable failure) { + fail(failure); + finishSource(); + } + }); + }); + } + + private void startBroadcast() { + long broadcastLogical = OOCUtils.estimateFullTileBytes(_broadcast.getDataCharacteristics()); + long outputLogical = OOCUtils.estimateFullTileBytes(_output.getDataCharacteristics()); + long broadcastPin = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(broadcastLogical); + long taskBytes = broadcastPin * 2 + outputLogical * 2; + OOCStream streamed = getInputReadStream(0); + AllocatedOOCStream admitted = new AllocatedOOCStream<>(streamed, _allowance, + ignored -> taskBytes); + getContext().addInStream(streamed, admitted); + admitted.setSubscriber(this::accept); + } + + private void accept(OOCStream.QueueCallback callback) { + if(callback.isEos() || callback.isFailure()) { + try(callback) { + if(callback.isFailure()) + callback.get(); + } + catch(Throwable failure) { + fail(failure); + } + finishSource(); + return; + } + + ReservationBudget budget = null; + OOCStream.QueueCallback retained = null; + _active.incrementAndGet(); + try(callback) { + budget = AllocatedOOCStream.detachBudget(callback); + if(budget == null) + throw new DMLRuntimeException("Missing admitted broadcast task budget."); + IndexedMatrixValue streamed = callback.get(); + int lookup = _lookup.applyAsInt(streamed); + retained = callback.keepOpen(); + OOCFuture> requested = _reader.request(lookup, budget); + OOCStream.QueueCallback pendingStreamed = retained; + ReservationBudget pendingBudget = budget; + retained = null; + budget = null; + requested.whenComplete( + (broadcast, error) -> broadcastReady(pendingStreamed, broadcast, pendingBudget, lookup, error)); + } + catch(Throwable failure) { + fail(failure); + completeOne(); + } + finally { + if(retained != null) + retained.close(); + if(budget != null) + budget.close(); + } + } + + private void broadcastReady(OOCStream.QueueCallback streamed, + StoreLease broadcast, ReservationBudget budget, int lookup, Throwable error) { + if(error != null || broadcast == null) { + try { + streamed.close(); + if(broadcast != null) + broadcast.close(); + budget.close(); + } + finally { + fail(error != null ? error : new IllegalStateException("Missing broadcast tile " + lookup)); + completeOne(); + } + return; + } + BroadcastWork work = new BroadcastWork(streamed, broadcast, budget); + try { + _ready.enqueue(work); + } + catch(Throwable failure) { + work.close(); + fail(failure); + completeOne(); + } + } + + private void process(BroadcastWork work) { + ReservationBudget budget = work.takeBudget(); + try { + IndexedMatrixValue output = _operation.apply(work._streamed.get(), work._broadcast.value()); + OOCUtils.enqueueExact(_outputStream, output, budget); + budget = null; + } + catch(Throwable failure) { + fail(failure); + } + finally { + work.close(); + if(budget != null) + budget.close(); + completeOne(); + } + } + + private void finishSource() { + if(_sourceComplete.compareAndSet(false, true)) + completeOne(); + } + + private void completeOne() { + if(_active.decrementAndGet() != 0) + return; + try { + _ready.closeInput(); + } + catch(IllegalStateException ignored) { + // Failure propagation may already have closed the ready stream. + } + } + + private void fail(Throwable error) { + if(!_failed.compareAndSet(false, true)) + return; + DMLRuntimeException failure = DMLRuntimeException.of(error); + _outputStream.propagateFailure(failure); + getContext().failAll(failure); + } + + private void cleanup() { + if(!_cleaned.compareAndSet(false, true)) + return; + try { + if(_reader != null) + _reader.close(); + } + finally { + try { + if(_store != null) + _store.close(); + } + finally { + onComplete(); + } + } + } + + private static final class BroadcastWork implements AutoCloseable { + private OOCStream.QueueCallback _streamed; + private StoreLease _broadcast; + private ReservationBudget _budget; + + private BroadcastWork(OOCStream.QueueCallback streamed, + StoreLease broadcast, ReservationBudget budget) { + _streamed = streamed; + _broadcast = broadcast; + _budget = budget; + } + + private ReservationBudget takeBudget() { + ReservationBudget budget = _budget; + _budget = null; + return budget; + } + + @Override + public void close() { + if(_streamed != null) { + _streamed.close(); + _streamed = null; + } + if(_broadcast != null) { + _broadcast.close(); + _broadcast = null; + } + if(_budget != null) { + _budget.close(); + _budget = null; + } + } + } +} 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 new file mode 100644 index 00000000000..ad126af0c71 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/GroupedReduceOOCPrimitive.java @@ -0,0 +1,359 @@ +/* + * 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.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; + +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.instructions.ooc.CachingStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +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.store.StateTable; +import org.apache.sysds.runtime.ooc.store.StoreLease; +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 GroupedReduceOOCPrimitive extends OOCPrimitive { + private final OOCStream _input; + private final OOCStreamable _output; + private final BiFunction _merge; + private final AtomicBoolean _cleaned; + private final AtomicBoolean _failed; + private final AtomicBoolean _sourceComplete; + private final AtomicInteger _active; + private final AtomicInteger _finalizedGroups; + private StateTable _table; + private OOCStream _ready; + private OOCStream _outputStream; + private int _numGroups; + private int _groupSize; + + public GroupedReduceOOCPrimitive(OOCStreamable input, OOCStreamable output, + BiFunction merge, StreamContext context) { + this(input.getReadStream(), output, merge, context); + } + + private GroupedReduceOOCPrimitive(OOCStream input, OOCStreamable output, + BiFunction merge, StreamContext context) { + super(context, input.getPrimitive() == null ? List.of() : List.of(input.getPrimitive())); + _input = input; + _output = output; + _merge = merge; + _cleaned = new AtomicBoolean(); + _failed = new AtomicBoolean(); + _sourceComplete = new AtomicBoolean(); + _active = new AtomicInteger(1); + _finalizedGroups = new AtomicInteger(); + } + + @Override + protected void inferPatternsInternal() { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = OOCAccessPattern.ROW_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(OOCAccessPattern.ROW_MAJOR); + } + + @Override + protected void startExecution() { + DataCharacteristics inputDc = _input.getDataCharacteristics(); + if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) + throw new DMLRuntimeException("Grouped OOC reduction requires known input dimensions and block size."); + _numGroups = Math.toIntExact(inputDc.getNumRowBlocks()); + _groupSize = Math.toIntExact(inputDc.getNumColBlocks()); + _outputStream = _output.getWriteStream(); + _ready = new SubscribableTaskQueue<>(); + getContext().addInStream(_input).addOutStream(_outputStream, _ready); + _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); + + OOCInstructionUtils.submitOOCTasks(_ready, callback -> process(callback.get()), getContext()) + .whenComplete((ignored, error) -> { + try { + _outputStream.closeInput(); + } + catch(Throwable failure) { + fail(failure); + } + finally { + cleanup(); + } + }); + + long logicalBytes = Math.max(OOCUtils.estimateFullTileBytes(inputDc), + OOCUtils.estimateFullTileBytes(_output.getDataCharacteristics())); + long pinBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(logicalBytes); + long taskBytes = pinBytes + logicalBytes * 2; + AllocatedOOCStream admitted = new AllocatedOOCStream<>(_input, _allowance, + ignored -> taskBytes); + getContext().addInStream(admitted); + admitted.setSubscriber(this::accept); + } + + private void accept(OOCStream.QueueCallback callback) { + if(callback.isEos() || callback.isFailure()) { + try(callback) { + if(callback.isFailure()) + callback.get(); + } + catch(Throwable failure) { + fail(failure); + } + finishSource(); + return; + } + + ReservationBudget budget = null; + ManagedPayload payload = null; + _active.incrementAndGet(); + 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); + reduce(group, payload, budget); + payload = null; + budget = null; + } + catch(Throwable failure) { + fail(failure); + completeOne(); + } + finally { + if(payload != null) + payload.release(); + if(budget != null) + budget.close(); + } + } + + private void reduce(int group, ManagedPayload incoming, ReservationBudget budget) { + if(multiplicity(incoming.value()) == _groupSize) { + finalizeGroup(group, incoming, budget); + return; + } + OOCFuture> match; + try { + match = _table.putOrTake(group, incoming, budget); + } + catch(Throwable failure) { + incoming.release(); + budget.close(); + fail(failure); + completeOne(); + return; + } + match.whenComplete((existing, error) -> { + if(error != null) { + incoming.release(); + budget.close(); + fail(error); + completeOne(); + } + else if(existing == null) { + budget.close(); + completeOne(); + } + else { + MergeWork work = new MergeWork(group, incoming, existing, budget); + try { + _ready.enqueue(work); + } + catch(Throwable failure) { + work.close(); + fail(failure); + completeOne(); + } + } + }); + } + + private void process(MergeWork work) { + ReservationBudget budget = work.takeBudget(); + ManagedPayload merged = null; + OOCFuture released; + try { + IndexedMatrixValue left = work._existing.value(); + 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)); + MatrixBlock value = _merge.apply((MatrixBlock) left.getValue(), (MatrixBlock) right.getValue()); + merged = payload(new IndexedMatrixValue(new MatrixIndexes(work._group + 1L, count), value), budget); + work.releaseIncoming(); + released = work.closeExistingAsync(); + } + catch(Throwable failure) { + if(merged != null) + merged.release(); + work.close(); + budget.close(); + fail(failure); + completeOne(); + return; + } + + ManagedPayload next = merged; + released.whenComplete((ignored, error) -> { + if(error != null) { + next.release(); + budget.close(); + fail(error); + completeOne(); + } + else + reduce(work._group, next, budget); + }); + } + + 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); + _finalizedGroups.incrementAndGet(); + } + catch(Throwable failure) { + budget.close(); + fail(failure); + } + completeOne(); + } + + private static ManagedPayload payload(IndexedMatrixValue value, ReservationBudget budget) { + long bytes = ((MatrixBlock) value.getValue()).getExactSerializedSize(); + budget.reserveBlocking(bytes); + return new ManagedPayload<>(value, bytes, budget); + } + + private static int multiplicity(IndexedMatrixValue value) { + return Math.toIntExact(value.getIndexes().getColumnIndex()); + } + + private void finishSource() { + if(_sourceComplete.compareAndSet(false, true)) + completeOne(); + } + + private void completeOne() { + int remaining = _active.decrementAndGet(); + if(remaining != 0) + return; + if(!_failed.get() && _finalizedGroups.get() != _numGroups) + fail(new DMLRuntimeException( + "Grouped reduction completed " + _finalizedGroups.get() + " of " + _numGroups + " row groups.")); + try { + _ready.closeInput(); + } + catch(IllegalStateException ignored) { + // Failure propagation may already have closed the ready stream. + } + } + + private void fail(Throwable error) { + if(!_failed.compareAndSet(false, true)) + return; + DMLRuntimeException failure = DMLRuntimeException.of(error); + _outputStream.propagateFailure(failure); + getContext().failAll(failure); + } + + private void cleanup() { + if(!_cleaned.compareAndSet(false, true)) + return; + try { + if(_table != null) + _table.close(); + } + finally { + onComplete(); + } + } + + private static final class MergeWork implements AutoCloseable { + private final int _group; + private ManagedPayload _incoming; + private StoreLease _existing; + private ReservationBudget _budget; + + private MergeWork(int group, ManagedPayload incoming, + StoreLease existing, ReservationBudget budget) { + _group = group; + _incoming = incoming; + _existing = existing; + _budget = budget; + } + + private ReservationBudget takeBudget() { + ReservationBudget budget = _budget; + _budget = null; + return budget; + } + + private void releaseIncoming() { + if(_incoming != null) { + _incoming.release(); + _incoming = null; + } + } + + private OOCFuture closeExistingAsync() { + StoreLease existing = _existing; + _existing = null; + return existing.closeAsync(); + } + + @Override + public void close() { + releaseIncoming(); + if(_existing != null) { + _existing.close(); + _existing = null; + } + if(_budget != null) + _budget.close(); + } + } +} 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 773ef9da834..f86c2360608 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 @@ -29,6 +29,8 @@ import java.util.function.BiFunction; import java.util.function.Consumer; import java.util.function.Function; +import java.util.function.Supplier; +import java.util.function.ToIntFunction; import org.apache.sysds.api.DMLScript; import org.apache.sysds.runtime.DMLRuntimeException; @@ -40,11 +42,14 @@ import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; import org.apache.sysds.runtime.ooc.memory.ReservationBudget; +import org.apache.sysds.runtime.ooc.primitives.BroadcastOOCPrimitive; +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.TransposeOOCPrimitive; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream; import org.apache.sysds.runtime.ooc.stream.StreamContext; import org.apache.sysds.runtime.ooc.stream.TaskContext; @@ -88,6 +93,19 @@ public static void equiJoin(OOCStreamable left, OOCStreamabl output.assignPrimitive(new JoinOOCPrimitive(left, right, output, operation, context)); } + public static void indexedBroadcastMap(OOCStreamable streamed, + OOCStreamable broadcast, OOCStream output, + ToIntFunction lookup, Supplier liveness, + BiFunction operation, StreamContext context) { + output.assignPrimitive( + new BroadcastOOCPrimitive(streamed, broadcast, output, lookup, liveness, operation, context)); + } + + public static void rowGroupedReduce(OOCStreamable input, OOCStream output, + BiFunction merge, StreamContext context) { + output.assignPrimitive(new GroupedReduceOOCPrimitive(input, output, merge, 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 5a98a05b78d..7981736753f 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 @@ -141,13 +141,17 @@ public MatrixIndexes next() { public static long estimateOutputTileBytes(DataCharacteristics dc) { if(dc == null || dc.getBlocksize() <= 0 || !dc.dimsKnown()) { - int blocksize = dc != null && dc.getBlocksize() > 0 ? dc.getBlocksize() : 1000; - return estimateMatrixBlockBytes(blocksize, blocksize); + return estimateFullTileBytes(dc); } return estimateMatrixBlockBytes(Math.min(dc.getBlocksize(), dc.getRows()), Math.min(dc.getBlocksize(), dc.getCols())); } + public static long estimateFullTileBytes(DataCharacteristics dc) { + int blocksize = dc != null && dc.getBlocksize() > 0 ? dc.getBlocksize() : 1000; + return estimateMatrixBlockBytes(blocksize, blocksize); + } + private static long estimateMatrixBlockBytes(long rows, long cols) { return Math.max(MatrixBlock.estimateSizeDenseInMemory(rows, cols), MatrixBlock.estimateSizeSparseInMemory(rows, cols, 1.0)); 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 c75458eef38..8660252e634 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 @@ -138,6 +138,15 @@ public void testAllocatedStreamReservations() { budget.close(); Assert.assertEquals(0, allowance.getUsedMemory()); + allowance.reserveBlocking(60); + ReservationBudget reusable = new ReservationBudget(allowance, 60).enableReuse(); + reusable.reserveBlocking(40); + reusable.release(40); + reusable.reserveBlocking(40); + reusable.release(40); + reusable.close(); + Assert.assertEquals(0, allowance.getUsedMemory()); + source.enqueue(2); OOCStream.QueueCallback second = allocated.dequeueCB(); OOCStream.QueueCallback retained = second.keepOpen();