diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStream.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStream.java index 49a41ce2365..f5f44fdd3f0 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStream.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/OOCStream.java @@ -51,8 +51,8 @@ static QueueCallback eos(DMLRuntimeException e) { default void start() { OOCPrimitive primitive = getPrimitive(); - if(primitive != null) - primitive.tryStartExecution(); + if(primitive != null && !primitive.hasStartedExecution()) + primitive.start(); } interface QueueCallback extends AutoCloseable { 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 new file mode 100644 index 00000000000..3db13c2691a --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/planning/OOCPlanner.java @@ -0,0 +1,54 @@ +/* + * 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.planning; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Set; + +import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; + +public final class OOCPlanner { + public static void compile(OOCPrimitive root) { + List primitives = new ArrayList<>(); + collect(root, Collections.newSetFromMap(new IdentityHashMap<>()), primitives); + if(primitives.isEmpty()) + return; + + for(int i = primitives.size() - 1; i >= 0; i--) + if(primitives.get(i).getAccessPattern().isUnset()) + primitives.get(i).inferPatterns(); + if(root.getAccessPattern() == OOCAccessPattern.ANY || root.getAccessPattern().isUnset()) + root.requestPattern(OOCAccessPattern.ROW_MAJOR); + + for(OOCPrimitive primitive : primitives) + primitive.tryStartExecution(); + } + + private static void collect(OOCPrimitive primitive, Set visited, List result) { + if(primitive.hasStartedExecution() || !visited.add(primitive)) + return; + result.add(primitive); + for(OOCPrimitive child : primitive.getChildren()) + collect(child, visited, result); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java new file mode 100644 index 00000000000..fe1285a63fa --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/MappingOOCPrimitive.java @@ -0,0 +1,76 @@ +/* + * 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.function.Function; + +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; + +public class MappingOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _input; + private final OOCStreamable _output; + private final Function _operation; + + public MappingOOCPrimitive(OOCStreamable input, OOCStreamable output, + Function operation, StreamContext context) { + this(input.getPrimitive(), input, output, operation, context); + } + + private MappingOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable input, + OOCStreamable output, Function operation, + StreamContext context) { + super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive)); + _input = input; + _output = output; + _operation = operation; + } + + @Override + protected void inferPatternsInternal() { + OOCAccessPattern inputPattern = getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0) + .getAccessPattern(); + _pattern = _pattern.preferred(inputPattern); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = _pattern.preferred(accessPattern); + if(!getChildren().isEmpty()) + getChildren().get(0).requestPattern(accessPattern); + } + + @Override + protected void startExecution() { + OOCStream input = _input.getReadStream(); + OOCStream output = _output.getWriteStream(); + OOCInstructionUtils + .submitAdmittedOOCTasks(input, output, + value -> new IndexedMatrixValue(value.getIndexes(), _operation.apply(value)), _allowance, getContext()) + .thenRun(this::onComplete); + } +} 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 67c3072de50..aacb35acdad 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 @@ -23,15 +23,24 @@ import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.planning.OOCPlanner; +import org.apache.sysds.runtime.ooc.stream.StreamContext; public abstract class OOCPrimitive { + private final StreamContext _context; private final List _children; private final List _parents; + private final AtomicBoolean _started; private final AtomicBoolean _executionStarted; protected OOCAccessPattern _pattern; + protected MemoryAllowance _allowance; - protected OOCPrimitive(List children) { + protected OOCPrimitive(StreamContext context, List children) { + _context = context; _parents = new ArrayList<>(); List uniqueChildren = new ArrayList<>(children.size()); for(OOCPrimitive child : children) { @@ -41,10 +50,15 @@ protected OOCPrimitive(List children) { child.addParent(this); } _children = List.copyOf(uniqueChildren); + _started = new AtomicBoolean(); _executionStarted = new AtomicBoolean(); _pattern = OOCAccessPattern.UNSET; } + public final StreamContext getContext() { + return _context; + } + public final List getChildren() { return _children; } @@ -72,9 +86,30 @@ public final boolean hasStartedExecution() { return _executionStarted.get(); } + public final void start() { + if(_started.compareAndSet(false, true)) + OOCPlanner.compile(this); + } + public final void tryStartExecution() { - if(_executionStarted.compareAndSet(false, true)) + if(_executionStarted.compareAndSet(false, true)) { + _allowance = new SyncMemoryAllowance(GlobalMemoryBroker.get()); startExecution(); + } + } + + public final void onComplete() { + _allowance.shutdown(); + } + + public final void inferPatterns() { + if(!hasStartedExecution()) + inferPatternsInternal(); + } + + public final void requestPattern(OOCAccessPattern accessPattern) { + if(!hasStartedExecution() && _pattern != accessPattern) + requestPatternInternal(accessPattern); } private static boolean containsIdentity(List primitives, OOCPrimitive primitive) { @@ -86,7 +121,7 @@ private static boolean containsIdentity(List primitives, OOCPrimit protected abstract void startExecution(); - public abstract void inferPatterns(); + protected abstract void inferPatternsInternal(); - public abstract void requestPattern(OOCAccessPattern accessPattern); + protected abstract void requestPatternInternal(OOCAccessPattern accessPattern); } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java new file mode 100644 index 00000000000..83ffcd2acf4 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/PlannableDataGenOOCPrimitive.java @@ -0,0 +1,91 @@ +/* + * 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.function.Function; + +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.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.ooc.memory.ReservationBudget; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; +import org.apache.sysds.runtime.ooc.util.OOCUtils; + +public class PlannableDataGenOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _output; + private final Function _operation; + + public PlannableDataGenOOCPrimitive(OOCStreamable output, + Function operation, StreamContext context) { + super(context, List.of()); + _output = output; + _operation = operation; + } + + @Override + protected void inferPatternsInternal() { + if(_pattern.isUnset()) + _pattern = OOCAccessPattern.ANY; + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = accessPattern; + } + + @Override + protected void startExecution() { + OOCStream work = new SubscribableTaskQueue<>(); + OOCStream output = _output.getWriteStream(); + long outputBytes = OOCUtils.estimateOutputTileBytes(_output.getDataCharacteristics()); + AllocatedOOCStream admitted = new AllocatedOOCStream<>(work, _allowance, ignored -> outputBytes); + getContext().addOutStream(output); + OOCInstructionUtils.submitOOCTasks(admitted, callback -> { + ReservationBudget budget = AllocatedOOCStream.detachBudget(callback); + try { + MatrixIndexes indexes = callback.get(); + OOCUtils.enqueueExact(output, new IndexedMatrixValue(indexes, _operation.apply(indexes)), budget); + budget = null; + } + finally { + if(budget != null) + budget.close(); + } + }, getContext()).thenRun(output::closeInput).exceptionally(error -> { + output.propagateFailure(DMLRuntimeException.of(error)); + return null; + }).thenRun(this::onComplete); + + OOCInstructionUtils.submitOOCTask(() -> { + for(MatrixIndexes indexes : OOCUtils.getAccessPattern(_output.getDataCharacteristics(), _pattern)) + work.enqueue(indexes); + work.closeInput(); + }, new StreamContext().addOutStream(work)); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java new file mode 100644 index 00000000000..938c753a197 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/TransposeOOCPrimitive.java @@ -0,0 +1,76 @@ +/* + * 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.function.Function; + +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.ooc.OOCStreamable; +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.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; + +public class TransposeOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _input; + private final OOCStreamable _output; + private final Function _operation; + + public TransposeOOCPrimitive(OOCStreamable input, OOCStreamable output, + Function operation, StreamContext context) { + this(input.getPrimitive(), input, output, operation, context); + } + + private TransposeOOCPrimitive(OOCPrimitive inputPrimitive, OOCStreamable input, + OOCStreamable output, Function operation, StreamContext context) { + super(context, inputPrimitive == null ? List.of() : List.of(inputPrimitive)); + _input = input; + _output = output; + _operation = operation; + } + + @Override + protected void inferPatternsInternal() { + _pattern = (getChildren().isEmpty() ? OOCAccessPattern.ANY : getChildren().get(0).getAccessPattern()) + .transposed(); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + _pattern = accessPattern; + if(!getChildren().isEmpty()) + getChildren().get(0).requestPattern(accessPattern.transposed()); + } + + @Override + protected void startExecution() { + OOCStream input = _input.getReadStream(); + OOCStream output = _output.getWriteStream(); + OOCInstructionUtils.submitAdmittedOOCTasks(input, output, value -> { + MatrixIndexes indexes = value.getIndexes(); + return new IndexedMatrixValue(new MatrixIndexes(indexes.getColumnIndex(), indexes.getRowIndex()), + _operation.apply((MatrixBlock) value.getValue())); + }, _allowance, getContext()).thenRun(this::onComplete); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCInstructionUtils.java index 75f6f9db2e5..98bab492d53 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 @@ -33,8 +33,18 @@ import org.apache.sysds.api.DMLScript; 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.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; 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.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.stream.AllocatedOOCStream; import org.apache.sysds.runtime.ooc.stream.StreamContext; import org.apache.sysds.runtime.ooc.stream.TaskContext; import org.apache.sysds.runtime.util.CommonThreadPool; @@ -46,6 +56,31 @@ public final class OOCInstructionUtils { private static final AtomicInteger COMPUTE_IN_FLIGHT = new AtomicInteger(); private static final int COMPUTE_BACKPRESSURE_THRESHOLD = 100; + public static void dataGen(OOCStream output, Function operation, + StreamContext context) { + output.assignPrimitive(new PlannableDataGenOOCPrimitive(output, operation, context)); + } + + public static void equiMapBlock(OOCStreamable input, OOCStream output, + Function operation, StreamContext context) { + equiMap(input, output, value -> operation.apply((MatrixBlock) value.getValue()), context); + } + + public static void equiMap(OOCStreamable input, OOCStream output, + Function operation, StreamContext context) { + output.assignPrimitive(new MappingOOCPrimitive(input, output, operation, context)); + } + + public static void transposedMap(OOCStream input, OOCStream output, + Function operation, StreamContext context) { + output.assignPrimitive(new TransposeOOCPrimitive(input, output, operation, context)); + } + + public static void transpose(OOCStream input, OOCStream output, + StreamContext context) { + transposedMap(input, output, MatrixBlock::transpose, context); + } + public static int getComputeInFlight() { return COMPUTE_IN_FLIGHT.get(); } @@ -59,6 +94,31 @@ public static CompletableFuture submitOOCTasks(OOCStream queue, return submitOOCTasks(List.of(queue), (i, callback) -> consumer.accept(callback), null, null, context); } + public static CompletableFuture submitAdmittedOOCTasks(OOCStream in, + OOCStream out, Function operation, + MemoryAllowance allowance, StreamContext context) { + context.addOutStream(out); + long outputBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + AllocatedOOCStream admitted = new AllocatedOOCStream<>(in, allowance, + ignored -> outputBytes); + return submitOOCTasks(admitted, callback -> { + ReservationBudget budget = AllocatedOOCStream.detachBudget(callback); + try { + if(budget == null) + throw new DMLRuntimeException("Missing admitted output budget"); + OOCUtils.enqueueExact(out, operation.apply(callback.get()), budget); + budget = null; + } + finally { + if(budget != null) + budget.close(); + } + }, context).thenRun(out::closeInput).exceptionally(error -> { + out.propagateFailure(DMLRuntimeException.of(error)); + return null; + }); + } + public static CompletableFuture submitOOCTasks(OOCStream queue, Consumer> consumer, Function, Boolean> predicate, BiConsumer> onNotProcessed, StreamContext context) { 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 74e6ed30a8f..7b8a8ae46b1 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 @@ -20,16 +20,23 @@ package org.apache.sysds.runtime.ooc.util; import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.meta.DataCharacteristics; +import org.apache.sysds.runtime.instructions.ooc.OOCStream; +import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; 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.memory.InMemoryQueueCallback; import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.runtime.ooc.memory.ReservationBudget; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; import org.apache.sysds.runtime.util.IndexRange; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.Iterator; import java.util.List; import java.util.function.BooleanSupplier; @@ -106,4 +113,61 @@ public static long getNumBlocks(DataCharacteristics dc) { } return -1; } + + public static Iterable getAccessPattern(DataCharacteristics dc, OOCAccessPattern pattern) { + long rows = dc.getRows() == 0 ? 0 : dc.getNumRowBlocks(); + long cols = dc.getCols() == 0 ? 0 : dc.getNumColBlocks(); + return getAccessPattern(rows, cols, pattern); + } + + public static Iterable getAccessPattern(long rows, long cols, OOCAccessPattern pattern) { + return () -> new Iterator<>() { + private final long _size = rows * cols; + private long _position; + + @Override + public boolean hasNext() { + return _position < _size; + } + + @Override + public MatrixIndexes next() { + long position = _position++; + return pattern == OOCAccessPattern.COL_MAJOR ? new MatrixIndexes(position % rows + 1, + position / rows + 1) : new MatrixIndexes(position / cols + 1, position % cols + 1); + } + }; + } + + 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 estimateMatrixBlockBytes(Math.min(dc.getBlocksize(), dc.getRows()), + Math.min(dc.getBlocksize(), dc.getCols())); + } + + private static long estimateMatrixBlockBytes(long rows, long cols) { + return Math.max(MatrixBlock.estimateSizeDenseInMemory(rows, cols), + MatrixBlock.estimateSizeSparseInMemory(rows, cols, 1.0)); + } + + public static void enqueueExact(OOCStream out, IndexedMatrixValue value, + ReservationBudget budget) { + long bytes = ((MatrixBlock) value.getValue()).getExactSerializedSize(); + OOCStream.QueueCallback callback = null; + try { + budget.reserveBlocking(bytes); + callback = new InMemoryQueueCallback(value, null, budget, bytes); + budget.close(); + out.enqueue(callback); + callback = null; + } + finally { + budget.close(); + if(callback != null) + callback.close(); + } + } } 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 9b6e435ca60..40abcf26348 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 @@ -19,12 +19,24 @@ package org.apache.sysds.test.component.ooc; +import java.util.HashMap; import java.util.List; +import java.util.Map; +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.OOCStream; 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.meta.MatrixCharacteristics; +import org.apache.sysds.runtime.meta.MetaDataFormat; import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; import org.apache.sysds.runtime.ooc.primitives.OOCPrimitive; import org.apache.sysds.runtime.ooc.stream.FilteredOOCStream; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; import org.junit.Assert; import org.junit.Test; @@ -50,28 +62,67 @@ public void testGraphPatternsAndExecution() { filtered.start(); Assert.assertTrue(sink.hasStartedExecution()); Assert.assertEquals(1, sink._executions); + Assert.assertEquals(1, source._executions); + Assert.assertEquals(OOCAccessPattern.ROW_MAJOR, sink.getAccessPattern()); + sink.inferPatterns(); + sink.requestPattern(OOCAccessPattern.COL_MAJOR); + Assert.assertEquals(OOCAccessPattern.ROW_MAJOR, sink.getAccessPattern()); + } + + @Test + public void testDataGenMapTransposePipeline() { + SubscribableTaskQueue generated = new SubscribableTaskQueue<>(); + SubscribableTaskQueue mapped = new SubscribableTaskQueue<>(); + SubscribableTaskQueue transposed = new SubscribableTaskQueue<>(); + generated.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 3, 1), FileFormat.BINARY))); + mapped.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(2, 3, 1), FileFormat.BINARY))); + transposed.setData(new MatrixObject(ValueType.FP64, "/dev/null", + new MetaDataFormat(new MatrixCharacteristics(3, 2, 1), FileFormat.BINARY))); + + OOCInstructionUtils.dataGen(generated, + indexes -> new MatrixBlock(1, 1, (double) indexes.getRowIndex() * 10 + indexes.getColumnIndex()), + new StreamContext()); + OOCInstructionUtils.equiMapBlock(generated, mapped, input -> new MatrixBlock(1, 1, input.get(0, 0) + 1), + new StreamContext()); + OOCInstructionUtils.transpose(mapped, transposed, new StreamContext()); + + transposed.start(); + Map values = new HashMap<>(); + OOCStream.QueueCallback callback; + while((callback = transposed.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", 12.0, "2,1", 13.0, "3,1", 14.0, "1,2", 22.0, "2,2", 23.0, "3,2", 24.0), + values); } private static final class TestPrimitive extends OOCPrimitive { private int _executions; private TestPrimitive(List children) { - super(children); + super(null, children); } @Override protected void startExecution() { _executions++; + onComplete(); } @Override - public void inferPatterns() { + protected void inferPatternsInternal() { _pattern = OOCAccessPattern.ANY; inferParentPatterns(); } @Override - public void requestPattern(OOCAccessPattern accessPattern) { + protected void requestPatternInternal(OOCAccessPattern accessPattern) { _pattern = accessPattern; } }