[mlir][sparse] avoid creating temporary unordered COO buffer when reshape sparse tensor.

[Peiming]

fix bug

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

387–388

These comment doc needs some updating

397

debug info?

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

391

I'm pretty sure that should be dstTp

[Peiming]

remove debugging code.

[Peiming]

update comments.

[Peiming]

update comments

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

390–391

I think it'd be cleaner to use SparseTensorType from the beginning, instead of having both RankedTensorType and SparseTensorType variants floating around. That's very easily doable by changing the first few lines of this method to:

const auto srcTp = getSparseTensorType(srcTensor);
const auto dstTp = getSparseTensorType(op.getResult());
if (!srcTp.hasEncoding() || !dstTp.hasEncoding())
  return failure();

Since RankedTensorType supports all the helper methods of SparseTensorEncodingAttr, you shouldn't need the srcEnc/dstEnc variables anymore (can just use srcTp/dstTp instead). But if there are places where you need them, you can always:

const auto srcEnc = srcTp.getEncoding();
const auto dstEnc = dstTp.getEndocing();

[wrengr]

(Assuming my previous comment is implemented)

[Peiming]

address comments.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

public:
   using OpRewritePattern<ReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ReshapeOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     Value srcTensor = op.getSrc();
     auto srcTp = getRankedTensorType(srcTensor);
     auto dstTp = getRankedTensorType(op.getResult());
-    SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
-    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
-    if (!encDst || !encSrc) {
+
+    SparseTensorType srcStt(srcTp);
+    SparseTensorType dstStt(dstTp);
+
+    const auto encSrc = srcStt.getEncoding();
+    if (!srcStt.hasEncoding() || !dstStt.hasEncoding()) {
       return failure();
     }
 
     // Generate code to represent the static dimension constants or compute
     // the dynamic dimension values.
     SmallVector<Value> srcSizes;
     sizesForTensor(rewriter, srcSizes, loc, srcTp, srcTensor);
     SmallVector<Value> dstSizes;
     SmallVector<Value> dstDynSizes;
     if (dstTp.hasStaticShape()) {
       for (auto d : dstTp.getShape())
         dstSizes.push_back(constantIndex(rewriter, loc, d));
     } else {
       ArrayRef<int64_t> dstShape = dstTp.getShape();
       genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
                          op.getReassociationIndices());
       for (auto [idx, shape] : llvm::enumerate(dstShape)) {
         if (shape == ShapedType::kDynamic)
           dstDynSizes.push_back(dstSizes[idx]);
       }
     }
-
-    // Implement the sparse2sparse reshape as follows:
-    //   %tmp = bufferization.alloc_tensor : unordered COO
-    //   foreach srcCoords %srcTensor
-    //     insert reshapeCvs(srcCoords), %tmp
-    //   %t = sparse_tensor.cast %tmp
     Value nnz = rewriter.create<NumberOfEntriesOp>(loc, srcTensor);
-    RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
-    Value cooBuffer =
+    // Only need a unordered COO buffer if input and output are not sorted
+    // in the same way.
+    Type bufferTp =
+        srcStt.isAllOrdered() && srcStt.isIdentity() && dstStt.isIdentity()
+            ? dstTp
+            : getUnorderedCOOFromType(dstTp);
+
+    Value buffer =
         rewriter
