[mlir][sparse] Add rewrite rules for sparse-to-sparse reshape operators.

[Peiming]

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

382

Do we really need to distinguish enableRT here? Because I see no problem of rewriting reshape/concatenate even when RT is enabled.

Since the from/toCOO <==> foreach operator, we can relies on RT only for insertion?

But I am okay with how it works right now, maybe it should not be addressed now...

[Peiming]

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

425–426

Maybe you should use the builder from callback

[bixia]

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

382

If enableRT means keep using the RT libs as we are right now, we should rerwrite sparse2sparse reshape here.

[bixia]

Use the foreach builder.

[bixia]

Rebase.

[Peiming]

LGTM.

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

385–386

Isn't it already guaranteed by verifier?

Anyway, no harm to do extra check

[aartbik]

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

301–366

not to nitpick, but we always use singular here, since each rule applies to one operator

382

But isn't it easier to have two ruless then, and add them below based on enableRT or not
(see how Peiming added rule conditionally at L540?

Having the bool leak into the constructors with different codepaths is harder to read than just applying different rule sets?

517

why is this so far away from the other reshape rewriter def?

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

3

shall we make a new test flag --test_sparse_rewriting that only calls populateSparseTensorRewriting() and applies those rules, so we can test this rewriting without having the full sparsification?

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

3

This lifts rewriting into its own pass

https://reviews.llvm.org/D135319

so we can test the rewrite in isolation from sparsification

[bixia]

Split the reshape rewrite rules.

[bixia]

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

301–366

I was thinking this is a template and is used to two reshape operators. But use singular instead.

382

Split into two rules anyway. It is a little different in the sense that we now add two rewrite rules for one operator.

385–386

Remove the check

[bixia]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

3

Thanks. Will update the test when this is landed.

[bixia]

Remove the check for element type.

[aartbik]

Yeah, this looks really good. I will wait for the CHECK test simplification, but then good to go!

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

301–366

Ah, yeah, that is a good point too actually ;-)

382

I love how it looks now. Having two rules for same op is fine, since the logic is mutually exclusive, right?!

[bixia]

Use sparse-tensor-rewrite for the test.

[bixia]

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

382

Yes, the two rewrite are mutually exclusive.

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

3

Fixed the test.

[bixia]

Rebase

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

if (auto arg = yieldOp.getOperand(0).dyn_cast<BlockArgument>()) {
     if (arg.getOwner()->getParentOp() == op) {
       OpOperand *t = op.getInputAndOutputOperands()[arg.getArgNumber()];
       return isZeroValue(t->get());
     }
   }
   return isZeroValue(yieldOp.getOperand(0));
 }
 
+/// Populates given sizes array from type (for static sizes) and from
+/// the tensor (for dynamic sizes).
+static void sizesForTensor(OpBuilder &builder, SmallVector<Value, 4> &sizes,
+                           Location loc, ShapedType stp, Value tensor) {
+  for (const auto &d : enumerate(stp.getShape())) {
+    Value dim;
+    if (d.value() == ShapedType::kDynamicSize)
+      dim = builder.create<tensor::DimOp>(loc, tensor, d.index());
+    else
+      dim = constantIndex(builder, loc, d.value());
+    sizes.push_back(dim);
+  }
+}
+
 // TODO: The dim level property of the COO type relies on input tensors, the
 // shape relies on the output tensor
 // Helpers to setup a COO type.
 static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
   auto *ctx = src.getContext();
   auto rank = src.getRank();
   SmallVector<SparseTensorEncodingAttr::DimLevelType, 4> dims;
 
-  // An unordered and non-unique compressed dim at beginning.
+  // An unordered and non-unique compressed dim at beginning unless the tensor
+  // is a 1D tensor.
+  if (rank > 1)
     dims.push_back(SparseTensorEncodingAttr::DimLevelType::CompressedNuNo);
+
   // TODO: it is actually ordered at the level for ordered input.
   // Followed by unordered non-unique n-2 singleton levels.
   std::fill_n(std::back_inserter(dims), rank - 2,
               SparseTensorEncodingAttr::DimLevelType::SingletonNuNo);
   // TODO: only if all the inputs (for concatentate) are unique at the last
   // level should the COO has a unique level at the end. Ends by a unordered
   // unique singleton level.
   dims.push_back(SparseTensorEncodingAttr::DimLevelType::SingletonNo);

 
 private:
   // Helper to add argument and record the mapping.
   static void addArg(BlockAndValueMapping &mapper, Block *b, BlockArgument a) {
     mapper.map(a, b->addArgument(a.getType(), a.getLoc()));
   }
 };
 
