tuluyhansozen commented on code in PR #2529:
URL: https://github.com/apache/systemds/pull/2529#discussion_r3537692764


##########
src/main/java/org/apache/sysds/runtime/instructions/ooc/TSMMOOCInstruction.java:
##########
@@ -59,39 +64,90 @@ public static TSMMOOCInstruction parseInstruction(String 
str) {
        }
 
        @Override
-       public void processInstruction( ExecutionContext ec ) {
+       public void processInstruction(ExecutionContext ec) {
                MatrixObject min = ec.getMatrixObject(input1);
-               int nRows = (int) min.getDataCharacteristics().getRows();
-               int nCols = (int) min.getDataCharacteristics().getCols();
-               int bLen = min.getDataCharacteristics().getBlocksize();
-               
-               OOCStream<IndexedMatrixValue> qIn = min.getStreamHandle();
+               int numRowBlocks = 
Math.toIntExact(min.getDataCharacteristics().getNumRowBlocks());
+               int numColBlocks = 
Math.toIntExact(min.getDataCharacteristics().getNumColBlocks());
+               int blocksPerJoinGroup = _type.isLeft() ? numColBlocks : 
numRowBlocks;
+               int partialsPerOutput = _type.isLeft() ? numRowBlocks : 
numColBlocks;
+
+               OOCStreamable<IndexedMatrixValue> inputStreamable = 
min.getStreamable();
+               final boolean createdCache = !inputStreamable.hasStreamCache();
+               final CachingStream inputCache = createdCache ? new 
CachingStream(min.getStreamHandle())
+                       : inputStreamable.getStreamCache();
+
+               OOCStream<List<IndexedMatrixValue>> groupedPartials = 
createWritableStream();
+               OOCStream<IndexedMatrixValue> partials = createWritableStream();
+               OOCStream<IndexedMatrixValue> out = createWritableStream();
+               addOutStream(out);
+               ec.getMatrixObject(output).setStreamHandle(out);
+
+               CompletableFuture<Void> joinFuture = 
joinManyOOC(inputCache.getReadStream(), inputCache.getReadStream(), 
groupedPartials,
+                       this::createPartialOutputTiles, this::getJoinIndex, 
this::getJoinIndex,
+                       blocksPerJoinGroup, blocksPerJoinGroup);
+               CompletableFuture<Void> expandFuture = 
expandOOC(groupedPartials, partials, values -> values);
+
                BinaryOperator plus = 
InstructionUtils.parseBinaryOperator(Opcodes.PLUS.toString());
+               CompletableFuture<Void> outFuture = groupedReduceOOC(partials, 
out, (left, right) -> {
+                       MatrixBlock result = ((MatrixBlock) 
left.getValue()).binaryOperations(plus, right.getValue());
+                       left.setValue(result);
+                       return left;
+               }, partialsPerOutput);
+
+               propagateFailuresToOutput(out, List.of(joinFuture, 
expandFuture, outFuture));
+
+               outFuture.whenComplete((result, error) -> {
+                       if(createdCache)
+                               inputCache.scheduleDeletion();
+               });
+       }
+
+       private long getJoinIndex(IndexedMatrixValue value) {
+               return _type.isLeft() ? value.getIndexes().getRowIndex() : 
value.getIndexes().getColumnIndex();
+       }
 
-               //validation check TODO extend compiler to not create OOC 
otherwise
-               if(    (_type.isLeft() && nCols > bLen)
-                       || (_type.isRight() && nRows > bLen) )
-               {
-                       throw new UnsupportedOperationException();
+       private long getOutputIndex(IndexedMatrixValue value) {
+               return _type.isLeft() ? value.getIndexes().getColumnIndex() : 
value.getIndexes().getRowIndex();
+       }
+
+       private List<IndexedMatrixValue> 
createPartialOutputTiles(IndexedMatrixValue left, IndexedMatrixValue right) {
+               long leftIndex = getOutputIndex(left);
+               long rightIndex = getOutputIndex(right);
+               if(leftIndex > rightIndex)
+                       return List.of();
+
+               MatrixBlock leftBlock = (MatrixBlock) left.getValue();
+               MatrixBlock rightBlock = (MatrixBlock) right.getValue();
+               if(leftIndex == rightIndex) {
+                       MatrixBlock diagonal = 
leftBlock.transposeSelfMatrixMultOperations(new MatrixBlock(), _type);
+                       return List.of(new IndexedMatrixValue(new 
MatrixIndexes(leftIndex, rightIndex), diagonal));
                }
-               
-               //int dim = _type.isLeft() ? nCols : nRows;
-               MatrixBlock resultBlock = null;
-
-               OOCStream<MatrixBlock> tmpStream = createWritableStream();
-
-               mapOOC(qIn, tmpStream,
-                       tmp -> ((MatrixBlock) tmp.getValue())
-                               .transposeSelfMatrixMultOperations(new 
MatrixBlock(), _type));
-
-               MatrixBlock tmp;
-               while ((tmp = tmpStream.dequeue()) != 
LocalTaskQueue.NO_MORE_TASKS) {
-                       if (resultBlock == null)
-                               resultBlock = tmp;
-                       else
-                               resultBlock.binaryOperationsInPlace(plus, tmp);

Review Comment:
   Thank you for your comment. I updated and commited TSMMOOCInstruction to 
keep the old lightweight path for single-output-tile TSMM.
   For LEFT TSMM, it now uses the fast path when `numColBlocks == 1`. For RIGHT 
TSMM, it uses the fast path when `numRowBlocks == 1`. For multi-tile outputs, 
it still uses the newer `CachingStream` / join / reduce path.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to