From 2729c721573a0c82f5392a2a69f7f6ebcae34abd Mon Sep 17 00:00:00 2001 From: Jessica Priebe Date: Mon, 21 Sep 2026 14:42:34 +0200 Subject: [PATCH 1/2] add reshape primitive --- .../ooc/ReshapeOOCInstruction.java | 362 +------- .../ooc/cache/packed/OOCPackedCache.java | 2 +- .../ooc/primitives/ReshapeOOCPrimitive.java | 796 ++++++++++++++++++ .../runtime/ooc/util/OOCInstructionUtils.java | 6 + .../sysds/runtime/ooc/util/OOCUtils.java | 11 +- .../sysds/test/functions/ooc/ReshapeTest.java | 5 +- 6 files changed, 816 insertions(+), 366 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java diff --git a/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReshapeOOCInstruction.java b/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReshapeOOCInstruction.java index 7590438b949..65b077d3fc7 100644 --- a/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReshapeOOCInstruction.java +++ b/src/main/java/org/apache/sysds/runtime/instructions/ooc/ReshapeOOCInstruction.java @@ -23,16 +23,12 @@ import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.controlprogram.caching.MatrixObject; import org.apache.sysds.runtime.controlprogram.context.ExecutionContext; -import org.apache.sysds.runtime.data.DenseBlockFP64; 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.MatrixBlock; -import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.matrix.operators.Operator; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; -import java.util.ArrayList; -import java.util.concurrent.CompletableFuture; public class ReshapeOOCInstruction extends ComputationOOCInstruction { private final CPOperand _opRows; @@ -77,361 +73,7 @@ public void processInstruction(ExecutionContext ec) { MatrixObject in = ec.getMatrixObject(input1); OOCStream qIn = in.getStreamHandle(); - int blen = in.getBlocksize(); - long rlen = in.getNumRows(); - long clen = in.getNumColumns(); - if(rlen * clen != rows * cols) - throw new DMLRuntimeException("Reshape matrix requires consistent numbers of input/output cells (" + rlen - + ":" + clen + ", " + rows + ":" + cols + ")."); - - if(rlen == rows) { - mapOOC(qIn, qOut, tmp -> tmp); - return; - } - - if(clen <= blen && rlen <= blen && cols <= blen && rows <= blen) { - mapOOC(qIn, qOut, tmp -> { - MatrixBlock res = ((MatrixBlock) tmp.getValue()).reshape((int) rows, (int) cols, byRow); - return new IndexedMatrixValue(tmp.getIndexes(), res); - }); - return; - } - - int numBlocksPerRowIn = (int) Math.ceil((double) clen / blen); - int numBlocksPerColIn = (int) Math.ceil((double) rlen / blen); - int numBlocksPerRowOut = (int) Math.ceil((double) cols / blen); - int numBlocksPerColOut = (int) Math.ceil((double) rows / blen); - - if(byRow) { - OOCStream singleRowBlocks = new SubscribableTaskQueue<>(); - // split blocks into single rows and adapt index - CompletableFuture f = expandOOC(qIn, singleRowBlocks, tmp -> { - ArrayList out = new ArrayList<>(); - MatrixBlock blk = (MatrixBlock) tmp.getValue(); - for(int i = 0; i < blk.getNumRows(); i++) { - MatrixBlock slice = blk.slice(i, i); - long r = tmp.getIndexes().getRowIndex(); - long c = tmp.getIndexes().getColumnIndex(); - r = (r - 1) * blen + i + 1; - MatrixIndexes idx = new MatrixIndexes(r, c); - out.add(new IndexedMatrixValue(idx, slice)); - } - return out; - }); - - if(clen % blen == 0 && cols % blen == 0) { - // singleRowBlocks do not need to be split - if(rows == 1) { - // result is one single row - mapOOC(singleRowBlocks.getReadStream(), qOut, tmp -> { - long r = tmp.getIndexes().getRowIndex(); - long c = tmp.getIndexes().getColumnIndex(); - // adapt index to new position in row - return new IndexedMatrixValue(new MatrixIndexes(1, (r - 1) * numBlocksPerRowIn + c), tmp.getValue()); - }); - } - else { - f.join(); - reshapeFullColBlocks(rows, cols, blen, numBlocksPerRowIn, numBlocksPerRowOut, numBlocksPerColOut, singleRowBlocks, qOut); - } - } - else { - f.join(); - reshapePartialColBlocks(rlen, clen, rows, cols, blen, numBlocksPerRowIn, numBlocksPerRowOut, numBlocksPerColOut, singleRowBlocks, qOut); - } - } - else { - OOCStream singleColBlocks = new SubscribableTaskQueue<>(); - // split blocks into single cols and adapt index - CompletableFuture f = expandOOC(qIn, singleColBlocks, tmp -> { - ArrayList out = new ArrayList<>(); - MatrixBlock blk = (MatrixBlock) tmp.getValue(); - for(int i = 0; i < blk.getNumColumns(); i++) { - MatrixBlock slice = blk.slice(0, blk.getNumRows() - 1, i, i); - long r = tmp.getIndexes().getRowIndex(); - long c = tmp.getIndexes().getColumnIndex(); - c = (c - 1) * blen + i + 1; - MatrixIndexes idx = new MatrixIndexes(r, c); - out.add(new IndexedMatrixValue(idx, slice)); - } - return out; - }); - - if(rlen % blen == 0 && rows % blen == 0) { - // cols do not need to be split - if(cols == 1) { - // result is one single col - mapOOC(singleColBlocks.getReadStream(), qOut, tmp -> { - long r = tmp.getIndexes().getRowIndex(); - long c = tmp.getIndexes().getColumnIndex(); - // adapt index to new position in col - return new IndexedMatrixValue(new MatrixIndexes((c - 1) * numBlocksPerColIn + r, 1), tmp.getValue()); - }); - } - else { - f.join(); - reshapeFullRowBlocks(rows, cols, blen, numBlocksPerRowOut, numBlocksPerColIn, numBlocksPerColOut, singleColBlocks, qOut); - } - } - else { - f.join(); - reshapePartialRowBlocks(rlen, clen, rows, cols, blen, numBlocksPerRowOut, numBlocksPerColIn, numBlocksPerColOut, singleColBlocks, qOut); - } - } - } - - private void reshapeFullColBlocks(long rows, long cols, int blen, int numBlocksPerRowIn, int numBlocksPerRowOut, - int numBlocksPerColOut, OOCStream singleRowBlocks, OOCStream qOut) { - // use cache for accessing input rows by index - CachingStream singleRowBlockCache = new CachingStream(singleRowBlocks); - singleRowBlockCache.incrSubscriberCount(1); - singleRowBlockCache.scheduleDeletion(); - - // totalRowIdx corresponds to index of row block when all aligned in one row - // br * numBlocksPerRowOut * blen + b + r * numBlocksPerRowOut; - // with numBlocksPerRowOut * blen = cols - long totalIdx = -cols - 1 - numBlocksPerRowOut; - - // iterate through rows of output blocks - for(int br = 0; br < numBlocksPerColOut; br++) { - totalIdx += cols; - long tmp = totalIdx; - // for each block in row - for(int b = 0; b < numBlocksPerRowOut; b++) { - totalIdx += 1; - int localRows = (br == numBlocksPerColOut - 1 && rows % blen != 0) ? (int) rows % blen : blen; - MatrixBlock res = new MatrixBlock(localRows, blen, false); - long tmp2 = totalIdx; - // for each row in block - for(int r = 0; r < blen && r < localRows; r++) { - totalIdx += numBlocksPerRowOut; - // calc col idx for input - long colBlockIn = totalIdx % numBlocksPerRowIn + 1; - // calc row idx for input - long rowBlockIn = totalIdx / numBlocksPerRowIn + 1; - - try(OOCStream.QueueCallback cb = singleRowBlockCache - .findCached(new MatrixIndexes(rowBlockIn, colBlockIn))) { - MatrixBlock blk = (MatrixBlock) cb.get().getValue(); - res.setRow(r, blk.getDenseBlockValues()); - } - } - totalIdx = tmp2; - qOut.enqueue(new IndexedMatrixValue(new MatrixIndexes(br + 1, b + 1), res)); - } - totalIdx = tmp; - } - qOut.closeInput(); - } - - private void reshapePartialColBlocks(long rlen, long clen, long rows, long cols, int blen, int numBlocksPerRowIn, - int numBlocksPerRowOut, int numBlocksPerColOut, OOCStream singleRowBlocks, OOCStream qOut) { - // use cache for accessing input rows by index - CachingStream singleRowBlockCache = new CachingStream(singleRowBlocks); - singleRowBlockCache.incrSubscriberCount(1); - singleRowBlockCache.scheduleDeletion(); - - int br = 0; - int bc = 0; - int r = 0; - - // allocate row of output blocks - MatrixBlock[] outputBlockRow = allocateSliceBlocks(br, rows, cols, blen, numBlocksPerRowOut, numBlocksPerColOut,true); - - int offsetOut = 0; - int localColsOut = (cols > blen) ? blen : (int) cols; - // iterate through input rows and add to row of output blocks - for(int i = 1; i <= rlen; i++) { - for(int j = 1; j <= numBlocksPerRowIn; j++) { - try(OOCStream.QueueCallback qcb = singleRowBlockCache.findCached(new MatrixIndexes(i, j))) { - MatrixBlock blk = (MatrixBlock) qcb.get().getValue(); - - int offsetIn = 0; - int localColsIn = (j == numBlocksPerRowIn && clen % blen != 0) ? (int) clen % blen : blen; - while(offsetIn < localColsIn) { - // until input row fully processed - int remIn = localColsIn - offsetIn; - int remOut = localColsOut - offsetOut; - if(remIn < remOut) { - // next input - setOutputEntries(blk, outputBlockRow[bc], r, offsetIn, offsetOut, remIn, true); - offsetIn += remIn; - offsetOut += remIn; - continue; - } - else if(remIn == remOut) { - // next input and next row - setOutputEntries(blk, outputBlockRow[bc], r, offsetIn, offsetOut, remIn, true); - offsetIn += remIn; - } - else { - // next row - setOutputEntries(blk, outputBlockRow[bc], r, offsetIn, offsetOut, remOut, true); - offsetIn += remOut; - } - bc++; - offsetOut = 0; - if(bc == numBlocksPerRowOut) { - // next row - r++; - if(r == outputBlockRow[0].getNumRows()) { - // enqueue filled output blocks and allocate new ones - for(int b = 0; b < outputBlockRow.length; b++) - qOut.enqueue(new IndexedMatrixValue(new MatrixIndexes(br + 1, b + 1), outputBlockRow[b])); - br++; - // allocate new block row - outputBlockRow = allocateSliceBlocks(br, rows, cols, blen, numBlocksPerRowOut, numBlocksPerColOut, true); - r = 0; - } - bc = 0; - } - localColsOut = (bc == numBlocksPerRowOut - 1 && cols % blen != 0) ? (int) cols % blen : blen; - } - } - } - } - qOut.closeInput(); - } - - private void reshapeFullRowBlocks(long rows, long cols, int blen, int numBlocksPerRowOut, int numBlocksPerColIn, - int numBlocksPerColOut, OOCStream singleColBlocks, OOCStream qOut) { - // use cache for accessing input cols by index - CachingStream singleColBlockCache = new CachingStream(singleColBlocks); - singleColBlockCache.incrSubscriberCount(1); - singleColBlockCache.scheduleDeletion(); - - // totalColIdx corresponds to index of col block when all aligned in one col - // bc * numBlocksPerColOut * blen + b + c * numBlocksPerColOut; - // with numBlocksPerColOut * blen = rows - long totalIdx = -rows - 1 - numBlocksPerColOut; - - // iterate through cols of output blocks - for(int bc = 0; bc < numBlocksPerRowOut; bc++) { - totalIdx += rows; - long tmp = totalIdx; - // for each block in col - for(int b = 0; b < numBlocksPerColOut; b++) { - totalIdx += 1; - int localCols = (bc == numBlocksPerRowOut - 1 && cols % blen != 0) ? (int) cols % blen : blen; - MatrixBlock res = new MatrixBlock(blen, localCols, false); - res.allocateDenseBlock(); - long tmp2 = totalIdx; - // for each col in block - for(int c = 0; c < blen && c < localCols; c++) { - totalIdx += numBlocksPerColOut; - // calc col idx for input - long colBlockIn = totalIdx / numBlocksPerColIn + 1; - // calc row idx for input - long rowBlockIn = totalIdx % numBlocksPerColIn + 1; - - try(OOCStream.QueueCallback cb = singleColBlockCache - .findCached(new MatrixIndexes(rowBlockIn, colBlockIn))) { - MatrixBlock blk = (MatrixBlock) cb.get().getValue(); - res.getDenseBlock().set(0, blen, c, c + 1, blk.getDenseBlock()); - } - } - totalIdx = tmp2; - res.recomputeNonZeros(); - qOut.enqueue(new IndexedMatrixValue(new MatrixIndexes(b + 1, bc + 1), res)); - } - totalIdx = tmp; - } - qOut.closeInput(); - } - - private void reshapePartialRowBlocks(long rlen, long clen, long rows, long cols, int blen, int numBlocksPerRowOut, - int numBlocksPerColIn, int numBlocksPerColOut, OOCStream singleColBlocks, OOCStream qOut) { - // use cache for accessing input cols by index - CachingStream singleRowBlockCache = new CachingStream(singleColBlocks); - singleRowBlockCache.incrSubscriberCount(1); - singleRowBlockCache.scheduleDeletion(); - - int br = 0; - int bc = 0; - int c = 0; - - // allocate col of output blocks - MatrixBlock[] outputBlockCol = allocateSliceBlocks(bc, rows, cols, blen, numBlocksPerRowOut, numBlocksPerColOut, false); - - int offsetOut = 0; - int localRowsOut = (rows > blen) ? blen : (int) rows; - // iterate through input cols and add to col of output blocks - for(int j = 1; j <= clen; j++) { - for(int i = 1; i <= numBlocksPerColIn; i++) { - try(OOCStream.QueueCallback qcb = singleRowBlockCache.findCached(new MatrixIndexes(i, j))) { - MatrixBlock blk = (MatrixBlock) qcb.get().getValue(); - - int offsetIn = 0; - int localRowsIn = (i == numBlocksPerColIn && rlen % blen != 0) ? (int) rlen % blen : blen; - while(offsetIn < localRowsIn) { - // until input col fully processed - int remIn = localRowsIn - offsetIn; - int remOut = localRowsOut - offsetOut; - if(remIn < remOut) { - // next input - setOutputEntries(blk, outputBlockCol[br], c, offsetIn, offsetOut, remIn, false); - offsetIn += remIn; - offsetOut += remIn; - continue; - } - else if(remIn == remOut) { - // next input and next col - setOutputEntries(blk, outputBlockCol[br], c, offsetIn, offsetOut, remIn, false); - offsetIn += remIn; - } - else { - // next col - setOutputEntries(blk, outputBlockCol[br], c, offsetIn, offsetOut, remOut, false); - offsetIn += remOut; - } - br++; - offsetOut = 0; - if(br == numBlocksPerColOut) { - // next col - c++; - if(c == outputBlockCol[0].getNumColumns()) { - // enqueue filled output blocks and allocate new ones - for(int b = 0; b < outputBlockCol.length; b++) { - outputBlockCol[b].recomputeNonZeros(); - qOut.enqueue(new IndexedMatrixValue(new MatrixIndexes(b + 1, bc + 1), outputBlockCol[b])); - } - bc++; - // allocate new block col - outputBlockCol = allocateSliceBlocks(bc, rows, cols, blen, numBlocksPerRowOut, numBlocksPerColOut, false); - c = 0; - } - br = 0; - } - localRowsOut = (br == numBlocksPerColOut - 1 && rows % blen != 0) ? (int) rows % blen : blen; - } - } - } - } - qOut.closeInput(); - } - - private MatrixBlock[] allocateSliceBlocks(int idx, long rows, long cols, int blen, int numBlocksPerRow, int numBlocksPerCol, boolean isBlockRowSlice) { - int num = isBlockRowSlice ? numBlocksPerRow : numBlocksPerCol; - MatrixBlock[] res = new MatrixBlock[num]; - - // full inner blocks, adjust for outer blocks - int localRows = ((!isBlockRowSlice || idx == numBlocksPerCol - 1) && rows % blen != 0) ? (int) rows % blen : blen; - int localCols = ((isBlockRowSlice || idx == numBlocksPerRow - 1) && cols % blen != 0) ? (int) cols % blen : blen; - - for(int k = 0; k < num - 1; k++) { - res[k] = isBlockRowSlice ? new MatrixBlock(localRows, blen, false) : new MatrixBlock(blen, localCols, false); - res[k].allocateDenseBlock(); - } - res[num - 1] = new MatrixBlock(localRows, localCols, false); - res[num - 1].allocateDenseBlock(); - return res; - } - - private void setOutputEntries(MatrixBlock src, MatrixBlock dest, int idx, int srcOffset, int destOffset, int length, boolean rowWise) { - if(rowWise) - ((DenseBlockFP64) dest.getDenseBlock()).setPartialRow(src.getDenseBlock(), idx, srcOffset, destOffset, length); - else - ((DenseBlockFP64) dest.getDenseBlock()).setPartialCol(src.getDenseBlock(), idx, srcOffset, destOffset, length); + OOCInstructionUtils.reshape(qIn, qOut, rows, cols, byRow, getContext()); } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java index bdc4d69f263..13635b86b8e 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java @@ -47,7 +47,7 @@ public final class OOCPackedCache implements OOCCache { private static final long PACKED_STREAM_ID = CachingStream._streamSeq.getNextID(); - private static final long DEFAULT_PACK_THRESHOLD_BYTES = 1L << 18; + private static final long DEFAULT_PACK_THRESHOLD_BYTES = 1; private static final long DEFAULT_PACK_TARGET_BYTES = 1L << 19; // 512 KB tile packing private static final long DEFAULT_MAX_STAGING_BYTES = 1L << 26; private static final int DEFAULT_MAX_OPEN_BUILDERS = 64; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java new file mode 100644 index 00000000000..13c8a1b2756 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java @@ -0,0 +1,796 @@ +/* + * 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 org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.data.DenseBlockFP64; +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.spark.data.IndexedMatrixValue; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.meta.DataCharacteristics; +import org.apache.sysds.runtime.ooc.cache.OOCCacheManager; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.memory.ManagedPayload; +import org.apache.sysds.runtime.ooc.memory.ReservationBudget; +import org.apache.sysds.runtime.ooc.planning.OOCAccessPattern; +import org.apache.sysds.runtime.ooc.store.StateTable; +import org.apache.sysds.runtime.ooc.store.StoreLease; +import org.apache.sysds.runtime.ooc.stream.AllocatedOOCStream; +import org.apache.sysds.runtime.ooc.stream.StreamContext; +import org.apache.sysds.runtime.ooc.util.OOCInstructionUtils; +import org.apache.sysds.runtime.ooc.util.OOCUtils; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +public class ReshapeOOCPrimitive extends OOCPrimitive { + private final OOCStreamable _input; + private final OOCStreamable _output; + private final boolean _byRow; + private final long _rows; + private final long _cols; + private long _rlen; + private long _clen; + private int _blen; + + public ReshapeOOCPrimitive(OOCStreamable input, OOCStreamable output, + long rows, long cols, boolean byRow, StreamContext context) { + super(context, input); + _input = input; + _output = output; + _byRow = byRow; + _rows = rows; + _cols = cols; + _pattern = byRow ? OOCAccessPattern.ROW_MAJOR : OOCAccessPattern.COL_MAJOR; + } + + @Override + protected void inferPatternsInternal() { + for(OOCPrimitive child : getChildren()) + child.requestPattern(_pattern); + inferParentPatterns(); + } + + @Override + protected void requestPatternInternal(OOCAccessPattern accessPattern) { + for(OOCPrimitive child : getChildren()) + child.requestPattern(_pattern); + } + + @Override + protected void startExecution() { + DataCharacteristics inputDc = _input.getDataCharacteristics(); + if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) + throw new DMLRuntimeException("Reshape OOC reduction requires known input dimensions and block size."); + + OOCStream input = getInputReadStream(0); + OOCStream output = _output.getWriteStream(); + getContext().addOutStream(output); + + _rlen = inputDc.getRows(); + _clen = inputDc.getCols(); + _blen = Math.toIntExact(inputDc.getBlocksize()); + + if(_rlen * _clen != _rows * _cols) { + onComplete(); + throw new DMLRuntimeException("Reshape matrix requires consistent numbers of input/output cells (" + _rlen + + ":" + _clen + ", " + _rows + ":" + _cols + ")."); + } + + if(_rlen == _rows) { + OOCInstructionUtils + .submitAdmittedOOCTasks(input, output, + value -> new IndexedMatrixValue(value.getIndexes(), value.getValue()), _allowance, getContext()) + .thenRun(this::onComplete); + return; + } + + if(_clen <= _blen && _rlen <= _blen && _cols <= _blen && _rows <= _blen) { + OOCInstructionUtils.submitAdmittedOOCTasks(input, output, + value -> new IndexedMatrixValue(value.getIndexes(), + ((MatrixBlock) value.getValue()).reshape((int) _rows, (int) _cols, _byRow)), + _allowance, getContext()).thenRun(this::onComplete); + return; + } + + int numColBlocksIn = Math.toIntExact(inputDc.getNumColBlocks()); + int numRowBlocksIn = Math.toIntExact(inputDc.getNumRowBlocks()); + int numColBlocksOut = (int) Math.ceil((double) _cols / _blen); + int numRowBlocksOut = (int) Math.ceil((double) _rows / _blen); + + long sliceBytes = (OOCUtils.estimateFullTileBytes(input.getDataCharacteristics()) + + (_blen - 1) * MatrixBlock.getHeaderSize()) / _blen; + + StateTable table = new StateTable<>(OOCCacheManager.getGlobalCache(), + CachingStream._streamSeq.getNextID()); + + if(_byRow) { + if(_clen % _blen == 0 && _cols % _blen == 0) { + // singleRowBlocks do not need to be split + if(_rows == 1) { + // result is one single row + submitSingleRowColTask(input, output, numColBlocksIn, numRowBlocksIn); + } + else { + CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask( + () -> reshapeFullColBlocks(table, output, numColBlocksOut, numRowBlocksOut, sliceBytes), + getContext())); + } + } + else { + CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(() -> reshapePartialColBlocks(table, output, + numColBlocksIn, numColBlocksOut, numRowBlocksOut, sliceBytes), getContext())); + } + } + else { + if(_rlen % _blen == 0 && _rows % _blen == 0) { + // singleColBlocks do not need to be split + if(_cols == 1) { + // result is one single col + submitSingleRowColTask(input, output, numColBlocksIn, numRowBlocksIn); + } + else { + CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask( + () -> reshapeFullRowBlocks(table, output, numColBlocksOut, numRowBlocksOut, sliceBytes), + getContext())); + } + } + else { + CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(() -> reshapePartialRowBlocks(table, output, + numRowBlocksIn, numColBlocksOut, numRowBlocksOut, sliceBytes), getContext())); + } + } + } + + private CompletableFuture splitIntoTable(OOCStream in, + StateTable table, int numColBlocks, int numRowBlocks) { + + long blockBytes = OOCUtils.estimateFullTileBytes(in.getDataCharacteristics()) + + (_blen - 1) * MatrixBlock.getHeaderSize(); + long singleSliceBytes = blockBytes / _blen; + + AllocatedOOCStream allocated = new AllocatedOOCStream<>(in, _allowance, ignored -> blockBytes); + + return OOCInstructionUtils.submitOOCTasks(allocated, callback -> { + try(ReservationBudget budget = AllocatedOOCStream.detachBudget(callback)) { + if(budget == null) + throw new DMLRuntimeException("Missing admitted output budget"); + + IndexedMatrixValue imv = callback.get(); + MatrixBlock blk = (MatrixBlock) imv.getValue(); + long r = imv.getIndexes().getRowIndex(); + long c = imv.getIndexes().getColumnIndex(); + long rIdx; + long cIdx; + + int n = _byRow ? blk.getNumRows() : blk.getNumColumns(); + for(int i = 0; i < n; i++) { + MatrixBlock slice; + if(_byRow) { + slice = blk.slice(i, i); + rIdx = (r - 1) * _blen + i + 1; + cIdx = c; + } + else { + slice = blk.slice(0, blk.getNumRows() - 1, i, i); + cIdx = (c - 1) * _blen + i + 1; + rIdx = r; + } + + long targetIdx = _byRow ? (rIdx - 1) * numColBlocks + c - 1 : (cIdx - 1) * numRowBlocks + r - 1; + IndexedMatrixValue sliceImv = new IndexedMatrixValue(new MatrixIndexes(rIdx, cIdx), slice); + budget.reserveBlocking(singleSliceBytes); + table.put((int) targetIdx, new ManagedPayload<>(sliceImv, singleSliceBytes, budget)); + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + }, getContext()); + } + + private void submitSingleRowColTask(OOCStream in, OOCStream out, + int numColBlocks, int numRowBlocks) { + + // one input block is split into blen output blocks + long outputBytes = _blen * OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + + AllocatedOOCStream allocated = new AllocatedOOCStream<>(in, _allowance, ignored -> outputBytes); + + OOCInstructionUtils.submitOOCTasks(allocated, callback -> { + try(ReservationBudget budget = AllocatedOOCStream.detachBudget(callback)) { + if(budget == null) + throw new DMLRuntimeException("Missing admitted output budget"); + + IndexedMatrixValue imv = callback.get(); + MatrixBlock blk = (MatrixBlock) imv.getValue(); + long r = imv.getIndexes().getRowIndex(); + long c = imv.getIndexes().getColumnIndex(); + long rIdx; + long cIdx; + + int n = _byRow ? blk.getNumRows() : blk.getNumColumns(); + for(int i = 0; i < n; i++) { + MatrixBlock slice; + if(_byRow) { + // split and adjust idx + slice = blk.slice(i, i); + // total row, 1 based + rIdx = (r - 1) * _blen + i + 1; + // all in single row + cIdx = (rIdx - 1) * numColBlocks + c; + rIdx = 1; + } + else { + // split and adjust idx + slice = blk.slice(0, blk.getNumRows() - 1, i, i); + cIdx = (c - 1) * _blen + i + 1; + // all in single col + rIdx = (cIdx - 1) * numRowBlocks + r; + cIdx = 1; + } + + IndexedMatrixValue sliceImv = new IndexedMatrixValue(new MatrixIndexes(rIdx, cIdx), slice); + OOCUtils.enqueueExact(out, sliceImv, budget, false); + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + }, getContext()).thenRun(this::onComplete).thenRun(out::closeInput).exceptionally(error -> { + out.propagateFailure(DMLRuntimeException.of(error)); + return null; + }); + } + + private void reshapeFullColBlocks(StateTable table, OOCStream out, + int numColBlocksOut, int numRowBlocksOut, long rowBytes) { + + List>> futures = new ArrayList<>(); + ReservationBudget budget = null; + + long numRowsBlockOut = Math.min(out.getDataCharacteristics().getRows(), _blen); + long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + long outputBytes = numRowsBlockOut * rowBytes + blockBytes; + + // totalRowIdx corresponds to index of row block when all aligned in one row + // br * numColBlocksOut * blen + b + r * numColBlocksOut; + // with numColBlocksOut * blen = cols + long totalIdx = -_cols - 1 - numColBlocksOut; + + try { + // iterate through rows of output blocks + for(int br = 0; br < numRowBlocksOut; br++) { + totalIdx += _cols; + long tmp = totalIdx; + int localRows = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + // for each block in row + for(int b = 0; b < numColBlocksOut; b++) { + totalIdx += 1; + long tmp2 = totalIdx; + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + // for each row in block + for(int r = 0; r < _blen && r < localRows; r++) { + totalIdx += numColBlocksOut; + OOCFuture> rowFuture = table.take((int) totalIdx, budget); + futures.add(rowFuture); + } + totalIdx = tmp2; + OOCFuture>> future = OOCFuture.allOf(futures, StoreLease::close); + MatrixIndexes idx = new MatrixIndexes(br + 1, b + 1); + + ReservationBudget finalBudget = budget; + future.whenComplete((leases, error) -> { + MatrixBlock block = new MatrixBlock(localRows, _blen, false); + for(int r = 0; r < leases.size(); r++) { + StoreLease lease = leases.get(r); + MatrixBlock row = (MatrixBlock) lease.value().getValue(); + block.setRow(r, row.getDenseBlockValues()); + lease.close(); + } + block.recomputeNonZeros(); + OOCUtils.enqueueExact(out, new IndexedMatrixValue(idx, block), finalBudget, true); + futures.clear(); + }); + budget = null; + } + totalIdx = tmp; + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + if(budget != null) { + budget.close(); + } + try { + table.close(); + onComplete(); + } + finally { + out.closeInput(); + } + } + } + + private void reshapeFullRowBlocks(StateTable table, OOCStream out, + int numColBlocksOut, int numRowBlocksOut, long colBytes) { + + List>> futures = new ArrayList<>(); + ReservationBudget budget = null; + + long numColsBlockOut = Math.min(out.getDataCharacteristics().getCols(), _blen); + long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + long outputBytes = numColsBlockOut * colBytes + blockBytes; + + // totalColIdx corresponds to index of col block when all aligned in one col + // bc * numRowBlocksOut * blen + b + c * numRowBlocksOut; + // with numRowBlocksOut * blen = rows + long totalIdx = -_rows - 1 - numRowBlocksOut; + + try { + // iterate through cols of output blocks + for(int bc = 0; bc < numColBlocksOut; bc++) { + totalIdx += _rows; + long tmp = totalIdx; + int localCols = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + // for each block in col + for(int b = 0; b < numRowBlocksOut; b++) { + totalIdx += 1; + long tmp2 = totalIdx; + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + // for each col in block + for(int c = 0; c < _blen && c < localCols; c++) { + totalIdx += numRowBlocksOut; + OOCFuture> colFuture = table.take((int) totalIdx, budget); + futures.add(colFuture); + } + totalIdx = tmp2; + OOCFuture>> future = OOCFuture.allOf(futures, StoreLease::close); + MatrixIndexes idx = new MatrixIndexes(b + 1, bc + 1); + + ReservationBudget finalBudget = budget; + future.whenComplete((leases, error) -> { + MatrixBlock block = new MatrixBlock(_blen, localCols, false); + block.allocateDenseBlock(); + for(int c = 0; c < leases.size(); c++) { + StoreLease lease = leases.get(c); + MatrixBlock col = (MatrixBlock) lease.value().getValue(); + block.getDenseBlock().set(0, _blen, c, c + 1, col.getDenseBlock()); + lease.close(); + } + block.recomputeNonZeros(); + OOCUtils.enqueueExact(out, new IndexedMatrixValue(idx, block), finalBudget, true); + futures.clear(); + }); + budget = null; + } + totalIdx = tmp; + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + if(budget != null) { + budget.close(); + } + try { + table.close(); + onComplete(); + } + finally { + out.closeInput(); + } + } + } + + private void reshapePartialColBlocks(StateTable table, OOCStream out, + int numColBlocksIn, int numColBlocksOut, int numRowBlocksOut, long rowBytes) { + + long numRowsBlockOut = Math.min(_rows, _blen); + long numColsBlockOut = Math.min(_cols, _blen); + + long numNeededRowsIn = 2 + + (long) Math.ceil((((double) numRowsBlockOut * numColsBlockOut) / _clen) * numColBlocksIn * numColBlocksOut); + long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + long outputBytes = numNeededRowsIn * rowBytes + numColBlocksOut * blockBytes; + + ReservationBudget budget = null; + int br = 0; + + try { + List>> futures = new ArrayList<>(); + // new row of output blocks + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + int missing = getNumMissingRowSlices(1, br, 0, numColBlocksIn, numRowBlocksOut); + int startJ = 1; + final int[] offset = {0}; + + // iterate through input rows and add to row of output blocks + for(int i = 1; i <= _rlen; i++) { + for(int j = 1; j <= numColBlocksIn; j++) { + int totalIdx = (i - 1) * numColBlocksIn + j - 1; + OOCFuture> blkFuture = table.take(totalIdx, budget); + futures.add(blkFuture); + + if(futures.size() < missing) + continue; + + final ReservationBudget finalBudget = budget; + final int finalJ = startJ; + final int finalBr = br; + + OOCFuture>> future = OOCFuture.allOf(futures, StoreLease::close); + future.whenComplete((leases, error) -> { + int localJ = finalJ; + int offsetIn = offset[0]; + + int bc = 0; + int r = 0; + int offsetOut = 0; + MatrixBlock[] outputBlockRow = allocateSliceBlocks(finalBr, numColBlocksOut, numRowBlocksOut); + int localColsOut = (_cols > _blen) ? _blen : (int) _cols; + + for(int k = 0; k < leases.size(); k++) { + StoreLease lease = leases.get(k); + IndexedMatrixValue slice = lease.value(); + MatrixBlock sliceVal = (MatrixBlock) slice.getValue(); + + int localColsIn = (localJ == numColBlocksIn && _clen % _blen != 0) ? (int) _clen % _blen : _blen; + while(offsetIn < localColsIn) { + // until input row fully processed + int remIn = localColsIn - offsetIn; + int remOut = localColsOut - offsetOut; + if(remIn < remOut) { + // next input + setOutputEntries(sliceVal, outputBlockRow[bc], r, offsetIn, offsetOut, remIn); + offsetIn += remIn; + offsetOut += remIn; + continue; + } + else if(remIn == remOut) { + // next input and next row + setOutputEntries(sliceVal, outputBlockRow[bc], r, offsetIn, offsetOut, remIn); + offsetIn += remIn; + } + else { + // next row + setOutputEntries(sliceVal, outputBlockRow[bc], r, offsetIn, offsetOut, remOut); + offsetIn += remOut; + } + bc++; + offsetOut = 0; + if(bc == numColBlocksOut) { + // next row + r++; + if(r == outputBlockRow[0].getNumRows()) { + offset[0] = offsetIn == localColsIn ? 0 : offsetIn; + break; + } + bc = 0; + } + localColsOut = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + } + localJ++; + if(localJ == numColBlocksIn + 1) + localJ = 1; + + lease.close(); + offsetIn = 0; + + if(k == leases.size() - 1 && offset[0] != 0) { + // put current slice back into table, to be able to reserve new budget + finalBudget.reserveBlocking(rowBytes); + table.put(totalIdx, new ManagedPayload<>(slice, rowBytes, finalBudget)); + } + } + + // enqueue filled output blocks and allocate new ones + for(int b = 0; b < outputBlockRow.length; b++) { + outputBlockRow[b].recomputeNonZeros(); + OOCUtils.enqueueExact(out, + new IndexedMatrixValue(new MatrixIndexes(finalBr + 1, b + 1), outputBlockRow[b]), + finalBudget, false); + } + }); + + budget.close(); + futures.clear(); + + if(br == numRowBlocksOut - 1) + break; + + // new block row + br++; + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + + if(offset[0] != 0) { + // get slice back from table + blkFuture = table.take(totalIdx, budget); + futures.add(blkFuture); + startJ = j; + } + else { + startJ = (j == numColBlocksIn) ? 1 : j + 1; + } + + missing = getNumMissingRowSlices(startJ, br, offset[0], numColBlocksIn, numRowBlocksOut); + } + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + if(budget != null) { + budget.close(); + } + try { + table.close(); + onComplete(); + } + finally { + out.closeInput(); + } + } + } + + private int getNumMissingRowSlices(int j, int br, int offsetIn, int numColBlocksIn, int numRowBlocksOut) { + long localColsIn = (j == numColBlocksIn && _clen % _blen != 0) ? _clen % _blen : _blen; + long localRowsOut = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? _rows % _blen : _blen; + + long totalIdx = _cols * localRowsOut; + int cnt = 0; + + if(offsetIn != 0) { + // reuse table entry + totalIdx -= (localColsIn - offsetIn); + cnt++; + } + + int numRows = (int) Math.floor((double) totalIdx / _clen); + cnt += numRows * numColBlocksIn; + totalIdx -= numRows * _clen; + + long numBlenSlices = Math.max(0, numColBlocksIn - (j - 1)); + long restRow = numBlenSlices * _blen + localColsIn; + if(totalIdx - restRow > 0) { + totalIdx -= restRow; + cnt += numColBlocksIn - j; + } + cnt += (int) Math.ceil((double) totalIdx / _blen); + + return cnt; + } + + private void reshapePartialRowBlocks(StateTable table, OOCStream out, + int numRowBlocksIn, int numColBlocksOut, int numRowBlocksOut, long colBytes) { + + long numRowsBlockOut = Math.min(_rows, _blen); + long numColsBlockOut = Math.min(_cols, _blen); + + long numNeededColsIn = 2 + + (long) Math.ceil((((double) numRowsBlockOut * numColsBlockOut) / _rlen) * numRowBlocksIn * numRowBlocksOut); + long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); + long outputBytes = numNeededColsIn * colBytes + numRowBlocksOut * blockBytes; + + ReservationBudget budget = null; + int bc = 0; + + try { + List>> futures = new ArrayList<>(); + // new col of output blocks + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + int missing = getNumMissingColSlices(1, bc, 0, numRowBlocksIn, numColBlocksOut); + int startI = 1; + final int[] offset = {0}; + + // iterate through input cols and add to col of output blocks + for(int j = 1; j <= _clen; j++) { + for(int i = 1; i <= numRowBlocksIn; i++) { + int totalIdx = (j - 1) * numRowBlocksIn + i - 1; + OOCFuture> blkFuture = table.take(totalIdx, budget); + futures.add(blkFuture); + + if(futures.size() < missing) + continue; + + final ReservationBudget finalBudget = budget; + final int finalI = startI; + final int finalBc = bc; + + OOCFuture>> future = OOCFuture.allOf(futures, StoreLease::close); + future.whenComplete((leases, error) -> { + int localI = finalI; + int offsetIn = offset[0]; + + int br = 0; + int c = 0; + int offsetOut = 0; + + MatrixBlock[] outputBlockCol = allocateSliceBlocks(finalBc, numColBlocksOut, numRowBlocksOut); + int localRowsOut = (_rows > _blen) ? _blen : (int) _rows; + + for(int k = 0; k < leases.size(); k++) { + StoreLease lease = leases.get(k); + IndexedMatrixValue slice = lease.value(); + MatrixBlock sliceVal = (MatrixBlock) slice.getValue(); + + int localRowsIn = (localI == numRowBlocksIn && _rlen % _blen != 0) ? (int) _rlen % _blen : _blen; + while(offsetIn < localRowsIn) { + // until input col fully processed + int remIn = localRowsIn - offsetIn; + int remOut = localRowsOut - offsetOut; + if(remIn < remOut) { + // next input + setOutputEntries(sliceVal, outputBlockCol[br], c, offsetIn, offsetOut, remIn); + offsetIn += remIn; + offsetOut += remIn; + continue; + } + else if(remIn == remOut) { + // next input and next col + setOutputEntries(sliceVal, outputBlockCol[br], c, offsetIn, offsetOut, remIn); + offsetIn += remIn; + } + else { + // next col + setOutputEntries(sliceVal, outputBlockCol[br], c, offsetIn, offsetOut, remOut); + offsetIn += remOut; + } + br++; + offsetOut = 0; + if(br == numRowBlocksOut) { + // next col + c++; + if(c == outputBlockCol[0].getNumColumns()) { + offset[0] = offsetIn == localRowsIn ? 0 : offsetIn; + break; + } + br = 0; + } + localRowsOut = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + } + localI++; + if(localI == numRowBlocksIn + 1) + localI = 1; + + lease.close(); + offsetIn = 0; + + if(k == leases.size() - 1 && offset[0] != 0) { + // put current slice back into table, to be able to reserve new budget + finalBudget.reserveBlocking(colBytes); + table.put(totalIdx, new ManagedPayload<>(slice, colBytes, finalBudget)); + } + } + + // enqueue filled output blocks and allocate new ones + for(int b = 0; b < outputBlockCol.length; b++) { + outputBlockCol[b].recomputeNonZeros(); + OOCUtils.enqueueExact(out, + new IndexedMatrixValue(new MatrixIndexes(b + 1, finalBc + 1), outputBlockCol[b]), + finalBudget, false); + } + }); + + budget.close(); + futures.clear(); + + if(bc == numColBlocksOut - 1) + break; + + // new block row + bc++; + budget = OOCUtils.reserveBudget(_allowance, outputBytes); + + if(offset[0] != 0) { + // get slice back from table + blkFuture = table.take(totalIdx, budget); + futures.add(blkFuture); + startI = i; + } + else { + startI = (i == numRowBlocksIn) ? 1 : i + 1; + } + missing = getNumMissingColSlices(startI, bc, offset[0], numRowBlocksIn, numColBlocksOut); + } + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + if(budget != null) { + budget.close(); + } + try { + table.close(); + onComplete(); + } + finally { + out.closeInput(); + } + } + } + + private int getNumMissingColSlices(int i, int bc, int offsetIn, int numRowBlocksIn, int numColBlocksOut) { + long localRowsIn = (i == numRowBlocksIn && _rlen % _blen != 0) ? _rlen % _blen : _blen; + long localColsOut = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? _cols % _blen : _blen; + + long totalIdx = _rows * localColsOut; + int cnt = 0; + + if(offsetIn != 0) { + // reuse table entry + totalIdx -= (localRowsIn - offsetIn); + cnt++; + } + + int numCols = (int) Math.floor((double) totalIdx / _rlen); + cnt += numCols * numRowBlocksIn; + totalIdx -= numCols * _rlen; + + long numBlenSlices = Math.max(0, numRowBlocksIn - (i - 1)); + long restCol = numBlenSlices * _blen + localRowsIn; + if(totalIdx - restCol > 0) { + totalIdx -= restCol; + cnt += numRowBlocksIn - i; + } + cnt += (int) Math.ceil((double) totalIdx / _blen); + + return cnt; + } + + private MatrixBlock[] allocateSliceBlocks(int idx, int numColBlocksOut, int numRowBlocksOut) { + int n = _byRow ? numColBlocksOut : numRowBlocksOut; + MatrixBlock[] res = new MatrixBlock[n]; + + // full inner blocks, adjust for outer blocks + int localRows = ((!_byRow || idx == numRowBlocksOut - 1) && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + int localCols = ((_byRow || idx == numColBlocksOut - 1) && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + + for(int k = 0; k < n - 1; k++) { + res[k] = _byRow ? new MatrixBlock(localRows, _blen, false) : new MatrixBlock(_blen, localCols, false); + res[k].allocateDenseBlock(); + } + res[n - 1] = new MatrixBlock(localRows, localCols, false); + res[n - 1].allocateDenseBlock(); + + return res; + } + + private void setOutputEntries(MatrixBlock src, MatrixBlock dest, int idx, int srcOffset, int destOffset, int length) { + if(_byRow) + ((DenseBlockFP64) dest.getDenseBlock()).setPartialRow(src.getDenseBlock(), idx, srcOffset, destOffset, length); + else + ((DenseBlockFP64) dest.getDenseBlock()).setPartialCol(src.getDenseBlock(), idx, srcOffset, destOffset, length); + } +} 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 e52fea34a6a..7a893c3c1aa 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 @@ -56,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.ReshapeOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.TSMMOOCPrimitive; import org.apache.sysds.runtime.ooc.primitives.TransposeOOCPrimitive; import org.apache.sysds.runtime.ooc.stats.OOCEventLog; @@ -172,6 +173,11 @@ public static void reduce(OOCStreamable input, OOCStream output, Fu output.assignPrimitive(new ReduceOOCPrimitive<>(input, output, partial, merge, size, context)); } + public static void reshape(OOCStreamable input, OOCStream output, + long rows, long cols, boolean byRow, StreamContext context) { + output.assignPrimitive(new ReshapeOOCPrimitive(input, output, rows, cols, byRow, 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 2b070b9365f..b8502404605 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 @@ -176,23 +176,28 @@ private static long estimateMatrixBlockBytes(long rows, long cols) { } public static void enqueueExact(OOCStream out, IndexedMatrixValue value, - ReservationBudget budget) { + ReservationBudget budget, boolean closeBudget) { long bytes = ((MatrixBlock) value.getValue()).getExactSerializedSize(); OOCStream.QueueCallback callback = null; try { budget.reserveBlocking(bytes); callback = new InMemoryQueueCallback<>(value, null, budget, bytes); - budget.close(); + if(closeBudget) budget.close(); out.enqueue(callback); callback = null; } finally { - budget.close(); + if(closeBudget) budget.close(); if(callback != null) callback.close(); } } + public static void enqueueExact(OOCStream out, IndexedMatrixValue value, + ReservationBudget budget) { + enqueueExact(out, value, budget, true); + } + 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/functions/ooc/ReshapeTest.java b/src/test/java/org/apache/sysds/test/functions/ooc/ReshapeTest.java index 770c5b7c5bf..e020afbc2ae 100644 --- a/src/test/java/org/apache/sysds/test/functions/ooc/ReshapeTest.java +++ b/src/test/java/org/apache/sysds/test/functions/ooc/ReshapeTest.java @@ -78,8 +78,9 @@ public static Iterable getParams() { int[][][] dims = { {{1000, 1000}, {1, 1000000}}, // single row/col - {{3000, 4000}, {1500, 8000}}, // partialBlocks - {{2400, 1400}, {800, 4200}} // fullBlocks + {{4000, 4000}, {2000, 8000}}, // full + {{814, 618}, {407, 1236}}, // partial + {{2814, 1618}, {1407, 3236}} }; ArrayList params = new ArrayList<>(); From 9b3140c6695dc3ae39a4c52fa4acadc2153179ae Mon Sep 17 00:00:00 2001 From: Jessica Priebe Date: Tue, 22 Sep 2026 02:19:03 +0200 Subject: [PATCH 2/2] address review comments - part 1 --- .../ooc/cache/packed/OOCPackedCache.java | 2 +- .../ooc/primitives/ReshapeOOCPrimitive.java | 367 +++++++++--------- 2 files changed, 174 insertions(+), 195 deletions(-) diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java index 13635b86b8e..a0ae3bc1aa3 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/packed/OOCPackedCache.java @@ -47,7 +47,7 @@ public final class OOCPackedCache implements OOCCache { private static final long PACKED_STREAM_ID = CachingStream._streamSeq.getNextID(); - private static final long DEFAULT_PACK_THRESHOLD_BYTES = 1; + private static final long DEFAULT_PACK_THRESHOLD_BYTES = 1; // disabled for now private static final long DEFAULT_PACK_TARGET_BYTES = 1L << 19; // 512 KB tile packing private static final long DEFAULT_MAX_STAGING_BYTES = 1L << 26; private static final int DEFAULT_MAX_OPEN_BUILDERS = 64; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java index 13c8a1b2756..f0e9233e148 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java @@ -50,10 +50,24 @@ public class ReshapeOOCPrimitive extends OOCPrimitive { private final boolean _byRow; private final long _rows; private final long _cols; + private long _rlen; private long _clen; private int _blen; + OOCStream _in; + OOCStream _out; + StateTable _table; + + private int _numColBlocksIn; + private int _numRowBlocksIn; + private int _numColBlocksOut; + private int _numRowBlocksOut; + private long _numRowsBlockOut; + private long _numColsBlockOut; + private long _blockBytesOut; + private long _sliceBytes; + public ReshapeOOCPrimitive(OOCStreamable input, OOCStreamable output, long rows, long cols, boolean byRow, StreamContext context) { super(context, input); @@ -80,101 +94,75 @@ protected void requestPatternInternal(OOCAccessPattern accessPattern) { @Override protected void startExecution() { - DataCharacteristics inputDc = _input.getDataCharacteristics(); - if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) - throw new DMLRuntimeException("Reshape OOC reduction requires known input dimensions and block size."); - - OOCStream input = getInputReadStream(0); - OOCStream output = _output.getWriteStream(); - getContext().addOutStream(output); - - _rlen = inputDc.getRows(); - _clen = inputDc.getCols(); - _blen = Math.toIntExact(inputDc.getBlocksize()); + initExecution(); if(_rlen * _clen != _rows * _cols) { + // non matching dims onComplete(); throw new DMLRuntimeException("Reshape matrix requires consistent numbers of input/output cells (" + _rlen + ":" + _clen + ", " + _rows + ":" + _cols + ")."); } if(_rlen == _rows) { + // same block dims OOCInstructionUtils - .submitAdmittedOOCTasks(input, output, + .submitAdmittedOOCTasks(_in, _out, value -> new IndexedMatrixValue(value.getIndexes(), value.getValue()), _allowance, getContext()) .thenRun(this::onComplete); return; } + initBlocking(); + if(_clen <= _blen && _rlen <= _blen && _cols <= _blen && _rows <= _blen) { - OOCInstructionUtils.submitAdmittedOOCTasks(input, output, + // single block + OOCInstructionUtils.submitAdmittedOOCTasks(_in, _out, value -> new IndexedMatrixValue(value.getIndexes(), ((MatrixBlock) value.getValue()).reshape((int) _rows, (int) _cols, _byRow)), _allowance, getContext()).thenRun(this::onComplete); return; } - int numColBlocksIn = Math.toIntExact(inputDc.getNumColBlocks()); - int numRowBlocksIn = Math.toIntExact(inputDc.getNumRowBlocks()); - int numColBlocksOut = (int) Math.ceil((double) _cols / _blen); - int numRowBlocksOut = (int) Math.ceil((double) _rows / _blen); - - long sliceBytes = (OOCUtils.estimateFullTileBytes(input.getDataCharacteristics()) + - (_blen - 1) * MatrixBlock.getHeaderSize()) / _blen; - - StateTable table = new StateTable<>(OOCCacheManager.getGlobalCache(), - CachingStream._streamSeq.getNextID()); - if(_byRow) { if(_clen % _blen == 0 && _cols % _blen == 0) { - // singleRowBlocks do not need to be split + // no need to split singleRowBlocks if(_rows == 1) { // result is one single row - submitSingleRowColTask(input, output, numColBlocksIn, numRowBlocksIn); + submitSingleRowColTask(); } else { - CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); - f.thenRun(() -> OOCInstructionUtils.submitOOCTask( - () -> reshapeFullColBlocks(table, output, numColBlocksOut, numRowBlocksOut, sliceBytes), - getContext())); + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapeFullColBlocks, getContext())); } } else { - CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); - f.thenRun(() -> OOCInstructionUtils.submitOOCTask(() -> reshapePartialColBlocks(table, output, - numColBlocksIn, numColBlocksOut, numRowBlocksOut, sliceBytes), getContext())); + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapePartialColBlocks, getContext())); } } else { if(_rlen % _blen == 0 && _rows % _blen == 0) { - // singleColBlocks do not need to be split + // no need to split singleColBlocks if(_cols == 1) { // result is one single col - submitSingleRowColTask(input, output, numColBlocksIn, numRowBlocksIn); + submitSingleRowColTask(); } else { - CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); - f.thenRun(() -> OOCInstructionUtils.submitOOCTask( - () -> reshapeFullRowBlocks(table, output, numColBlocksOut, numRowBlocksOut, sliceBytes), - getContext())); + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapeFullRowBlocks, getContext())); } } else { - CompletableFuture f = splitIntoTable(input, table, numColBlocksIn, numRowBlocksIn); - f.thenRun(() -> OOCInstructionUtils.submitOOCTask(() -> reshapePartialRowBlocks(table, output, - numRowBlocksIn, numColBlocksOut, numRowBlocksOut, sliceBytes), getContext())); + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapePartialRowBlocks, getContext())); } } } - private CompletableFuture splitIntoTable(OOCStream in, - StateTable table, int numColBlocks, int numRowBlocks) { + private CompletableFuture splitIntoTable() { - long blockBytes = OOCUtils.estimateFullTileBytes(in.getDataCharacteristics()) + - (_blen - 1) * MatrixBlock.getHeaderSize(); - long singleSliceBytes = blockBytes / _blen; - - AllocatedOOCStream allocated = new AllocatedOOCStream<>(in, _allowance, ignored -> blockBytes); + AllocatedOOCStream allocated = new AllocatedOOCStream<>(_in, _allowance, + ignored -> _blen * _sliceBytes); return OOCInstructionUtils.submitOOCTasks(allocated, callback -> { try(ReservationBudget budget = AllocatedOOCStream.detachBudget(callback)) { @@ -202,10 +190,10 @@ private CompletableFuture splitIntoTable(OOCStream in, rIdx = r; } - long targetIdx = _byRow ? (rIdx - 1) * numColBlocks + c - 1 : (cIdx - 1) * numRowBlocks + r - 1; + long targetIdx = _byRow ? (rIdx - 1) * _numColBlocksIn + c - 1 : (cIdx - 1) * _numRowBlocksIn + r - 1; IndexedMatrixValue sliceImv = new IndexedMatrixValue(new MatrixIndexes(rIdx, cIdx), slice); - budget.reserveBlocking(singleSliceBytes); - table.put((int) targetIdx, new ManagedPayload<>(sliceImv, singleSliceBytes, budget)); + budget.reserveBlocking(_sliceBytes); + _table.put((int) targetIdx, new ManagedPayload<>(sliceImv, _sliceBytes, budget)); } } catch(IllegalStateException e) { @@ -214,13 +202,11 @@ private CompletableFuture splitIntoTable(OOCStream in, }, getContext()); } - private void submitSingleRowColTask(OOCStream in, OOCStream out, - int numColBlocks, int numRowBlocks) { + private void submitSingleRowColTask() { // one input block is split into blen output blocks - long outputBytes = _blen * OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); - - AllocatedOOCStream allocated = new AllocatedOOCStream<>(in, _allowance, ignored -> outputBytes); + AllocatedOOCStream allocated = new AllocatedOOCStream<>(_in, _allowance, + ignored -> _blen * _blockBytesOut); OOCInstructionUtils.submitOOCTasks(allocated, callback -> { try(ReservationBudget budget = AllocatedOOCStream.detachBudget(callback)) { @@ -243,7 +229,7 @@ private void submitSingleRowColTask(OOCStream in, OOCStream< // total row, 1 based rIdx = (r - 1) * _blen + i + 1; // all in single row - cIdx = (rIdx - 1) * numColBlocks + c; + cIdx = (rIdx - 1) * _numColBlocksIn + c; rIdx = 1; } else { @@ -251,53 +237,49 @@ private void submitSingleRowColTask(OOCStream in, OOCStream< slice = blk.slice(0, blk.getNumRows() - 1, i, i); cIdx = (c - 1) * _blen + i + 1; // all in single col - rIdx = (cIdx - 1) * numRowBlocks + r; + rIdx = (cIdx - 1) * _numRowBlocksIn + r; cIdx = 1; } IndexedMatrixValue sliceImv = new IndexedMatrixValue(new MatrixIndexes(rIdx, cIdx), slice); - OOCUtils.enqueueExact(out, sliceImv, budget, false); + OOCUtils.enqueueExact(_out, sliceImv, budget, false); } } catch(IllegalStateException e) { throw new DMLRuntimeException(e); } - }, getContext()).thenRun(this::onComplete).thenRun(out::closeInput).exceptionally(error -> { - out.propagateFailure(DMLRuntimeException.of(error)); + }, getContext()).thenRun(this::onComplete).thenRun(_out::closeInput).exceptionally(error -> { + _out.propagateFailure(DMLRuntimeException.of(error)); return null; }); } - private void reshapeFullColBlocks(StateTable table, OOCStream out, - int numColBlocksOut, int numRowBlocksOut, long rowBytes) { + private void reshapeFullColBlocks() { List>> futures = new ArrayList<>(); + long outputBytes = _numRowsBlockOut * _sliceBytes + _blockBytesOut; ReservationBudget budget = null; - long numRowsBlockOut = Math.min(out.getDataCharacteristics().getRows(), _blen); - long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); - long outputBytes = numRowsBlockOut * rowBytes + blockBytes; - // totalRowIdx corresponds to index of row block when all aligned in one row // br * numColBlocksOut * blen + b + r * numColBlocksOut; // with numColBlocksOut * blen = cols - long totalIdx = -_cols - 1 - numColBlocksOut; + long totalIdx = -_cols - 1 - _numColBlocksOut; try { // iterate through rows of output blocks - for(int br = 0; br < numRowBlocksOut; br++) { + for(int br = 0; br < _numRowBlocksOut; br++) { totalIdx += _cols; long tmp = totalIdx; - int localRows = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + int localRows = (br == _numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; // for each block in row - for(int b = 0; b < numColBlocksOut; b++) { + for(int b = 0; b < _numColBlocksOut; b++) { totalIdx += 1; long tmp2 = totalIdx; budget = OOCUtils.reserveBudget(_allowance, outputBytes); // for each row in block for(int r = 0; r < _blen && r < localRows; r++) { - totalIdx += numColBlocksOut; - OOCFuture> rowFuture = table.take((int) totalIdx, budget); + totalIdx += _numColBlocksOut; + OOCFuture> rowFuture = _table.take((int) totalIdx, budget); futures.add(rowFuture); } totalIdx = tmp2; @@ -314,7 +296,7 @@ private void reshapeFullColBlocks(StateTable table, OOCStrea lease.close(); } block.recomputeNonZeros(); - OOCUtils.enqueueExact(out, new IndexedMatrixValue(idx, block), finalBudget, true); + OOCUtils.enqueueExact(_out, new IndexedMatrixValue(idx, block), finalBudget, true); futures.clear(); }); budget = null; @@ -326,49 +308,36 @@ private void reshapeFullColBlocks(StateTable table, OOCStrea throw new DMLRuntimeException(e); } finally { - if(budget != null) { - budget.close(); - } - try { - table.close(); - onComplete(); - } - finally { - out.closeInput(); - } + closeResourcesAndComplete(budget); } } - private void reshapeFullRowBlocks(StateTable table, OOCStream out, - int numColBlocksOut, int numRowBlocksOut, long colBytes) { + private void reshapeFullRowBlocks() { List>> futures = new ArrayList<>(); + long outputBytes = _numColsBlockOut * _sliceBytes + _blockBytesOut; ReservationBudget budget = null; - long numColsBlockOut = Math.min(out.getDataCharacteristics().getCols(), _blen); - long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); - long outputBytes = numColsBlockOut * colBytes + blockBytes; - // totalColIdx corresponds to index of col block when all aligned in one col // bc * numRowBlocksOut * blen + b + c * numRowBlocksOut; // with numRowBlocksOut * blen = rows - long totalIdx = -_rows - 1 - numRowBlocksOut; + long totalIdx = -_rows - 1 - _numRowBlocksOut; try { // iterate through cols of output blocks - for(int bc = 0; bc < numColBlocksOut; bc++) { + for(int bc = 0; bc < _numColBlocksOut; bc++) { totalIdx += _rows; long tmp = totalIdx; - int localCols = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + int localCols = (bc == _numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; // for each block in col - for(int b = 0; b < numRowBlocksOut; b++) { + for(int b = 0; b < _numRowBlocksOut; b++) { totalIdx += 1; long tmp2 = totalIdx; budget = OOCUtils.reserveBudget(_allowance, outputBytes); // for each col in block for(int c = 0; c < _blen && c < localCols; c++) { - totalIdx += numRowBlocksOut; - OOCFuture> colFuture = table.take((int) totalIdx, budget); + totalIdx += _numRowBlocksOut; + OOCFuture> colFuture = _table.take((int) totalIdx, budget); futures.add(colFuture); } totalIdx = tmp2; @@ -386,7 +355,7 @@ private void reshapeFullRowBlocks(StateTable table, OOCStrea lease.close(); } block.recomputeNonZeros(); - OOCUtils.enqueueExact(out, new IndexedMatrixValue(idx, block), finalBudget, true); + OOCUtils.enqueueExact(_out, new IndexedMatrixValue(idx, block), finalBudget, true); futures.clear(); }); budget = null; @@ -398,29 +367,15 @@ private void reshapeFullRowBlocks(StateTable table, OOCStrea throw new DMLRuntimeException(e); } finally { - if(budget != null) { - budget.close(); - } - try { - table.close(); - onComplete(); - } - finally { - out.closeInput(); - } + closeResourcesAndComplete(budget); } } - private void reshapePartialColBlocks(StateTable table, OOCStream out, - int numColBlocksIn, int numColBlocksOut, int numRowBlocksOut, long rowBytes) { + private void reshapePartialColBlocks() { - long numRowsBlockOut = Math.min(_rows, _blen); - long numColsBlockOut = Math.min(_cols, _blen); - - long numNeededRowsIn = 2 + - (long) Math.ceil((((double) numRowsBlockOut * numColsBlockOut) / _clen) * numColBlocksIn * numColBlocksOut); - long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); - long outputBytes = numNeededRowsIn * rowBytes + numColBlocksOut * blockBytes; + long numNeededRowsIn = 2 + (long) Math + .ceil((((double) _numRowsBlockOut * _numColsBlockOut) / _clen) * _numColBlocksIn * _numColBlocksOut); + long outputBytes = numNeededRowsIn * _sliceBytes + _numColBlocksOut * _blockBytesOut; ReservationBudget budget = null; int br = 0; @@ -429,15 +384,15 @@ private void reshapePartialColBlocks(StateTable table, OOCSt List>> futures = new ArrayList<>(); // new row of output blocks budget = OOCUtils.reserveBudget(_allowance, outputBytes); - int missing = getNumMissingRowSlices(1, br, 0, numColBlocksIn, numRowBlocksOut); + int missing = getNumMissingRowSlices(1, br, 0); int startJ = 1; final int[] offset = {0}; // iterate through input rows and add to row of output blocks for(int i = 1; i <= _rlen; i++) { - for(int j = 1; j <= numColBlocksIn; j++) { - int totalIdx = (i - 1) * numColBlocksIn + j - 1; - OOCFuture> blkFuture = table.take(totalIdx, budget); + for(int j = 1; j <= _numColBlocksIn; j++) { + int totalIdx = (i - 1) * _numColBlocksIn + j - 1; + OOCFuture> blkFuture = _table.take(totalIdx, budget); futures.add(blkFuture); if(futures.size() < missing) @@ -455,7 +410,7 @@ private void reshapePartialColBlocks(StateTable table, OOCSt int bc = 0; int r = 0; int offsetOut = 0; - MatrixBlock[] outputBlockRow = allocateSliceBlocks(finalBr, numColBlocksOut, numRowBlocksOut); + MatrixBlock[] outputBlockRow = allocateSliceBlocks(finalBr); int localColsOut = (_cols > _blen) ? _blen : (int) _cols; for(int k = 0; k < leases.size(); k++) { @@ -463,7 +418,7 @@ private void reshapePartialColBlocks(StateTable table, OOCSt IndexedMatrixValue slice = lease.value(); MatrixBlock sliceVal = (MatrixBlock) slice.getValue(); - int localColsIn = (localJ == numColBlocksIn && _clen % _blen != 0) ? (int) _clen % _blen : _blen; + int localColsIn = (localJ == _numColBlocksIn && _clen % _blen != 0) ? (int) _clen % _blen : _blen; while(offsetIn < localColsIn) { // until input row fully processed int remIn = localColsIn - offsetIn; @@ -487,7 +442,7 @@ else if(remIn == remOut) { } bc++; offsetOut = 0; - if(bc == numColBlocksOut) { + if(bc == _numColBlocksOut) { // next row r++; if(r == outputBlockRow[0].getNumRows()) { @@ -496,10 +451,10 @@ else if(remIn == remOut) { } bc = 0; } - localColsOut = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + localColsOut = (bc == _numColBlocksOut - 1 && _cols % _blen != 0) ? (int) _cols % _blen : _blen; } localJ++; - if(localJ == numColBlocksIn + 1) + if(localJ == _numColBlocksIn + 1) localJ = 1; lease.close(); @@ -507,15 +462,15 @@ else if(remIn == remOut) { if(k == leases.size() - 1 && offset[0] != 0) { // put current slice back into table, to be able to reserve new budget - finalBudget.reserveBlocking(rowBytes); - table.put(totalIdx, new ManagedPayload<>(slice, rowBytes, finalBudget)); + finalBudget.reserveBlocking(_sliceBytes); + _table.put(totalIdx, new ManagedPayload<>(slice, _sliceBytes, finalBudget)); } } // enqueue filled output blocks and allocate new ones for(int b = 0; b < outputBlockRow.length; b++) { outputBlockRow[b].recomputeNonZeros(); - OOCUtils.enqueueExact(out, + OOCUtils.enqueueExact(_out, new IndexedMatrixValue(new MatrixIndexes(finalBr + 1, b + 1), outputBlockRow[b]), finalBudget, false); } @@ -524,7 +479,7 @@ else if(remIn == remOut) { budget.close(); futures.clear(); - if(br == numRowBlocksOut - 1) + if(br == _numRowBlocksOut - 1) break; // new block row @@ -533,15 +488,15 @@ else if(remIn == remOut) { if(offset[0] != 0) { // get slice back from table - blkFuture = table.take(totalIdx, budget); + blkFuture = _table.take(totalIdx, budget); futures.add(blkFuture); startJ = j; } else { - startJ = (j == numColBlocksIn) ? 1 : j + 1; + startJ = (j == _numColBlocksIn) ? 1 : j + 1; } - missing = getNumMissingRowSlices(startJ, br, offset[0], numColBlocksIn, numRowBlocksOut); + missing = getNumMissingRowSlices(startJ, br, offset[0]); } } } @@ -549,22 +504,13 @@ else if(remIn == remOut) { throw new DMLRuntimeException(e); } finally { - if(budget != null) { - budget.close(); - } - try { - table.close(); - onComplete(); - } - finally { - out.closeInput(); - } + closeResourcesAndComplete(budget); } } - private int getNumMissingRowSlices(int j, int br, int offsetIn, int numColBlocksIn, int numRowBlocksOut) { - long localColsIn = (j == numColBlocksIn && _clen % _blen != 0) ? _clen % _blen : _blen; - long localRowsOut = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? _rows % _blen : _blen; + private int getNumMissingRowSlices(int j, int br, int offsetIn) { + long localColsIn = (j == _numColBlocksIn && _clen % _blen != 0) ? _clen % _blen : _blen; + long localRowsOut = (br == _numRowBlocksOut - 1 && _rows % _blen != 0) ? _rows % _blen : _blen; long totalIdx = _cols * localRowsOut; int cnt = 0; @@ -576,30 +522,25 @@ private int getNumMissingRowSlices(int j, int br, int offsetIn, int numColBlocks } int numRows = (int) Math.floor((double) totalIdx / _clen); - cnt += numRows * numColBlocksIn; + cnt += numRows * _numColBlocksIn; totalIdx -= numRows * _clen; - long numBlenSlices = Math.max(0, numColBlocksIn - (j - 1)); + long numBlenSlices = Math.max(0, _numColBlocksIn - (j - 1)); long restRow = numBlenSlices * _blen + localColsIn; if(totalIdx - restRow > 0) { totalIdx -= restRow; - cnt += numColBlocksIn - j; + cnt += _numColBlocksIn - j; } cnt += (int) Math.ceil((double) totalIdx / _blen); return cnt; } - private void reshapePartialRowBlocks(StateTable table, OOCStream out, - int numRowBlocksIn, int numColBlocksOut, int numRowBlocksOut, long colBytes) { - - long numRowsBlockOut = Math.min(_rows, _blen); - long numColsBlockOut = Math.min(_cols, _blen); + private void reshapePartialRowBlocks() { - long numNeededColsIn = 2 + - (long) Math.ceil((((double) numRowsBlockOut * numColsBlockOut) / _rlen) * numRowBlocksIn * numRowBlocksOut); - long blockBytes = OOCUtils.estimateOutputTileBytes(out.getDataCharacteristics()); - long outputBytes = numNeededColsIn * colBytes + numRowBlocksOut * blockBytes; + long numNeededColsIn = 2 + (long) Math + .ceil((((double) _numRowsBlockOut * _numColsBlockOut) / _rlen) * _numRowBlocksIn * _numRowBlocksOut); + long outputBytes = numNeededColsIn * _sliceBytes + _numRowBlocksOut * _blockBytesOut; ReservationBudget budget = null; int bc = 0; @@ -608,15 +549,15 @@ private void reshapePartialRowBlocks(StateTable table, OOCSt List>> futures = new ArrayList<>(); // new col of output blocks budget = OOCUtils.reserveBudget(_allowance, outputBytes); - int missing = getNumMissingColSlices(1, bc, 0, numRowBlocksIn, numColBlocksOut); + int missing = getNumMissingColSlices(1, bc, 0); int startI = 1; final int[] offset = {0}; // iterate through input cols and add to col of output blocks for(int j = 1; j <= _clen; j++) { - for(int i = 1; i <= numRowBlocksIn; i++) { - int totalIdx = (j - 1) * numRowBlocksIn + i - 1; - OOCFuture> blkFuture = table.take(totalIdx, budget); + for(int i = 1; i <= _numRowBlocksIn; i++) { + int totalIdx = (j - 1) * _numRowBlocksIn + i - 1; + OOCFuture> blkFuture = _table.take(totalIdx, budget); futures.add(blkFuture); if(futures.size() < missing) @@ -635,7 +576,7 @@ private void reshapePartialRowBlocks(StateTable table, OOCSt int c = 0; int offsetOut = 0; - MatrixBlock[] outputBlockCol = allocateSliceBlocks(finalBc, numColBlocksOut, numRowBlocksOut); + MatrixBlock[] outputBlockCol = allocateSliceBlocks(finalBc); int localRowsOut = (_rows > _blen) ? _blen : (int) _rows; for(int k = 0; k < leases.size(); k++) { @@ -643,7 +584,7 @@ private void reshapePartialRowBlocks(StateTable table, OOCSt IndexedMatrixValue slice = lease.value(); MatrixBlock sliceVal = (MatrixBlock) slice.getValue(); - int localRowsIn = (localI == numRowBlocksIn && _rlen % _blen != 0) ? (int) _rlen % _blen : _blen; + int localRowsIn = (localI == _numRowBlocksIn && _rlen % _blen != 0) ? (int) _rlen % _blen : _blen; while(offsetIn < localRowsIn) { // until input col fully processed int remIn = localRowsIn - offsetIn; @@ -667,7 +608,7 @@ else if(remIn == remOut) { } br++; offsetOut = 0; - if(br == numRowBlocksOut) { + if(br == _numRowBlocksOut) { // next col c++; if(c == outputBlockCol[0].getNumColumns()) { @@ -676,10 +617,10 @@ else if(remIn == remOut) { } br = 0; } - localRowsOut = (br == numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + localRowsOut = (br == _numRowBlocksOut - 1 && _rows % _blen != 0) ? (int) _rows % _blen : _blen; } localI++; - if(localI == numRowBlocksIn + 1) + if(localI == _numRowBlocksIn + 1) localI = 1; lease.close(); @@ -687,15 +628,15 @@ else if(remIn == remOut) { if(k == leases.size() - 1 && offset[0] != 0) { // put current slice back into table, to be able to reserve new budget - finalBudget.reserveBlocking(colBytes); - table.put(totalIdx, new ManagedPayload<>(slice, colBytes, finalBudget)); + finalBudget.reserveBlocking(_sliceBytes); + _table.put(totalIdx, new ManagedPayload<>(slice, _sliceBytes, finalBudget)); } } // enqueue filled output blocks and allocate new ones for(int b = 0; b < outputBlockCol.length; b++) { outputBlockCol[b].recomputeNonZeros(); - OOCUtils.enqueueExact(out, + OOCUtils.enqueueExact(_out, new IndexedMatrixValue(new MatrixIndexes(b + 1, finalBc + 1), outputBlockCol[b]), finalBudget, false); } @@ -704,7 +645,7 @@ else if(remIn == remOut) { budget.close(); futures.clear(); - if(bc == numColBlocksOut - 1) + if(bc == _numColBlocksOut - 1) break; // new block row @@ -713,14 +654,14 @@ else if(remIn == remOut) { if(offset[0] != 0) { // get slice back from table - blkFuture = table.take(totalIdx, budget); + blkFuture = _table.take(totalIdx, budget); futures.add(blkFuture); startI = i; } else { - startI = (i == numRowBlocksIn) ? 1 : i + 1; + startI = (i == _numRowBlocksIn) ? 1 : i + 1; } - missing = getNumMissingColSlices(startI, bc, offset[0], numRowBlocksIn, numColBlocksOut); + missing = getNumMissingColSlices(startI, bc, offset[0]); } } } @@ -728,22 +669,13 @@ else if(remIn == remOut) { throw new DMLRuntimeException(e); } finally { - if(budget != null) { - budget.close(); - } - try { - table.close(); - onComplete(); - } - finally { - out.closeInput(); - } + closeResourcesAndComplete(budget); } } - private int getNumMissingColSlices(int i, int bc, int offsetIn, int numRowBlocksIn, int numColBlocksOut) { - long localRowsIn = (i == numRowBlocksIn && _rlen % _blen != 0) ? _rlen % _blen : _blen; - long localColsOut = (bc == numColBlocksOut - 1 && _cols % _blen != 0) ? _cols % _blen : _blen; + private int getNumMissingColSlices(int i, int bc, int offsetIn) { + long localRowsIn = (i == _numRowBlocksIn && _rlen % _blen != 0) ? _rlen % _blen : _blen; + long localColsOut = (bc == _numColBlocksOut - 1 && _cols % _blen != 0) ? _cols % _blen : _blen; long totalIdx = _rows * localColsOut; int cnt = 0; @@ -755,27 +687,27 @@ private int getNumMissingColSlices(int i, int bc, int offsetIn, int numRowBlocks } int numCols = (int) Math.floor((double) totalIdx / _rlen); - cnt += numCols * numRowBlocksIn; + cnt += numCols * _numRowBlocksIn; totalIdx -= numCols * _rlen; - long numBlenSlices = Math.max(0, numRowBlocksIn - (i - 1)); + long numBlenSlices = Math.max(0, _numRowBlocksIn - (i - 1)); long restCol = numBlenSlices * _blen + localRowsIn; if(totalIdx - restCol > 0) { totalIdx -= restCol; - cnt += numRowBlocksIn - i; + cnt += _numRowBlocksIn - i; } cnt += (int) Math.ceil((double) totalIdx / _blen); return cnt; } - private MatrixBlock[] allocateSliceBlocks(int idx, int numColBlocksOut, int numRowBlocksOut) { - int n = _byRow ? numColBlocksOut : numRowBlocksOut; + private MatrixBlock[] allocateSliceBlocks(int idx) { + int n = _byRow ? _numColBlocksOut : _numRowBlocksOut; MatrixBlock[] res = new MatrixBlock[n]; // full inner blocks, adjust for outer blocks - int localRows = ((!_byRow || idx == numRowBlocksOut - 1) && _rows % _blen != 0) ? (int) _rows % _blen : _blen; - int localCols = ((_byRow || idx == numColBlocksOut - 1) && _cols % _blen != 0) ? (int) _cols % _blen : _blen; + int localRows = ((!_byRow || idx == _numRowBlocksOut - 1) && _rows % _blen != 0) ? (int) _rows % _blen : _blen; + int localCols = ((_byRow || idx == _numColBlocksOut - 1) && _cols % _blen != 0) ? (int) _cols % _blen : _blen; for(int k = 0; k < n - 1; k++) { res[k] = _byRow ? new MatrixBlock(localRows, _blen, false) : new MatrixBlock(_blen, localCols, false); @@ -793,4 +725,51 @@ private void setOutputEntries(MatrixBlock src, MatrixBlock dest, int idx, int sr else ((DenseBlockFP64) dest.getDenseBlock()).setPartialCol(src.getDenseBlock(), idx, srcOffset, destOffset, length); } + + private void closeResourcesAndComplete(ReservationBudget budget) { + if(budget != null) { + budget.close(); + } + try { + _table.close(); + onComplete(); + } + finally { + _out.closeInput(); + } + } + + private void initExecution() { + DataCharacteristics inputDc = _input.getDataCharacteristics(); + if(inputDc == null || !inputDc.dimsKnown() || inputDc.getBlocksize() <= 0) + throw new DMLRuntimeException( + "Reshape OOC reduction requires known input dimensions and block size."); + + _in = getInputReadStream(0); + _out = _output.getWriteStream(); + getContext().addOutStream(_out); + + _rlen = inputDc.getRows(); + _clen = inputDc.getCols(); + _blen = Math.toIntExact(inputDc.getBlocksize()); + } + + private void initBlocking() { + DataCharacteristics inputDc = _input.getDataCharacteristics(); + + _numColBlocksIn = Math.toIntExact(inputDc.getNumColBlocks()); + _numRowBlocksIn = Math.toIntExact(inputDc.getNumRowBlocks()); + + _numColBlocksOut = (int) Math.ceil((double) _cols / _blen); + _numRowBlocksOut = (int) Math.ceil((double) _rows / _blen); + + _numRowsBlockOut = Math.min(_rows, _blen); + _numColsBlockOut = Math.min(_cols, _blen); + + _blockBytesOut = OOCUtils.estimateOutputTileBytes(_out.getDataCharacteristics()); + _sliceBytes = (OOCUtils.estimateFullTileBytes(_in.getDataCharacteristics()) + + (_blen - 1) * MatrixBlock.getHeaderSize()) / _blen; + + _table = new StateTable<>(OOCCacheManager.getGlobalCache(), CachingStream._streamSeq.getNextID()); + } }