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..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 = 1L << 18; + 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 new file mode 100644 index 00000000000..3ac363c1871 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/primitives/ReshapeOOCPrimitive.java @@ -0,0 +1,782 @@ +/* + * 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; + + 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); + _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() { + 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(_in, _out, + value -> new IndexedMatrixValue(value.getIndexes(), value.getValue()), _allowance, getContext()) + .thenRun(this::onComplete); + return; + } + + if(_clen <= _blen && _rlen <= _blen && _cols <= _blen && _rows <= _blen) { + // 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; + } + + initBlocking(); + + if(_byRow) { + if(_clen % _blen == 0 && _cols % _blen == 0) { + // no need to split singleRowBlocks + if(_rows == 1) { + // result is one single row + submitSingleRowColTask(); + } + else { + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapeFullColBlocks, getContext())); + } + } + else { + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapePartialColBlocks, getContext())); + } + } + else { + if(_rlen % _blen == 0 && _rows % _blen == 0) { + // no need to split singleColBlocks + if(_cols == 1) { + // result is one single col + submitSingleRowColTask(); + } + else { + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapeFullRowBlocks, getContext())); + } + } + else { + CompletableFuture f = splitIntoTable(); + f.thenRun(() -> OOCInstructionUtils.submitOOCTask(this::reshapePartialRowBlocks, getContext())); + } + } + } + + private CompletableFuture splitIntoTable() { + AllocatedOOCStream allocated = new AllocatedOOCStream<>(_in, _allowance, + ignored -> _blen * _sliceBytes); + return OOCInstructionUtils.submitOOCTasks(allocated, this::splitBlockIntoTable, getContext()); + } + + private void splitBlockIntoTable(OOCStream.QueueCallback 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(); + + int n = _byRow ? blk.getNumRows() : blk.getNumColumns(); + for(int i = 0; i < n; i++) + putSliceIntoTable(blk, r, c, i, budget); + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + } + + private void putSliceIntoTable(MatrixBlock blk, long r, long c, int i, ReservationBudget budget) { + MatrixBlock slice = createSlice(blk, i); + long rIdx = _byRow ? (r - 1) * _blen + i + 1 : r; + long cIdx = _byRow ? c : (c - 1) * _blen + i + 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(_sliceBytes); + _table.put((int) targetIdx, new ManagedPayload<>(sliceImv, _sliceBytes, budget)); + } + + private void submitSingleRowColTask() { + // one input block is split into blen output blocks + AllocatedOOCStream allocated = new AllocatedOOCStream<>(_in, _allowance, + ignored -> _blen * _blockBytesOut); + + OOCInstructionUtils.submitOOCTasks(allocated, this::processSingleRowColBlock, getContext()) + .thenRun(this::onComplete).thenRun(_out::closeInput).exceptionally(error -> { + _out.propagateFailure(DMLRuntimeException.of(error)); + return null; + }); + } + + private void processSingleRowColBlock(OOCStream.QueueCallback 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(); + + int n = _byRow ? blk.getNumRows() : blk.getNumColumns(); + for(int i = 0; i < n; i++) + enqueueSlice(blk, r, c, i, budget); + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + } + + private void enqueueSlice(MatrixBlock blk, long r, long c, int i, ReservationBudget budget) { + MatrixBlock slice = createSlice(blk, i); + long rIdx = _byRow ? 1 : ((c - 1) * _blen + i) * _numRowBlocksIn + r; + long cIdx = _byRow ? ((r - 1) * _blen + i) * _numColBlocksIn + c : 1; + IndexedMatrixValue sliceImv = new IndexedMatrixValue(new MatrixIndexes(rIdx, cIdx), slice); + OOCUtils.enqueueExact(_out, sliceImv, budget, false); + } + + private MatrixBlock createSlice(MatrixBlock block, int i) { + return _byRow ? block.slice(i, i) : block.slice(0, block.getNumRows() - 1, i, i); + } + + private void reshapeFullColBlocks() { + + List>> futures = new ArrayList<>(); + long outputBytes = _numRowsBlockOut * _sliceBytes + _blockBytesOut; + ReservationBudget budget = null; + + // 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) -> processFullColBlock(idx, localRows, leases, finalBudget)); + futures.clear(); + budget = null; + } + totalIdx = tmp; + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + closeResourcesAndComplete(budget); + } + } + + private void processFullColBlock(MatrixIndexes idx, int localRows, List> leases, + ReservationBudget budget) { + + 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), budget, true); + } + + private void reshapeFullRowBlocks() { + + List>> futures = new ArrayList<>(); + long outputBytes = _numColsBlockOut * _sliceBytes + _blockBytesOut; + ReservationBudget budget = null; + + // 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) -> processFullRowBlock(idx, localCols, leases, finalBudget)); + futures.clear(); + budget = null; + } + totalIdx = tmp; + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + closeResourcesAndComplete(budget); + } + } + + private void processFullRowBlock(MatrixIndexes idx, int localCols, List> leases, + ReservationBudget budget) { + + 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), budget, true); + } + + private void reshapePartialColBlocks() { + long numNeededRowsIn = 2 + (long) Math + .ceil((((double) _numRowsBlockOut * _numColsBlockOut) / _clen) * _numColBlocksIn * _numColBlocksOut); + long outputBytes = numNeededRowsIn * _sliceBytes + _numColBlocksOut * _blockBytesOut; + + 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); + 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) -> processPartialColBlocks(finalBr, finalJ, offset, totalIdx, + leases, finalBudget)); + + futures.clear(); + budget = null; + + 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]); + } + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + closeResourcesAndComplete(budget); + } + } + + private void processPartialColBlocks(int br, int j, int[] offset, int totalIdx, + List> leases, ReservationBudget budget) { + + int offsetIn = offset[0]; + int offsetOut = 0; + int bc = 0; + int r = 0; + + MatrixBlock[] outputBlockRow = allocateSliceBlocks(br); + 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 = (j == _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; + } + j++; + if(j == _numColBlocksIn + 1) + j = 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 + budget.reserveBlocking(_sliceBytes); + _table.put(totalIdx, new ManagedPayload<>(slice, _sliceBytes, budget)); + } + } + + // 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(br + 1, b + 1), outputBlockRow[b]), + budget, false); + } + + budget.close(); + } + + 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; + + 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() { + long numNeededColsIn = 2 + (long) Math + .ceil((((double) _numRowsBlockOut * _numColsBlockOut) / _rlen) * _numRowBlocksIn * _numRowBlocksOut); + long outputBytes = numNeededColsIn * _sliceBytes + _numRowBlocksOut * _blockBytesOut; + + 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); + 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) -> processPartialRowBlocks(finalBc, finalI, offset, totalIdx, + leases, finalBudget)); + + futures.clear(); + budget = null; + + 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]); + } + } + } + catch(IllegalStateException e) { + throw new DMLRuntimeException(e); + } + finally { + closeResourcesAndComplete(budget); + } + } + + private void processPartialRowBlocks(int bc, int i, int[] offset, int totalIdx, + List> leases, ReservationBudget budget) { + + int offsetIn = offset[0]; + int offsetOut = 0; + int br = 0; + int c = 0; + + MatrixBlock[] outputBlockCol = allocateSliceBlocks(bc); + 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 = (i == _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; + } + i++; + if(i == _numRowBlocksIn + 1) + i = 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 + budget.reserveBlocking(_sliceBytes); + _table.put(totalIdx, new ManagedPayload<>(slice, _sliceBytes, budget)); + } + } + + // 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, bc + 1), outputBlockCol[b]), + budget, false); + } + + budget.close(); + } + + 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; + + 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 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); + } + + 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()); + } +} 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..d8122456d26 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 + {{1400, 1200}, {1200, 1400}}, // partial + {{814, 618}, {407, 1236}} }; ArrayList params = new ArrayList<>();