-/// Sparse rewriting rule for reshape operator.
+/// Sparse rewriting rule for sparse-to-sparse reshape operator.
+template <typename ReshapeOp>
+struct Sparse2SparseReshapeRewriter : public OpRewritePattern<ReshapeOp> {
+public:
+  using OpRewritePattern<ReshapeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ReshapeOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    Value srcTensor = op.getSrc();
+    auto srcTp = srcTensor.getType().template cast<RankedTensorType>();
+    auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
+    SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
+    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
+    if (!encDst || !encSrc) {
+      return failure();
+    }
+
+    // Generate code to represent the static dimension constants or compute
+    // the dynamic dimension values.
+    SmallVector<Value, 4> srcSizes;
+    sizesForTensor(rewriter, srcSizes, loc, srcTp, srcTensor);
+    SmallVector<Value, 4> dstSizes;
+    SmallVector<Value, 4> 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 &d : llvm::enumerate(dstShape)) {
+        if (d.value() == ShapedType::kDynamicSize)
+          dstDynSizes.push_back(dstSizes[d.index()]);
+      }
+    }
+
+    // Implement the sparse2sparse reshape as follows:
+    //   %tmp = bufferization.alloc_tensor : unordered COO
+    //   foreach srcCoords %srcTensor
+    //     insert translateIndicesArray(srcCoords), %tmp
+    //   %t = sparse_tensor.cast %tmp
+    RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
+    auto cooBuffer =
+        rewriter.create<AllocTensorOp>(loc, cooTp, dstDynSizes).getResult();
+    rewriter.create<ForeachOp>(
+        loc, srcTensor, [&](OpBuilder &builder, Location loc, ValueRange args) {
+          SmallVector<Value, 4> srcIndices;
+          SmallVector<Value, 4> dstIndices;
+          for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
+            uint64_t dim = toStoredDim(encSrc, i);
+            srcIndices.push_back(args[dim]);
+          }
+          translateIndicesArray(builder, loc, op.getReassociationIndices(),
+                                srcIndices, srcSizes, dstSizes, dstIndices);
+          builder.create<InsertOp>(loc, args.back(), cooBuffer, dstIndices);
+          builder.create<sparse_tensor::YieldOp>(loc);
+        });
+
+    rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, cooBuffer);
+    return success();
+  }
+};
+
+/// Sparse rewriting rule for sparse-to-dense and dense-to-sparse reshape
+/// operator.
 template <typename ReshapeOp>
 struct ReshapeRewriter : public OpRewritePattern<ReshapeOp> {
 public:
   using OpRewritePattern<ReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ReshapeOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     auto encDst = getSparseTensorEncoding(op.getResult().getType());
     auto encSrc = getSparseTensorEncoding(op.getSrc().getType());
     // Since a pure dense expansion is very cheap (change of view), for
     // a sparse2dense or dense2sparse, we can simply unfuse a sparse
     // conversion from the reshape operation itself.
     // All other cases are handled elsewhere.
     if (encDst && encSrc) {
       return failure();
     }
     if (encSrc) {
       RankedTensorType rtp =
           op.getSrc().getType().template cast<RankedTensorType>();
       auto denseTp =
           RankedTensorType::get(rtp.getShape(), rtp.getElementType());
       auto convert = rewriter.create<ConvertOp>(loc, denseTp, op.getSrc());
       op->setOperand(0, convert);
       return success();
     }
     if (encDst) {
       RankedTensorType rtp =
LogicalResult matchAndRewrite(ConcatenateOp op,
     size_t conDim = op.getDimension().getZExtValue();
     // %t = concatenate %s1, %s2, %s3 {dim = 1}
     // ==>
     // %tmp = bufferization.alloc_tensor : unordered COO
     // foreach in %s1 : insert d0, d1, %tmp
     // foreach in %s2 : insert d0, d1 + size(s1), %tmp
     // foreach in %s3 : insert d0, d1 + size(s1) + size(s2), %tmp
     // %t = sparse_tensor.cast %tmp
     auto cooTp = getUnorderedCOOFromType(rtp);
     auto cooBuffer =
         rewriter.create<AllocTensorOp>(loc, cooTp, ValueRange()).getResult();
 
     Value offset = constantIndex(rewriter, loc, 0);
     for (Value input : op.getInputs()) {
       // Builds the indexing map.
 
       // Build a for op for each input tensor to append new values into the
       // output tensor.
LogicalResult matchAndRewrite(ForeachOp op,
     // Inline body.
     rewriter.mergeBlockBefore(srcBlock, &*inlinePos, args);
     // delete the foreach operator.
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
-
 void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns,
                                          bool enableRT) {
   patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
                ReshapeRewriter<tensor::ExpandShapeOp>,
                ReshapeRewriter<tensor::CollapseShapeOp>, ForeachRewriter>(
       patterns.getContext());
   // TODO: If RT not enabled, rewrite concatenate ops, etc here.
   if (!enableRT)
-    patterns.add<ConcatenateRewriter>(patterns.getContext());
+    patterns.add<ConcatenateRewriter,
+                 Sparse2SparseReshapeRewriter<tensor::ExpandShapeOp>,
+                 Sparse2SparseReshapeRewriter<tensor::CollapseShapeOp>>(
+        patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s --check-prefix=CHECK-ROUND
 // RUN: mlir-opt %s --sparse-tensor-conversion --cse --canonicalize | FileCheck %s --check-prefix=CHECK-CONV
+// RUN: mlir-opt %s --sparse-tensor-rewrite=enable-runtime-library=false --cse --canonicalize  | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 #SparseMatrix = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 //
 // roundtrip:
 //
 // CHECK-ROUND-LABEL: func.func @sparse_expand(

 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
+// rewrite for codegen:
+//
+// CHECK-RWT-LABEL:   func.func @sparse_expand(
+// CHECK-RWT-SAME:    %[[S:.*]]:
+// CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
+// CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK-RWT-DAG:     %[[C1:.*]] = arith.constant 1 : index
+// CHECK-RWT:         %[[B:.*]] = bufferization.alloc_tensor()
+// CHECK-RWT:         %[[P0:.*]] = sparse_tensor.pointers %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[V:.*]] = sparse_tensor.values %[[S]]
+// CHECK-RWT:         %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
+// CHECK-RWT:         %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
+// CHECK-RWT:         scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
+// 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:           sparse_tensor.insert %[[SV]] into %[[B]]{{\[}}%[[DI0]], %[[DI1]]]
+// CHECK-RWT:         }
+// CHECK-RWT:         %[[T:.*]] = sparse_tensor.convert %[[B]]
+// CHECK-RWT:         return %[[T]] : 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>
 }
 
 //
 // roundtrip:

 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
+// rewrite for codegen:
+//
+// CHECK-RWT-LABEL:   func.func @sparse_collapse(
+// CHECK-RWT-SAME:    %[[S:.*]]:
+// CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
+// CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK-RWT-DAG:     %[[C1:.*]] = arith.constant 1 : index
+// CHECK-RWT:         %[[B:.*]] = bufferization.alloc_tensor()
+// CHECK-RWT:         %[[P0:.*]] = sparse_tensor.pointers %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
+// CHECK-RWT:          %[[P1:.*]] = sparse_tensor.pointers %[[S]] {dimension = 1 : index}
+// CHECK-RWT:          %[[I1:.*]] = sparse_tensor.indices %[[S]] {dimension = 1 : index}
+// CHECK-RWT:          %[[V:.*]] = sparse_tensor.values %[[S]]
+// CHECK-RWT:          %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
+// CHECK-RWT:          %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
+// CHECK-RWT:          scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
+// CHECK-RWT:            %[[SI0:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
+// CHECK-RWT:            %[[S1:.*]] = memref.load %[[P1]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[E1:.*]] = memref.load %[[P1]]{{\[}}%[[PE1]]] : memref<?xindex>
+// CHECK-RWT:            scf.for %[[J:.*]] = %[[S1]] to %[[E1]] step %[[C1]] {
+// CHECK-RWT:              %[[SI1:.*]] = memref.load %[[I1]]{{\[}}%[[J]]] : memref<?xindex>
+// CHECK-RWT:              %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[J]]] : memref<?xf64>
+// CHECK-RWT:              %[[T:.*]] = arith.muli %[[SI0]], %[[C10]] : index
+// CHECK-RWT:              %[[DI:.*]] = arith.addi %[[T]], %[[SI1]] : index
+// CHECK-RWT:              sparse_tensor.insert %[[SV]] into %[[B]]{{\[}}%[[DI]]]
+// CHECK-RWT             }
+// CHECK-RWT:          }
+// CHECK-RWT:        %[[T:.*]] = sparse_tensor.convert %[[B]]
+// CHECK-RWT:        return %[[T]] : 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>
 }
 
 //
 // roundtrip:

 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
+// rewrite for codegen:
+//
+// CHECK-RWT-LABEL:   func.func @dynamic_sparse_expand(
+// CHECK-RWT-SAME:    %[[S:.*]]:
+// CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
+// CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK-RWT-DAG:     %[[C1:.*]] = arith.constant 1 : index
+// CHECK-RWT:         %[[SD:.*]] = tensor.dim %[[S]], %[[C0]]
+// CHECK-RWT:         %[[DD0:.*]] = arith.divui %[[SD]], %[[C10]] : index
+// CHECK-RWT:         %[[B:.*]] = bufferization.alloc_tensor(%[[DD0]])
+// CHECK-RWT:         %[[P0:.*]] = sparse_tensor.pointers %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[V:.*]] = sparse_tensor.values %[[S]]
+// CHECK-RWT:         %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
+// CHECK-RWT:         %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
+// CHECK-RWT:         scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
+// CHECK-RWT:           %[[SI:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:           %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[I]]] : memref<?xf64>
+// CHECK-RWT:           %[[T1:.*]] = arith.muli %[[DD0]], %[[C10]] : index
+// CHECK-RWT:           %[[T2:.*]] = arith.divui %[[T1]], %[[DD0]] : index
+// 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:           sparse_tensor.insert %[[SV]] into %[[B]]{{\[}}%[[DI0]], %[[DI1]]]
+// CHECK-RWT:         }
+// CHECK-RWT:         %[[T:.*]] = sparse_tensor.convert %[[B]]
+// CHECK-RWT:         return %[[T]] : 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>
 }
 
 //
 // roundtrip:

 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      call @delSparseTensorCOOF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
+// rewrite for codegen:
+//
+// CHECK-RWT-LABEL:   func.func @dynamic_sparse_collapse(
+// CHECK-RWT-SAME:    %[[S:.*]]:
+// CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
+// CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK-RWT-DAG:     %[[C1:.*]] = arith.constant 1 : index
+// CHECK-RWT:         %[[SD1:.*]] = tensor.dim %[[S]], %[[C1]]
+// CHECK-RWT:         %[[DD0:.*]] = arith.muli %[[SD1]], %[[C10]] : index
+// CHECK-RWT:         %[[B:.*]] = bufferization.alloc_tensor(%[[DD0]])
+// CHECK-RWT:         %[[P0:.*]] = sparse_tensor.pointers %[[S]] {dimension = 0 : index}
+// CHECK-RWT:         %[[I0:.*]] = sparse_tensor.indices %[[S]] {dimension = 0 : index}
+// CHECK-RWT:          %[[P1:.*]] = sparse_tensor.pointers %[[S]] {dimension = 1 : index}
+// CHECK-RWT:          %[[I1:.*]] = sparse_tensor.indices %[[S]] {dimension = 1 : index}
+// CHECK-RWT:          %[[V:.*]] = sparse_tensor.values %[[S]]
+// CHECK-RWT:          %[[S0:.*]] = memref.load %[[P0]]{{\[}}%[[C0]]] : memref<?xindex>
+// CHECK-RWT:          %[[E0:.*]] = memref.load %[[P0]]{{\[}}%[[C1]]] : memref<?xindex>
+// CHECK-RWT:          scf.for %[[I:.*]] = %[[S0]] to %[[E0]] step %[[C1]] {
+// CHECK-RWT:            %[[SI0:.*]] = memref.load %[[I0]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[PE1:.*]] = arith.addi %[[I]], %[[C1]] : index
+// CHECK-RWT:            %[[S1:.*]] = memref.load %[[P1]]{{\[}}%[[I]]] : memref<?xindex>
+// CHECK-RWT:            %[[E1:.*]] = memref.load %[[P1]]{{\[}}%[[PE1]]] : memref<?xindex>
+// CHECK-RWT:            scf.for %[[J:.*]] = %[[S1]] to %[[E1]] step %[[C1]] {
+// CHECK-RWT:              %[[SI1:.*]] = memref.load %[[I1]]{{\[}}%[[J]]] : memref<?xindex>
+// CHECK-RWT:              %[[SV:.*]] = memref.load %[[V]]{{\[}}%[[J]]] : memref<?xf64>
+// CHECK-RWT:               %[[T1:.*]] = arith.divui %[[DD0]], %[[C10]] : index
+// CHECK-RWT:               %[[T2:.*]] = arith.muli %[[SI0]], %[[T1]] : index
+// CHECK-RWT:               %[[T3:.*]] = arith.divui %[[T1]], %[[SD1]] : index
+// CHECK-RWT:               %[[T4:.*]] = arith.muli %[[SI1]], %[[T3]] : index
+// CHECK-RWT:               %[[DI:.*]] = arith.addi %[[T2]], %[[T4]] : index
+// CHECK-RWT:              sparse_tensor.insert %[[SV]] into %[[B]]{{\[}}%[[DI]]]
+// CHECK-RWT             }
+// CHECK-RWT:          }
+// CHECK-RWT:        %[[T:.*]] = sparse_tensor.convert %[[B]]
+// CHECK-RWT:        return %[[T]] : 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>
 }