diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/BinaryOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/BinaryOOCInstruction.java index 8a5f7cf49a7..c252cd3d1ec 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/BinaryOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/BinaryOOCInstruction.java @@ -30,6 +30,7 @@ import org.apache.sysds.runtime.matrix.operators.BinaryOperator; import org.apache.sysds.runtime.matrix.operators.Operator; import org.apache.sysds.runtime.matrix.operators.ScalarOperator; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; public class BinaryOOCInstruction extends ComputationOOCInstruction { @@ -63,8 +64,6 @@ protected void processMatrixMatrixInstruction(ExecutionContext ec) { MatrixObject m1 = ec.getMatrixObject(input1); MatrixObject m2 = ec.getMatrixObject(input2); - OOCStream qIn1 = m1.getStreamHandle(); - OOCStream qIn2 = m2.getStreamHandle(); OOCStream qOut = new SubscribableTaskQueue<>(); ec.getMatrixObject(output).setStreamHandle(qOut); @@ -74,6 +73,8 @@ protected void processMatrixMatrixInstruction(ExecutionContext ec) { // If dimensions are unknown, we cannot safely detect broadcasting. // Fall back to strict key-based join and let downstream operators validate as needed. if(!known1 || !known2) { + OOCStream qIn1 = m1.getStreamHandle(); + OOCStream qIn2 = m2.getStreamHandle(); if(LOG.isWarnEnabled()) { LOG.warn("Falling back to key-wise OOC binary join for opcode '" + getOpcode() + "' due to unknown matrix dimensions: " + input1.getName() + "=" + m1.getNumRows() + "x" @@ -93,6 +94,8 @@ protected void processMatrixMatrixInstruction(ExecutionContext ec) { boolean isRowBroadcast = m1.getNumRows() > 1 && m2.getNumRows() == 1; if (isColBroadcast && !isRowBroadcast) { + OOCStream qIn1 = m1.getStreamHandle(); + OOCStream qIn2 = m2.getStreamHandle(); final long maxProcessesPerBroadcast = (m1.getNumColumns() + m1.getBlocksize() - 1) / m1.getBlocksize(); broadcastJoinOOC(qIn1, qIn2, qOut, (tmp1, b) -> { @@ -107,6 +110,8 @@ protected void processMatrixMatrixInstruction(ExecutionContext ec) { }, tmp -> tmp.getIndexes().getRowIndex()); } else if (isRowBroadcast && !isColBroadcast) { + OOCStream qIn1 = m1.getStreamHandle(); + OOCStream qIn2 = m2.getStreamHandle(); final long maxProcessesPerBroadcast = (m1.getNumRows() + m1.getBlocksize() - 1) / m1.getBlocksize(); broadcastJoinOOC(qIn1, qIn2, qOut, (tmp1, b) -> { @@ -126,12 +131,8 @@ else if (isRowBroadcast && !isColBroadcast) { + m1.getNumRows() + "x" + m1.getNumColumns() + " <=> " + m2.getNumRows() + "x" + m2.getNumColumns()); - joinOOC(qIn1, qIn2, qOut, (tmp1, tmp2) -> { - IndexedMatrixValue tmpOut = new IndexedMatrixValue(); - tmpOut.set(tmp1.getIndexes(), - tmp1.getValue().binaryOperations((BinaryOperator)_optr, tmp2.getValue(), tmpOut.getValue())); - return tmpOut; - }, IndexedMatrixValue::getIndexes); + OOCInstructionUtils.equiJoin(m1.getStreamable(), m2.getStreamable(), qOut, + (left, right) -> left.binaryOperations((BinaryOperator) _optr, right, new MatrixBlock()), getContext()); } } 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 07e1206a293..5cec3ae981d 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 @@ -30,6 +30,7 @@ import org.apache.sysds.runtime.ooc.cache.io.OOCMatrixIOHandler; import org.apache.sysds.runtime.ooc.cache.legacy.OOCCacheScheduler; import org.apache.sysds.runtime.ooc.cache.legacy.OOCLRUCacheScheduler; +import org.apache.sysds.runtime.ooc.cache.packed.OOCPackedCache; import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; import org.apache.sysds.utils.Statistics; @@ -52,22 +53,27 @@ public class OOCCacheManager { private static final AtomicReference _ioHandler; private static final AtomicReference _scheduler; + private static final AtomicReference _globalCache; static { _evictionLimit = (long)(Runtime.getRuntime().maxMemory() * OOC_BUFFER_PERCENTAGE); _hardLimit = (long)(Runtime.getRuntime().maxMemory() * OOC_BUFFER_PERCENTAGE_HARD); _ioHandler = new AtomicReference<>(); _scheduler = new AtomicReference<>(); + _globalCache = new AtomicReference<>(); } public static void reset() { TeeOOCInstruction.reset(); OOCIOHandler ioHandler = _ioHandler.getAndSet(null); OOCCacheScheduler cacheScheduler = _scheduler.getAndSet(null); + OOCPackedCache globalCache = _globalCache.getAndSet(null); if (ioHandler != null) ioHandler.shutdown(); if (cacheScheduler != null) cacheScheduler.shutdown(); + if(globalCache != null) + globalCache.shutdown(); if (DMLScript.OOC_STATISTICS) Statistics.resetOOCEvictionStats(); @@ -118,6 +124,18 @@ public static OOCCacheScheduler getCacheIfInitialized() { return _scheduler.get(); } + public static OOCPackedCache getGlobalCache() { + while(true) { + OOCPackedCache cache = _globalCache.get(); + if(cache != null) + return cache; + cache = new OOCPackedCache(new OOCMatrixIOHandler(), _hardLimit, _evictionLimit); + if(_globalCache.compareAndSet(null, cache)) + return cache; + cache.shutdown(); + } + } + public static OOCIOHandler getIOHandler() { OOCIOHandler io = _ioHandler.get(); if(io != null) diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java new file mode 100644 index 00000000000..99996eb7dac --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/JoinOOCPrimitive.java @@ -0,0 +1,220 @@ +/* + * 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.Objects; +import java.util.concurrent.ExecutionException; +import java.util.function.BiFunction; +import java.util.stream.Stream; + +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.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +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.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; +import org.apache.sysds.runtime.ooc.util.OOCUtils; +import org.apache.sysds.runtime.ooc.util.StateTableUtils; + +public class JoinOOCPrimitive extends OOCPrimitive { + private final OOCStream _left; + private final OOCStream _right; + private final OOCStreamable _output; + private final BiFunction _operation; + private StateTable _table; + + public JoinOOCPrimitive(OOCStreamable left, OOCStreamable right, + OOCStreamable output, BiFunction operation, + StreamContext context) { + this(left.getReadStream(), right.getReadStream(), output, operation, context); + } + + private JoinOOCPrimitive(OOCStream left, OOCStream right, + OOCStreamable output, BiFunction operation, + StreamContext context) { + super(context, Stream.of(left.getPrimitive(), right.getPrimitive()).filter(Objects::nonNull).toList()); + _left = left; + _right = right; + _output = output; + _operation = operation; + } + + @Override + protected void inferPatternsInternal() { + _pattern = OOCAccessPattern.ANY; + for(OOCPrimitive child : getChildren()) + _pattern = _pattern.fused(child.getAccessPattern()); + if(_pattern.isPlannable() && _pattern != OOCAccessPattern.ANY) + for(OOCPrimitive child : getChildren()) + child.requestPattern(_pattern); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = accessPattern; + for(OOCPrimitive child : getChildren()) + child.requestPattern(accessPattern); + } + + @Override + protected void startExecution() { + _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); + OOCStream output = _output.getWriteStream(); + OOCStream matches = new SubscribableTaskQueue<>(); + long inputBytes = Math.max(OOCUtils.estimateOutputTileBytes(_left.getDataCharacteristics()), + OOCUtils.estimateOutputTileBytes(_right.getDataCharacteristics())); + long outputBytes = OOCUtils.estimateOutputTileBytes(_output.getDataCharacteristics()); + long taskBytes = outputBytes + 2 * inputBytes; + + getContext().addOutStream(output); + OOCInstructionUtils.submitOOCTasks(matches, callback -> { + try(JoinWork work = callback.get()) { + IndexedMatrixValue left = work._left.get(); + IndexedMatrixValue right = work._right.get(); + OOCUtils.enqueueExact(output, new IndexedMatrixValue(left.getIndexes(), + _operation.apply((MatrixBlock) left.getValue(), (MatrixBlock) right.getValue())), work._budget); + } + }, callback -> true, (index, callback) -> callback.get().close(), getContext()).thenRun(() -> { + try { + _table.close(); + onComplete(); + } + finally { + output.closeInput(); + } + }); + + OOCInstructionUtils.submitOOCTask(() -> drive(matches, taskBytes), new StreamContext().addOutStream(output)); + } + + private void drive(OOCStream matches, long taskBytes) { + long cols = _right.getDataCharacteristics().getNumColBlocks(); + int unmatched = 0; + try { + while(true) { + OOCStream.QueueCallback left = _left.dequeueCB(); + OOCStream.QueueCallback right = _right.dequeueCB(); + boolean leftEos = left == null || left.isEos(); + boolean rightEos = right == null || right.isEos(); + if(leftEos || rightEos) { + if(left != null) + left.close(); + if(right != null) + right.close(); + if(leftEos != rightEos) + throw new DMLRuntimeException("Join inputs contain a different number of blocks"); + break; + } + unmatched += accept(left, true, cols, taskBytes, matches); + unmatched += accept(right, false, cols, taskBytes, matches); + } + if(unmatched != 0) + throw new DMLRuntimeException("Join inputs contain " + unmatched + " unmatched blocks"); + } + finally { + matches.closeInput(); + } + } + + private int accept(OOCStream.QueueCallback callback, boolean left, long cols, long taskBytes, + OOCStream matches) { + if(callback == null) + return 0; + OOCStream.QueueCallback owned = null; + ReservationBudget budget = null; + try { + owned = callback.keepOpen(); + callback.close(); + callback = null; + budget = OOCUtils.reserveBudget(_allowance, taskBytes); + IndexedMatrixValue value = owned.get(); + long row = value.getIndexes().getRowIndex() - 1; + long col = value.getIndexes().getColumnIndex() - 1; + int slot = Math.toIntExact(row * cols + col); + OOCFuture future = StateTableUtils.putOrTake(_table, slot, owned, budget); + owned = null; + StateTableUtils.Match match = await(future); + if(match == null) + return 1; + JoinWork work = left ? new JoinWork(match.left(), match.right(), budget) : new JoinWork(match.right(), + match.left(), budget); + budget = null; + try { + matches.enqueue(work); + work = null; + } + finally { + if(work != null) + work.close(); + } + return -1; + } + finally { + if(callback != null) + callback.close(); + if(owned != null) + owned.close(); + if(budget != null) + budget.close(); + } + } + + private static StateTableUtils.Match await(OOCFuture future) { + try { + return future.get(); + } + catch(InterruptedException error) { + Thread.currentThread().interrupt(); + throw new DMLRuntimeException(error); + } + catch(ExecutionException error) { + throw DMLRuntimeException.of(error.getCause()); + } + } + + private static final class JoinWork implements AutoCloseable { + private final OOCStream.QueueCallback _left; + private final OOCStream.QueueCallback _right; + private final ReservationBudget _budget; + + private JoinWork(OOCStream.QueueCallback left, + OOCStream.QueueCallback right, ReservationBudget budget) { + _left = left; + _right = right; + _budget = budget; + } + + @Override + public void close() { + try(_left; _right; _budget) { + // Release + } + } + } +} 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 781cdd04ea2..773ef9da834 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 @@ -40,6 +40,7 @@ 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.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; @@ -81,6 +82,12 @@ public static void transpose(OOCStreamable input, OOCStream< transposedMap(input, output, MatrixBlock::transpose, context); } + public static void equiJoin(OOCStreamable left, OOCStreamable right, + OOCStream output, BiFunction operation, + StreamContext context) { + output.assignPrimitive(new JoinOOCPrimitive(left, right, output, operation, 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 7b8a8ae46b1..5a98a05b78d 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 @@ -170,4 +170,9 @@ public static void enqueueExact(OOCStream out, IndexedMatrix callback.close(); } } + + public static ReservationBudget reserveBudget(MemoryAllowance allowance, long bytes) { + allowance.reserveBlocking(bytes); + return new ReservationBudget(allowance, bytes); + } } 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 0ce1e1ad38c..be6886ce1a6 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 @@ -26,6 +26,7 @@ import org.apache.sysds.common.Types.FileFormat; import org.apache.sysds.common.Types.ValueType; import org.apache.sysds.runtime.controlprogram.caching.MatrixObject; +import org.apache.sysds.runtime.instructions.ooc.CachingStream; import org.apache.sysds.runtime.instructions.ooc.OOCStream; import org.apache.sysds.runtime.instructions.ooc.SubscribableTaskQueue; import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; @@ -119,6 +120,42 @@ public void testDataGenMapTransposePipeline() { values); } + @Test + public void testJoinOutOfOrder() { + SubscribableTaskQueue left = new SubscribableTaskQueue<>(); + SubscribableTaskQueue right = new SubscribableTaskQueue<>(); + SubscribableTaskQueue joined = new SubscribableTaskQueue<>(); + SubscribableTaskQueue addends = new SubscribableTaskQueue<>(); + SubscribableTaskQueue output = new SubscribableTaskQueue<>(); + for(SubscribableTaskQueue stream : List.of(left, right, joined, addends, output)) + stream.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(1, 2, 1), FileFormat.BINARY))); + CachingStream cachedLeft = new CachingStream(left); + left.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 10d))); + left.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 2), new MatrixBlock(1, 1, 20d))); + right.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 2), new MatrixBlock(1, 1, 2d))); + right.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 1d))); + addends.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(1, 1, 100d))); + addends.enqueue(new IndexedMatrixValue(new MatrixIndexes(1, 2), new MatrixBlock(1, 1, 200d))); + left.closeInput(); + right.closeInput(); + addends.closeInput(); + OOCInstructionUtils.equiJoin(cachedLeft, right, joined, + (l, r) -> new MatrixBlock(1, 1, l.get(0, 0) + r.get(0, 0)), new StreamContext()); + OOCInstructionUtils.equiJoin(joined, addends, output, + (l, r) -> new MatrixBlock(1, 1, l.get(0, 0) + r.get(0, 0)), new StreamContext()); + + output.start(); + Map values = new HashMap<>(); + OOCStream.QueueCallback callback; + while((callback = output.dequeueCB()) != null) + try(OOCStream.QueueCallback current = callback) { + values.put(current.get().getIndexes().getColumnIndex(), current.get().getValue().get(0, 0)); + } + Assert.assertEquals(Map.of(1L, 111.0, 2L, 222.0), values); + cachedLeft.scheduleDeletion(); + } + private static final class TestPrimitive extends OOCPrimitive { private int _executions;