-            .create<AllocTensorOp>(loc, cooTp, dstDynSizes, Value(),
+            .create<AllocTensorOp>(loc, bufferTp, dstDynSizes, Value(),
                                    /*sizeHint=*/nnz, Attribute())
             .getResult();
 
+    // Implement the sparse2sparse reshape as follows:
+    //   foreach srcCoords %srcTensor
+    //     insert reshapeCvs(srcCoords), %buffer
+    //
+    // followed by an optional
+    //   %t = sparse_tensor.cast %tmp
+    // depending on whether the input/output are sorted in the same way.
     ForeachOp foreachOp = rewriter.create<ForeachOp>(
-        loc, srcTensor, cooBuffer,
+        loc, srcTensor, buffer,
         [&](OpBuilder &builder, Location loc, ValueRange srcLcvs, Value v,
             ValueRange reduc) {
           const Dimension dimRank = srcTp.getRank();
           SmallVector<Value> srcDcvs;
           srcDcvs.reserve(dimRank);
           for (Dimension d = 0; d < dimRank; d++) {
             // FIXME: `toStoredDim` is deprecated
             Level lvl = toStoredDim(encSrc, d);
             srcDcvs.push_back(srcLcvs[lvl]);
           }
           SmallVector<Value> dstDcvs;
           reshapeCvs(builder, loc, op.getReassociationIndices(), srcSizes,
                      srcDcvs, dstSizes, dstDcvs);
           auto t = builder.create<InsertOp>(loc, v, reduc.front(), dstDcvs);
           builder.create<sparse_tensor::YieldOp>(loc, t);
         });
-    auto t = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
-    auto converted = rewriter.create<ConvertOp>(loc, dstTp, t).getResult();
+
+    Value t = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
+    if (bufferTp != dstTp) {
+      Value converted = rewriter.create<ConvertOp>(loc, dstTp, t).getResult();
       rewriter.create<DeallocTensorOp>(loc, t);
-    rewriter.replaceOp(op, converted);
+      t = converted;
+    }
+    rewriter.replaceOp(op, t);
     return success();
   }
 };
 
 /// Sparse rewriting rule for sparse-to-dense and dense-to-sparse reshape
 /// operator.
 template <typename ReshapeOp>
 struct ReshapeRewriter : public OpRewritePattern<ReshapeOp> {

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // CHECK-RWT:           %[[SI:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
 // CHECK-RWT:           %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[I]]] : memref<?xf64>
 // CHECK-RWT:           %[[DI0:.*]] = arith.divui %[[SI]], %[[C10]] : index
 // CHECK-RWT:           %[[DI1:.*]] = arith.remui %[[SI]], %[[C10]] : index
 // CHECK-RWT:           %[[NT:.*]] = sparse_tensor.insert %[[SV]] into %[[R]]{{\[}}%[[DI0]], %[[DI1]]]
 // CHECK-RWT:           scf.yield %[[NT:.*]]
 // CHECK-RWT:         }
 // CHECK-RWT:         %[[NT1:.*]] = sparse_tensor.load %[[RET]] hasInserts
-// CHECK-RWT:         %[[T:.*]] = sparse_tensor.convert %[[NT1]]
-// CHECK-RWT:         return %[[T]] : tensor<10x10xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
+// CHECK-RWT-NOT:     sparse_tensor.convert
+// CHECK-RWT:         return %[[NT1]] : tensor<10x10xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
 //
 func.func @sparse_expand(%arg0: tensor<100xf64, #SparseVector>) -> tensor<10x10xf64, #SparseMatrix> {
   %0 = tensor.expand_shape %arg0 [[0, 1]] :
     tensor<100xf64, #SparseVector> into tensor<10x10xf64, #SparseMatrix>
   return %0 : tensor<10x10xf64, #SparseMatrix>
 }
 
 //

 // CHECK-RWT:             %[[T:.*]] = arith.muli %[[SI0]], %[[C10]] : index
 // CHECK-RWT:             %[[DI:.*]] = arith.addi %[[T]], %[[SI1]] : index
 // CHECK-RWT:             %[[R1:.*]] = sparse_tensor.insert %[[SV]] into %[[A1]]{{\[}}%[[DI]]]
 // CHECK-RWT              scf.yield %[[R1]]
 // CHECK-RWT            }
 // CHECK-RWT            scf.yield %[[RET_1]]
 // CHECK-RWT:         }
 // CHECK-RWT:        %[[NT1:.*]] = sparse_tensor.load %[[RET]] hasInserts
-// CHECK-RWT:        %[[T:.*]] = sparse_tensor.convert %[[NT1]]
-// CHECK-RWT:        return %[[T]] : tensor<100xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
+// CHECK-RWT-NOT:    sparse_tensor.convert
+// CHECK-RWT:        return %[[NT1]] : tensor<100xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
 //
 func.func @sparse_collapse(%arg0: tensor<10x10xf64, #SparseMatrix>) -> tensor<100xf64, #SparseVector> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] :
     tensor<10x10xf64, #SparseMatrix> into tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 //

 // CHECK-RWT:           %[[DI0:.*]] = arith.divui %[[SI]], %[[T2]] : index
 // CHECK-RWT:           %[[T3:.*]] = arith.remui %[[SI]], %[[T2]] : index
 // CHECK-RWT:           %[[T4:.*]] = arith.divui %[[T2]], %[[C10]] : index
 // CHECK-RWT:           %[[DI1:.*]] = arith.divui %[[T3]], %[[T4]] : index
 // CHECK-RWT:           %[[NT:.*]] = sparse_tensor.insert %[[SV]] into %[[R]]{{\[}}%[[DI0]], %[[DI1]]]
 // CHECK-RWT:           scf.yield %[[NT]]
 // CHECK-RWT:         }
 // CHECK-RWT:         %[[NT1:.*]] = sparse_tensor.load %[[RET]] hasInserts
-// CHECK-RWT:         %[[T:.*]] = sparse_tensor.convert %[[NT1]]
-// CHECK-RWT:         return %[[T]] : tensor<?x10xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
+// CHECK-RWT-NOT:     sparse_tensor.convert
+// CHECK-RWT:         return %[[NT1]] : tensor<?x10xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
 //
 func.func @dynamic_sparse_expand(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?x10xf64, #SparseMatrix> {
   %0 = tensor.expand_shape %arg0 [[0, 1]] :
     tensor<?xf64, #SparseVector> into tensor<?x10xf64, #SparseMatrix>
   return %0 : tensor<?x10xf64, #SparseMatrix>
 }
 
 //

 // CHECK-RWT:             %[[T4:.*]] = arith.muli %[[SI1]], %[[T3]] : index
 // CHECK-RWT:             %[[DI:.*]] = arith.addi %[[T2]], %[[T4]] : index
 // CHECK-RWT:             %[[NT:.*]] = sparse_tensor.insert %[[SV]] into %[[R1]]{{\[}}%[[DI]]]
 // CHECK-RWT              scf.yield %[[NT]]
 // CHECK-RWT            }
 // CHECK-RWT            scf.yield %[[RET_1]]
 // CHECK-RWT:        }
 // CHECK-RWT:        %[[NT1:.*]] = sparse_tensor.load %[[RET]] hasInserts
-// CHECK-RWT:        %[[T:.*]] = sparse_tensor.convert %[[NT1]]
-// CHECK-RWT:        return %[[T]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
+// CHECK-RWT-NOT:    sparse_tensor.convert
+// CHECK-RWT:        return %[[NT1]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
 //
 func.func @dynamic_sparse_collapse(%arg0: tensor<10x?xf64, #SparseMatrix>) -> tensor<?xf64, #SparseVector> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] :
     tensor<10x?xf64, #SparseMatrix> into tensor<?xf64, #SparseVector>
   return %0 : tensor<?xf64, #SparseVector>
 }