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 0707601b12f..dfa348a2b1b 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 @@ -19,9 +19,6 @@ package org.apache.sysds.runtime.instructions.ooc; -import java.util.List; -import java.util.concurrent.CompletableFuture; - import org.apache.sysds.common.Opcodes; import org.apache.sysds.lops.MMTSJ; import org.apache.sysds.lops.MMTSJ.MMTSJType; @@ -33,7 +30,6 @@ import org.apache.sysds.runtime.instructions.InstructionUtils; import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; -import org.apache.sysds.runtime.matrix.data.LibMatrixReorg; import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.matrix.operators.AggregateBinaryOperator; @@ -67,6 +63,10 @@ public static TSMMOOCInstruction parseInstruction(String str) { @Override public void processInstruction(ExecutionContext ec) { MatrixObject min = ec.getMatrixObject(input1); + if(!min.getDataCharacteristics().dimsKnown() || min.getBlocksize() <= 0) + throw new DMLRuntimeException("OOC TSMM requires known dimensions and a positive block size: " + + min.getNumRows() + "x" + min.getNumColumns() + " (blocksize " + min.getBlocksize() + ")"); + int numRowBlocks = Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks()); int numColBlocks = Math.toIntExact(min.getDataCharacteristics().getNumColBlocks()); if((_type.isLeft() && numColBlocks == 1) || (_type.isRight() && numRowBlocks == 1)) { @@ -74,38 +74,10 @@ public void processInstruction(ExecutionContext ec) { return; } - int blocksPerJoinGroup = _type.isLeft() ? numColBlocks : numRowBlocks; - int partialsPerOutput = _type.isLeft() ? numRowBlocks : numColBlocks; - - OOCStreamable inputStreamable = min.getStreamable(); - final boolean createdCache = !inputStreamable.hasStreamCache(); - final CachingStream inputCache = createdCache ? new CachingStream(min.getStreamHandle()) : inputStreamable - .getStreamCache(); - - OOCStream> groupedPartials = createWritableStream(); - OOCStream partials = createWritableStream(); OOCStream out = createWritableStream(); - addOutStream(out); ec.getMatrixObject(output).setStreamHandle(out); - - CompletableFuture joinFuture = joinManyOOC(inputCache.getReadStream(), inputCache.getReadStream(), - groupedPartials, this::createPartialOutputTiles, this::getJoinIndex, this::getJoinIndex, blocksPerJoinGroup, - blocksPerJoinGroup); - CompletableFuture expandFuture = expandOOC(groupedPartials, partials, values -> values); - BinaryOperator plus = InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString()); - CompletableFuture outFuture = groupedReduceOOC(partials, out, (left, right) -> { - MatrixBlock result = ((MatrixBlock) left.getValue()).binaryOperations(plus, right.getValue()); - left.setValue(result); - return left; - }, partialsPerOutput); - - propagateFailuresToOutput(out, List.of(joinFuture, expandFuture, outFuture)); - - outFuture.whenComplete((result, error) -> { - if(createdCache) - inputCache.scheduleDeletion(); - }); + OOCInstructionUtils.tsmm(min.getStreamable(), out, _type, (AggregateBinaryOperator) _optr, plus, getContext()); } private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObject min) { @@ -119,52 +91,4 @@ private void processSingleOutputTileInstruction(ExecutionContext ec, MatrixObjec ((MatrixBlock) left.getValue()).binaryOperationsInPlace(plus, right.getValue())), value -> ((MatrixBlock) value.getValue()).getExactSerializedSize(), getContext()); } - - private long getJoinIndex(IndexedMatrixValue value) { - return _type.isLeft() ? value.getIndexes().getRowIndex() : value.getIndexes().getColumnIndex(); - } - - private long getOutputIndex(IndexedMatrixValue value) { - return _type.isLeft() ? value.getIndexes().getColumnIndex() : value.getIndexes().getRowIndex(); - } - - private List createPartialOutputTiles(IndexedMatrixValue left, IndexedMatrixValue right) { - long leftIndex = getOutputIndex(left); - long rightIndex = getOutputIndex(right); - if(leftIndex > rightIndex) - return List.of(); - - MatrixBlock leftBlock = (MatrixBlock) left.getValue(); - MatrixBlock rightBlock = (MatrixBlock) right.getValue(); - if(leftIndex == rightIndex) { - MatrixBlock diagonal = leftBlock.transposeSelfMatrixMultOperations(new MatrixBlock(), _type); - return List.of(new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), diagonal)); - } - - MatrixBlock partial = multiplyOffDiagonal(leftBlock, rightBlock); - MatrixBlock mirror = LibMatrixReorg.transpose(partial); - return List.of(new IndexedMatrixValue(new MatrixIndexes(leftIndex, rightIndex), partial), - new IndexedMatrixValue(new MatrixIndexes(rightIndex, leftIndex), mirror)); - } - - private MatrixBlock multiplyOffDiagonal(MatrixBlock leftBlock, MatrixBlock rightBlock) { - if(_type.isLeft()) { - MatrixBlock leftTranspose = LibMatrixReorg.transpose(leftBlock); - return leftTranspose.aggregateBinaryOperations(leftTranspose, rightBlock, new MatrixBlock(), - (AggregateBinaryOperator) _optr); - } - - MatrixBlock rightTranspose = LibMatrixReorg.transpose(rightBlock); - return leftBlock.aggregateBinaryOperations(leftBlock, rightTranspose, new MatrixBlock(), - (AggregateBinaryOperator) _optr); - } - - private static void propagateFailuresToOutput(OOCStream out, List> futures) { - for(CompletableFuture future : futures) { - future.exceptionally(error -> { - out.propagateFailure(DMLRuntimeException.of(error)); - return null; - }); - } - } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java index 03eea8dcf31..838dd18206c 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java @@ -62,7 +62,9 @@ private static void injectMaterializations(OOCPrimitive primitive, Set> _store; private final AtomicBoolean _finished; private final boolean _reusable; + private final List>> _liveConsumers; + private MaterializedStore _materializedStore; private int _expectedReaders; private int _consumers; @@ -59,21 +64,32 @@ private MaterializeOOCPrimitive(OOCStreamable source, OOCSto _store = new OOCFuture<>(); _finished = new AtomicBoolean(); _reusable = reusable; + _liveConsumers = new ArrayList<>(); } public static MaterializeOOCPrimitive reusable(OOCStreamable source) { - return new MaterializeOOCPrimitive(source, OOCStoreLayout.ROW_MAJOR, null, true); + return reusable(source, OOCStoreLayout.ROW_MAJOR); } - public synchronized void registerRequest(int expectedReaders) { - if(_reusable) - throw new IllegalStateException("Reusable materialization registers readers dynamically."); + public static MaterializeOOCPrimitive reusable(OOCStreamable source, OOCStoreLayout layout) { + return new MaterializeOOCPrimitive(source, layout, null, true); + } + + public synchronized boolean registerRequest(int expectedReaders, + Consumer> liveConsumer) { if(expectedReaders <= 0) throw new IllegalArgumentException("Materialization request requires at least one reader."); - if(hasStartedExecution()) - throw new IllegalStateException("Cannot register a consumer after materialization started."); - _expectedReaders = Math.addExact(_expectedReaders, expectedReaders); - _consumers = Math.addExact(_consumers, 1); + boolean live = !hasStartedExecution(); + if(_materializedStore == null) { + if(!_reusable) + _expectedReaders = Math.addExact(_expectedReaders, expectedReaders); + _consumers = Math.addExact(_consumers, 1); + } + else + _materializedStore.registerConsumer(expectedReaders); + if(live && liveConsumer != null) + _liveConsumers.add(liveConsumer); + return live; } public OOCFuture> store() { @@ -99,19 +115,24 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { protected void startExecution() { try { OOCStream source = getInputReadStream(0); - MaterializedStore store = _reusable ? new MaterializedStore<>( - OOCCacheManager.getGlobalCache(), - CachingStream._streamSeq.getNextID()) : new MaterializedStore<>(OOCCacheManager.getGlobalCache(), - CachingStream._streamSeq.getNextID(), _expectedReaders, _consumers); DataCharacteristics characteristics = _source.getDataCharacteristics(); + boolean logicalLayout = characteristics != null && characteristics.dimsKnown() && + characteristics.getBlocksize() > 0; + ToIntFunction linearize = logicalLayout ? indexes -> _layout.linearize(indexes, + characteristics) : null; + MaterializedStore store; + synchronized(this) { + int consumers = _reusable ? 1 + _consumers : _consumers; + store = new MaterializedStore<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID(), + _reusable ? -1 : _expectedReaders, consumers, logicalLayout ? _layout : null, + logicalLayout ? characteristics : null); + _materializedStore = store; + } AtomicInteger nextIndex = new AtomicInteger(); - ToIntFunction linearize; - if(_reusable && - (characteristics == null || !characteristics.dimsKnown() || characteristics.getBlocksize() <= 0)) - linearize = ignored -> nextIndex.getAndIncrement(); - else - linearize = indexes -> _layout.linearize(indexes, characteristics); - OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, linearize, _allowance); + ToIntFunction publicationIndex = linearize != null ? linearize : ignored -> nextIndex + .getAndIncrement(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, publicationIndex, _allowance, + _liveConsumers); materializer.completion().whenComplete((ignored, error) -> { if(error != null) fail(error); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java index 45b7a3ff031..1e87184d243 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/OOCPrimitive.java @@ -25,6 +25,7 @@ import java.util.Objects; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.OOCStream; @@ -229,6 +230,15 @@ private InputSlot(OOCStreamable source) { } } - public record OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders) { + public record OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders, + Consumer> liveConsumer, Consumer liveRegistration) { + public OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders) { + this(inputIndex, layout, expectedReaders, null, null); + } + + public OOCMaterializedInputRequest(int inputIndex, OOCStoreLayout layout, int expectedReaders, + Consumer> liveConsumer) { + this(inputIndex, layout, expectedReaders, liveConsumer, null); + } } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TSMMOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TSMMOOCPrimitive.java new file mode 100644 index 00000000000..6ed43140f2f --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TSMMOOCPrimitive.java @@ -0,0 +1,571 @@ +/* + * 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.concurrent.atomic.AtomicIntegerArray; + +import org.apache.sysds.lops.MMTSJ.MMTSJType; +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.LibMatrixReorg; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.matrix.operators.AggregateBinaryOperator; +import org.apache.sysds.runtime.matrix.operators.BinaryOperator; +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.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.planning.OOCStoreLayout; +import org.apache.sysds.runtime.ooc.store.CountingLiveness; +import org.apache.sysds.runtime.ooc.store.IndexedMaterializedStoreReader; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; +import org.apache.sysds.runtime.ooc.store.StateTable; +import org.apache.sysds.runtime.ooc.store.StoreLease; +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 TSMMOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _output; + private final MMTSJType _type; + private final AggregateBinaryOperator _multiply; + private final BinaryOperator _plus; + private final AtomicBoolean _inputComplete = new AtomicBoolean(); + private final AtomicBoolean _liveInputEnded = new AtomicBoolean(); + private final AtomicBoolean _liveSchedulingReady = new AtomicBoolean(); + private final AtomicInteger _active = new AtomicInteger(1); + private boolean _liveInput; + private MaterializedStore _inputStore; + private volatile IndexedMaterializedStoreReader _inputReader; + private StateTable _accumulators; + private AtomicIntegerArray _tilesSeen; + private AtomicIntegerArray _tilesArrived; + private AtomicIntegerArray _groupsScheduled; + private OOCStream _ready; + private OOCStream _outputStream; + private int _groups; + private int _width; + private long _taskBytes; + + public TSMMOOCPrimitive(OOCStreamable input, OOCStreamable output, + MMTSJType type, AggregateBinaryOperator multiply, BinaryOperator plus, StreamContext context) { + super(context, input); + _output = output; + _type = type; + _multiply = multiply; + _plus = plus; + } + + @Override + public List requiredMaterializedInputs() { + OOCStoreLayout layout = _type.isLeft() ? OOCStoreLayout.ROW_MAJOR : OOCStoreLayout.COL_MAJOR; + return List.of(new OOCMaterializedInputRequest(0, layout, 1, this::accept, live -> _liveInput = live)); + } + + @Override + protected void inferPatternsInternal() { + _pattern = _type.isLeft() ? OOCAccessPattern.ROW_MAJOR : OOCAccessPattern.COL_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(_pattern); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = _type.isLeft() ? OOCAccessPattern.ROW_MAJOR : OOCAccessPattern.COL_MAJOR; + for(OOCPrimitive child : getChildren()) + child.requestPattern(_pattern); + } + + private static long taskBytes(DataCharacteristics inputDc, DataCharacteristics outputDc) { + long inputBytes = OOCUtils.estimateFullTileBytes(inputDc); + long outputBytes = OOCUtils.estimateFullTileBytes(outputDc); + long multiplyBytes = 2 * OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(inputBytes) + 2 * outputBytes; + long mergeBytes = OOCCacheManager.getGlobalCache().maxPhysicalPinBytes(outputBytes) + 2 * outputBytes; + return Math.max(multiplyBytes, mergeBytes); + } + + @Override + protected void startExecution() { + DataCharacteristics inputDc = getInput(0).getDataCharacteristics(); + if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) + throw new DMLRuntimeException("TSMM OOC requires known input dimensions and block size."); + _groups = Math.toIntExact(_type.isLeft() ? inputDc.getNumRowBlocks() : inputDc.getNumColBlocks()); + _width = Math.toIntExact(_type.isLeft() ? inputDc.getNumColBlocks() : inputDc.getNumRowBlocks()); + if(_groups <= 0 || _width <= 0) + throw new DMLRuntimeException("TSMM OOC requires non-empty input block geometry."); + _tilesSeen = new AtomicIntegerArray(_groups); + _tilesArrived = new AtomicIntegerArray(_groups * _width); + _groupsScheduled = new AtomicIntegerArray(_groups); + _accumulators = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); + long accumulatorPriorityOffset = (long) _width * _width; + _accumulators.addEvictionPolicy(slot -> slot - accumulatorPriorityOffset); + _outputStream = _output.getWriteStream(); + _ready = new SubscribableTaskQueue<>(); + + _taskBytes = taskBytes(inputDc, _output.getDataCharacteristics()); + + getContext().addOutStream(_outputStream, _ready); + OOCInstructionUtils.submitCloseableOOCTasks(_ready, work -> { + if(work instanceof MultiplyWork multiply) + multiply(multiply); + else + merge((MergeWork) work); + }, getContext()).whenComplete((ignored, error) -> { + try { + if(error != null) + fail(error); + _outputStream.closeInput(); + } + catch(Throwable failure) { + fail(failure); + } + finally { + cleanup(); + } + }); + + getMaterializedInput(0).whenComplete((store, error) -> { + if(error != null) { + fail(error); + finishInput(); + return; + } + _inputStore = store; + CountingLiveness liveness = new CountingLiveness(_groups * _width, _width + 1); + if(_liveInput) { + _inputReader = store.openLiveIndexedReader(liveness); + for(int group = 0; group < _groups; group++) + tryScheduleGroup(group); + _liveSchedulingReady.set(true); + finishLiveInput(); + } + else { + store.completion().whenComplete((ignored, completionError) -> { + if(completionError != null) { + fail(completionError); + finishInput(); + return; + } + _inputReader = store.openIndexedReader(liveness); + OOCInstructionUtils.submitOOCTask(this::drain, new StreamContext().addOutStream(_outputStream)); + }); + } + }); + } + + private void accept(OOCStream.QueueCallback callback) { + if(callback.isEos() || callback.isFailure()) { + try(callback) { + if(callback.isFailure()) + callback.get(); + } + catch(Throwable failure) { + fail(failure); + } + _liveInputEnded.set(true); + finishLiveInput(); + return; + } + + try(callback) { + IndexedMatrixValue tile = callback.get(); + long rowIndex = tile.getIndexes().getRowIndex(); + long colIndex = tile.getIndexes().getColumnIndex(); + int group = Math.toIntExact((_type.isLeft() ? rowIndex : colIndex) - 1); + int position = Math.toIntExact((_type.isLeft() ? colIndex : rowIndex) - 1); + if(group < 0 || group >= _groups || position < 0 || position >= _width) + throw new DMLRuntimeException("TSMM live tile " + tile.getIndexes() + " is outside the input geometry " + + _groups + "x" + _width + "."); + if(!_tilesArrived.compareAndSet(group * _width + position, 0, 1)) + return; + _tilesSeen.incrementAndGet(group); + tryScheduleGroup(group); + } + catch(Throwable failure) { + fail(failure); + finishInput(); + } + } + + private void finishLiveInput() { + if(_liveInputEnded.get() && _liveSchedulingReady.get()) + finishInput(); + } + + private void drain() { + try { + for(int group = 0; group < _groups && !hasFailed(); group++) + scheduleGroup(group); + } + catch(Throwable failure) { + fail(failure); + } + finally { + finishInput(); + } + } + + private void tryScheduleGroup(int group) { + if(_inputReader != null && _tilesSeen.get(group) == _width && _groupsScheduled.compareAndSet(group, 0, 1)) + scheduleGroup(group); + } + + private void scheduleGroup(int group) { + for(int left = 0; left < _width; left++) + for(int right = left; right < _width; right++) + schedulePair(group, left, right); + } + + private void schedulePair(int group, int left, int right) { + _active.incrementAndGet(); + _allowance.reserveAsync(_taskBytes).whenComplete((ignored, admissionError) -> { + if(admissionError != null) { + fail(admissionError); + completeOne(); + return; + } + ReservationBudget budget = new ReservationBudget(_allowance, _taskBytes).enableReuse(); + long groupIndex = group + 1L; + long leftIndex = left + 1L; + long rightIndex = right + 1L; + OOCFuture + .allOf(List.of( + _inputReader.request(_type.isLeft() ? groupIndex : leftIndex, + _type.isLeft() ? leftIndex : groupIndex, budget), + _inputReader.request(_type.isLeft() ? groupIndex : rightIndex, + _type.isLeft() ? rightIndex : groupIndex, budget)), + StoreLease::close) + .whenComplete((inputs, inputError) -> { + if(inputError != null) { + budget.close(); + fail(inputError); + completeOne(); + return; + } + try { + if(inputs.get(0) == null || inputs.get(1) == null) + throw new DMLRuntimeException("Missing buffered TSMM tiles for group " + (group + 1)); + _ready.enqueue(new MultiplyWork(group, left, right, inputs.get(0), inputs.get(1), budget)); + } + catch(Throwable failure) { + for(StoreLease lease : inputs) + if(lease != null) + lease.close(); + budget.close(); + fail(failure); + completeOne(); + } + }); + }); + } + + private void multiply(MultiplyWork work) { + ReservationBudget budget = work.takeBudget(); + ManagedPayload partial = null; + try { + MatrixBlock left = (MatrixBlock) work._left.value().getValue(); + MatrixBlock right = (MatrixBlock) work._right.value().getValue(); + MatrixBlock block; + if(work._leftPosition == work._rightPosition) + block = left.transposeSelfMatrixMultOperations(new MatrixBlock(), _type); + else if(_type.isLeft()) { + MatrixBlock transposed = LibMatrixReorg.transpose(left); + block = transposed.aggregateBinaryOperations(transposed, right, new MatrixBlock(), _multiply); + } + else { + MatrixBlock transposed = LibMatrixReorg.transpose(right); + block = left.aggregateBinaryOperations(left, transposed, new MatrixBlock(), _multiply); + } + int outputSlot = work._leftPosition * _width + work._rightPosition; + partial = payload(outputSlot, 1, block, budget); + OOCFuture> released = work.releaseInputsAsync(); + ManagedPayload result = partial; + partial = null; + released.whenComplete((ignored, error) -> { + if(error != null) { + result.release(); + budget.close(); + fail(error); + completeOne(); + } + else + reduce(work._group, outputSlot, result, budget); + }); + } + catch(Throwable failure) { + if(partial != null) + partial.release(); + budget.close(); + fail(failure); + completeOne(); + } + } + + private void reduce(int group, int slot, ManagedPayload incoming, ReservationBudget budget) { + if(count(incoming.value()) == _groups) { + finalizeOutput(slot, incoming, budget); + return; + } + OOCFuture> match; + try { + match = _accumulators.putOrTake(slot, 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 { + try { + _ready.enqueue(new MergeWork(group, slot, incoming, existing, budget)); + } + catch(Throwable failure) { + incoming.release(); + existing.close(); + budget.close(); + fail(failure); + completeOne(); + } + } + }); + } + + private void merge(MergeWork work) { + ReservationBudget budget = work.takeBudget(); + ManagedPayload merged = null; + try { + IndexedMatrixValue existing = work._existing.value(); + IndexedMatrixValue incoming = work._incoming.value(); + MatrixBlock block = ((MatrixBlock) existing.getValue()).binaryOperations(_plus, incoming.getValue(), + new MatrixBlock()); + merged = payload(work._slot, count(existing) + count(incoming), block, budget); + work.releaseIncoming(); + OOCFuture released = work.closeExistingAsync(); + ManagedPayload result = merged; + merged = null; + released.whenComplete((ignored, error) -> { + if(error != null) { + result.release(); + budget.close(); + fail(error); + completeOne(); + } + else + reduce(work._group, work._slot, result, budget); + }); + } + catch(Throwable failure) { + if(merged != null) + merged.release(); + budget.close(); + fail(failure); + completeOne(); + } + } + + private void finalizeOutput(int slot, ManagedPayload payload, ReservationBudget budget) { + OOCStream.QueueCallback upper = null; + OOCStream.QueueCallback lower = null; + try { + int row = slot / _width; + int col = slot % _width; + MatrixBlock block = (MatrixBlock) payload.value().getValue(); + payload.release(); + long upperBytes = block.getExactSerializedSize(); + budget.reserveBlocking(upperBytes); + upper = new InMemoryQueueCallback<>(new IndexedMatrixValue(new MatrixIndexes(row + 1L, col + 1L), block), + null, budget, upperBytes); + if(row != col) { + MatrixBlock mirror = LibMatrixReorg.transpose(block); + long lowerBytes = mirror.getExactSerializedSize(); + budget.reserveBlocking(lowerBytes); + lower = new InMemoryQueueCallback<>( + new IndexedMatrixValue(new MatrixIndexes(col + 1L, row + 1L), mirror), null, budget, lowerBytes); + } + budget.close(); + _outputStream.enqueue(upper); + upper = null; + if(lower != null) { + _outputStream.enqueue(lower); + lower = null; + } + } + catch(Throwable failure) { + payload.release(); + budget.close(); + fail(failure); + } + finally { + if(upper != null) + upper.close(); + if(lower != null) + lower.close(); + } + completeOne(); + } + + private static ManagedPayload payload(int slot, int count, MatrixBlock block, + ReservationBudget budget) { + long bytes = block.getExactSerializedSize(); + budget.reserveBlocking(bytes); + return new ManagedPayload<>(new IndexedMatrixValue(new MatrixIndexes(slot + 1L, count), block), bytes, budget); + } + + private static int count(IndexedMatrixValue value) { + return Math.toIntExact(value.getIndexes().getColumnIndex()); + } + + private void finishInput() { + if(_inputComplete.compareAndSet(false, true)) + completeOne(); + } + + private void completeOne() { + if(_active.decrementAndGet() != 0) + return; + try { + _ready.closeInput(); + } + catch(IllegalStateException ignored) { + } + } + + private void cleanup() { + _accumulators.close(); + if(_inputReader != null) + _inputReader.close(); + if(_inputStore != null) + _inputStore.close(); + onComplete(); + } + + private static final class MultiplyWork implements AutoCloseable { + private final int _group; + private final int _leftPosition; + private final int _rightPosition; + private StoreLease _left; + private StoreLease _right; + private ReservationBudget _budget; + + private MultiplyWork(int group, int leftPosition, int rightPosition, StoreLease left, + StoreLease right, ReservationBudget budget) { + _group = group; + _leftPosition = leftPosition; + _rightPosition = rightPosition; + _left = left; + _right = right; + _budget = budget; + } + + private ReservationBudget takeBudget() { + ReservationBudget budget = _budget; + _budget = null; + return budget; + } + + private OOCFuture> releaseInputsAsync() { + OOCFuture left = _left.closeAsync(); + OOCFuture right = _right.closeAsync(); + _left = null; + _right = null; + return OOCFuture.allOf(List.of(left, right), ignored -> { + }); + } + + @Override + public void close() { + if(_left != null) + _left.close(); + if(_right != null) + _right.close(); + if(_budget != null) + _budget.close(); + } + } + + private static final class MergeWork implements AutoCloseable { + private final int _group; + private final int _slot; + private ManagedPayload _incoming; + private StoreLease _existing; + private ReservationBudget _budget; + + private MergeWork(int group, int slot, ManagedPayload incoming, + StoreLease existing, ReservationBudget budget) { + _group = group; + _slot = slot; + _incoming = incoming; + _existing = existing; + _budget = budget; + } + + private ReservationBudget takeBudget() { + ReservationBudget budget = _budget; + _budget = null; + return budget; + } + + private void releaseIncoming() { + _incoming.release(); + _incoming = null; + } + + private OOCFuture closeExistingAsync() { + OOCFuture released = _existing.closeAsync(); + _existing = null; + return released; + } + + @Override + public void close() { + if(_incoming != null) + _incoming.release(); + if(_existing != null) + _existing.close(); + if(_budget != null) + _budget.close(); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java b/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java index 1c3a4c9005f..0074733c9b6 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java @@ -19,11 +19,14 @@ package org.apache.sysds.runtime.ooc.store; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.meta.DataCharacteristics; import org.apache.sysds.runtime.ooc.cache.BlockEntry; import org.apache.sysds.runtime.ooc.cache.OOCCache; import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; import org.apache.sysds.runtime.ooc.util.OOCUtils; import java.util.function.IntConsumer; @@ -34,16 +37,21 @@ public final class IndexedMaterializedStoreReader imp private final long _streamId; private final IntSupplier _completedSize; private final MaterializedStore.Liveness _liveness; + private final OOCStoreLayout _layout; + private final DataCharacteristics _characteristics; private final Runnable _afterClose; private final IntConsumer _afterRelease; private volatile boolean _closed; IndexedMaterializedStoreReader(OOCCache cache, long streamId, IntSupplier completedSize, - MaterializedStore.Liveness liveness, Runnable afterClose, IntConsumer afterRelease) { + MaterializedStore.Liveness liveness, OOCStoreLayout layout, DataCharacteristics characteristics, + Runnable afterClose, IntConsumer afterRelease) { _cache = cache; _streamId = streamId; _completedSize = completedSize; _liveness = liveness; + _layout = layout; + _characteristics = characteristics; _afterClose = afterClose; _afterRelease = afterRelease; } @@ -66,6 +74,16 @@ public void close() { _afterClose.run(); } + public OOCFuture> request(MatrixIndexes indexes, MemoryAllowance requestAllowance) { + return request(indexes.getRowIndex(), indexes.getColumnIndex(), requestAllowance); + } + + public OOCFuture> request(long row, long col, MemoryAllowance requestAllowance) { + if(_layout == null) + throw new IllegalStateException("Materialized reader has no logical matrix-index layout."); + return request(_layout.linearize(row, col, _characteristics), requestAllowance); + } + public OOCFuture> request(int index, MemoryAllowance requestAllowance) { checkReady(index); reserve(index); @@ -86,6 +104,16 @@ else if(entry == null) { return result; } + public StoreLease requestIfLive(MatrixIndexes indexes, MemoryAllowance requestAllowance) { + return requestIfLive(indexes.getRowIndex(), indexes.getColumnIndex(), requestAllowance); + } + + public StoreLease requestIfLive(long row, long col, MemoryAllowance requestAllowance) { + if(_layout == null) + throw new IllegalStateException("Materialized reader has no logical matrix-index layout."); + return requestIfLive(_layout.linearize(row, col, _characteristics), requestAllowance); + } + public StoreLease requestIfLive(int index, MemoryAllowance requestAllowance) { checkReady(index); reserve(index); diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java index aab137d8298..fdb3837dff5 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java @@ -25,7 +25,9 @@ import org.apache.sysds.runtime.ooc.cache.OOCFuture; import org.apache.sysds.runtime.ooc.cache.io.SpillableObject; import org.apache.sysds.runtime.ooc.memory.ManagedPayload; +import org.apache.sysds.runtime.meta.DataCharacteristics; import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.runtime.ooc.planning.OOCStoreLayout; import org.apache.sysds.runtime.ooc.util.OOCUtils; import java.util.ArrayList; @@ -44,6 +46,8 @@ public final class MaterializedStore { private final OOCFuture _completion; private final OOCFuture _readersSealedFuture; private final boolean _autoSealReaders; + private final OOCStoreLayout _layout; + private final DataCharacteristics _characteristics; private volatile List _readers; private volatile int _completedSize; @@ -54,10 +58,15 @@ public final class MaterializedStore { private int _consumers; public MaterializedStore(OOCCache cache, long streamId) { - this(cache, streamId, -1, 1); + this(cache, streamId, -1, 1, null, null); } public MaterializedStore(OOCCache cache, long streamId, int expectedReaders, int consumers) { + this(cache, streamId, expectedReaders, consumers, null, null); + } + + public MaterializedStore(OOCCache cache, long streamId, int expectedReaders, int consumers, OOCStoreLayout layout, + DataCharacteristics characteristics) { if(expectedReaders == 0 || expectedReaders < -1) throw new IllegalArgumentException("Expected reader count must be positive or disabled."); if(consumers <= 0) @@ -71,6 +80,8 @@ public MaterializedStore(OOCCache cache, long streamId, int expectedReaders, int _completion = new OOCFuture<>(); _readersSealedFuture = new OOCFuture<>(); _autoSealReaders = expectedReaders > 0; + _layout = layout; + _characteristics = characteristics; _pendingReaders = expectedReaders; _consumers = consumers; _readers = Collections.emptyList(); @@ -113,6 +124,15 @@ public synchronized void complete() { + " published items for logical range [0, " + _completedSize + ")"); _complete = true; _completion.complete(null); + if(_autoSealReaders && _pendingReaders == 0) + sealReaders(); + } + + public synchronized void registerConsumer(int expectedReaders) { + if(_readersSealed || _closed) + throw new IllegalStateException("Store no longer accepts consumers"); + _pendingReaders += expectedReaders; + _consumers++; } void failMaterialization(Throwable error) { @@ -151,7 +171,19 @@ public synchronized IndexedMaterializedStoreReader openIndexedReader(Liveness if(_readersSealed) throw new IllegalStateException("Store no longer accepts new readers"); IndexedMaterializedStoreReader reader = new IndexedMaterializedStoreReader<>(_cache, _streamId, - () -> _completedSize, liveness, this::forgetAfterReaderClose, this::tryForget); + () -> _completedSize, liveness, _layout, _characteristics, this::forgetAfterReaderClose, this::tryForget); + _registeredReaders.add(reader); + readerRegistered(); + return reader; + } + + public synchronized IndexedMaterializedStoreReader openLiveIndexedReader(Liveness liveness) { + if(_closed) + throw new IllegalStateException("Store is closed"); + if(_readersSealed) + throw new IllegalStateException("Store no longer accepts new readers"); + IndexedMaterializedStoreReader reader = new IndexedMaterializedStoreReader<>(_cache, _streamId, this::size, + liveness, _layout, _characteristics, this::forgetAfterReaderClose, this::tryForget); _registeredReaders.add(reader); readerRegistered(); return reader; @@ -194,7 +226,7 @@ private void readerRegistered() { return; if(_pendingReaders <= 0) throw new IllegalStateException("More materialized readers opened than declared."); - if(--_pendingReaders == 0) + if(--_pendingReaders == 0 && _complete) sealReaders(); } 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 a764fe03589..e52fea34a6a 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 @@ -34,6 +34,7 @@ import java.util.function.ToLongFunction; import org.apache.sysds.api.DMLScript; +import org.apache.sysds.lops.MMTSJ.MMTSJType; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.instructions.ooc.OOCStream; import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; @@ -55,6 +56,7 @@ import org.apache.sysds.runtime.ooc.primitives.NaryJoinOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.PlannableDataGenOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.ReduceOOCPrimitive; +import org.apache.sysds.runtime.ooc.primitives.TSMMOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.TransposeOOCPrimitive; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; import org.apache.sysds.runtime.ooc.store.MaterializedStore; @@ -133,6 +135,11 @@ public static void naryEquiJoin(List> inputs, new NaryJoinOOCPrimitive(inputs, output, key, size, operation, inputBytes, joinBytes, context)); } + public static void tsmm(OOCStreamable input, OOCStream output, + MMTSJType type, AggregateBinaryOperator multiply, BinaryOperator plus, StreamContext context) { + output.assignPrimitive(new TSMMOOCPrimitive(input, output, type, multiply, plus, context)); + } + public static void matrixMultiply(OOCStreamable left, OOCStreamable right, OOCStream output, AggregateBinaryOperator multiply, BinaryOperator plus, StreamContext context) { diff --git a/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java index 14d26db3ae2..6071af3f015 100644 --- a/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java +++ b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java @@ -89,10 +89,10 @@ public void testMaterializationReadersAndForgetting() throws Exception { OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 1.0), null, _producer, TILE_BYTES)); materializer.accept(new OOCStream.SimpleQueueCallback<>(tile(1, 2.0), null)); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(2, 3.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(2, 3.0), null, _producer, TILE_BYTES)); materializer.accept(OOCStream.eos(null)); materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); @@ -125,13 +125,54 @@ public void testMaterializationReadersAndForgetting() throws Exception { Assert.assertEquals(0, _readerAllowance.getUsedMemory()); } + @Test + public void testLiveIndexedReader() throws Exception { + IndexedMaterializedStoreReader reader = _store + .openLiveIndexedReader(new CountingLiveness(1, 1)); + _store.sealReaders(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, indexes -> 0, _materializerAllowance); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 7), null, _producer, TILE_BYTES)); + + Assert.assertFalse(_store.completion().isDone()); + try(StoreLease lease = reader.request(0, _readerAllowance).get(WAIT_SECONDS, + TimeUnit.SECONDS)) { + Assert.assertEquals(7, lease.value().getValue().get(0, 0), 0); + } + OOCCacheTestUtils.await(() -> _cache.getOwnedCacheSize() == 0, WAIT_SECONDS); + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + } + + @Test + public void testLateMaterializedReader() throws Exception { + MaterializedStore store = new MaterializedStore<>(_cache, + CachingStream._streamSeq.getNextID(), 1, 1); + IndexedMaterializedStoreReader live = store + .openLiveIndexedReader(new CountingLiveness(1, 1)); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(store, indexes -> 0, _materializerAllowance); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 7), null, _producer, TILE_BYTES)); + store.registerConsumer(1); + materializer.accept(OOCStream.eos(null)); + IndexedMaterializedStoreReader late = store.openIndexedReader(new CountingLiveness(1, 1)); + + for(IndexedMaterializedStoreReader reader : List.of(live, late)) + try(StoreLease lease = reader.request(0, _readerAllowance).get(WAIT_SECONDS, + TimeUnit.SECONDS)) { + Assert.assertEquals(7, lease.value().getValue().get(0, 0), 0); + } + store.close(); + store.close(); + } + @Test public void testOrderedReaderRetries() throws Exception { OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); for(int i = 0; i < 2; i++) { _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(i, i + 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(i, i + 1.0), null, _producer, TILE_BYTES)); } materializer.accept(OOCStream.eos(null)); materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); @@ -161,10 +202,10 @@ public void testSoftOrderingReturnsReadyRequestFirst() throws Exception { OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); _producer.reserveBlocking(largeBytes); - materializer.accept(new InMemoryQueueCallback(new IndexedMatrixValue(new MatrixIndexes(1, 1), largeBlock), null, - _producer, largeBytes)); + materializer.accept(new InMemoryQueueCallback<>(new IndexedMatrixValue(new MatrixIndexes(1, 1), largeBlock), + null, _producer, largeBytes)); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(1, 2.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(1, 2.0), null, _producer, TILE_BYTES)); materializer.accept(OOCStream.eos(null)); materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); @@ -195,7 +236,7 @@ public void testDirectRequests() throws Exception { OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 1.0), null, _producer, TILE_BYTES)); materializer.accept(OOCStream.eos(null)); materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); @@ -229,7 +270,7 @@ public void testCompletionMissingPublications() throws Exception { })); for(int index : new int[] {0, 2}) { _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(index, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(index, 1.0), null, _producer, TILE_BYTES)); } materializer.accept(OOCStream.eos(null)); @@ -257,7 +298,7 @@ public void testLiveCallbackKeepsPublicationPinned() throws Exception { })); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 1.0), null, _producer, TILE_BYTES)); Assert.assertEquals(TILE_BYTES, _producer.getUsedMemory()); Assert.assertNotNull(retained.get()); retained.get().close(); @@ -334,7 +375,7 @@ public void testFailurePropagation() throws Exception { materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); _producer.reserveBlocking(TILE_BYTES); - materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new InMemoryQueueCallback<>(tile(0, 1.0), null, _producer, TILE_BYTES)); try { materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); Assert.fail("Publishing into a closed store must fail"); 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 b408d3f1e9e..f7972268d45 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 @@ -23,10 +23,16 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.apache.sysds.common.Types.FileFormat; import org.apache.sysds.common.Types.ValueType; +import org.apache.sysds.lops.MMTSJ.MMTSJType; import org.apache.sysds.runtime.controlprogram.caching.MatrixObject; +import org.apache.sysds.runtime.functionobjects.Multiply; +import org.apache.sysds.runtime.functionobjects.Plus; import org.apache.sysds.runtime.instructions.ooc.CachingStream; import org.apache.sysds.runtime.instructions.ooc.OOCStream; import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; @@ -34,6 +40,9 @@ 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.matrix.operators.AggregateBinaryOperator; +import org.apache.sysds.runtime.matrix.operators.AggregateOperator; +import org.apache.sysds.runtime.matrix.operators.BinaryOperator; import org.apache.sysds.runtime.meta.MatrixCharacteristics; import org.apache.sysds.runtime.meta.MetaDataFormat; import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; @@ -249,6 +258,80 @@ private static Map runGroupedReduce(GroupedReduceOOCPrimitive.Gr return values; } + @Test + public void testTsmmOutOfOrderGroups() { + SubscribableTaskQueue input = new SubscribableTaskQueue<>(); + SubscribableTaskQueue output = new SubscribableTaskQueue<>(); + input.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 2, 1), FileFormat.BINARY))); + output.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 2, 1), FileFormat.BINARY))); + for(double[] tile : List.of(new double[] {2, 1, 3}, new double[] {1, 2, 2}, new double[] {2, 2, 4}, + new double[] {1, 1, 1})) + input.enqueue(new IndexedMatrixValue(new MatrixIndexes((long) tile[0], (long) tile[1]), + new MatrixBlock(1, 1, tile[2]))); + input.closeInput(); + + AggregateOperator aggregate = new AggregateOperator(0, Plus.getPlusFnObject()); + OOCInstructionUtils.tsmm(input, output, MMTSJType.LEFT, + new AggregateBinaryOperator(Multiply.getMultiplyFnObject(), aggregate), + new BinaryOperator(Plus.getPlusFnObject()), 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)); + } + Assert.assertEquals(Map.of("1,1", 10d, "1,2", 14d, "2,1", 14d, "2,2", 20d), values); + } + + @Test + public void testTsmmConsumesLiveMaterialization() throws InterruptedException { + SubscribableTaskQueue input = new SubscribableTaskQueue<>(); + SubscribableTaskQueue output = new SubscribableTaskQueue<>(); + input.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(1, 2, 1), FileFormat.BINARY))); + output.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 2, 1), FileFormat.BINARY))); + + AggregateOperator aggregate = new AggregateOperator(0, Plus.getPlusFnObject()); + OOCInstructionUtils.tsmm(input, output, MMTSJType.LEFT, + new AggregateBinaryOperator(Multiply.getMultiplyFnObject(), aggregate), + new BinaryOperator(Plus.getPlusFnObject()), new StreamContext()); + + Map values = new ConcurrentHashMap<>(); + CountDownLatch blocks = new CountDownLatch(4); + CountDownLatch complete = new CountDownLatch(1); + output.setSubscriber(callback -> { + try(callback) { + if(callback.isEos() || callback.isFailure()) { + complete.countDown(); + return; + } + IndexedMatrixValue value = callback.get(); + values.put(value.getIndexes().getRowIndex() + "," + value.getIndexes().getColumnIndex(), + value.getValue().get(0, 0)); + blocks.countDown(); + } + }); + output.start(); + + try { + input.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 1d))); + input.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 2), new MatrixBlock(1, 1, 2d))); + Assert.assertTrue("TSMM waited for materialization completion", blocks.await(10, TimeUnit.SECONDS)); + Assert.assertEquals(Map.of("1,1", 1d, "1,2", 2d, "2,1", 2d, "2,2", 4d), values); + } + finally { + input.closeInput(); + } + Assert.assertTrue("TSMM output did not close", complete.await(10, TimeUnit.SECONDS)); + } + @Test public void testNaryJoinOutOfOrder() { SubscribableTaskQueue first = new SubscribableTaskQueue<>();