[mlir][sparse] Add rewriting rule for the convert operator.

[aartbik]

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

159

This feels like it could be moved to IR/SparseTensor.h where we have several related tests on ranked tensor type and sparsity. WDYT?

162

I don't think we can have enc + rank == 0, so second test is strictly speaking not needed

166

Even for rank == 1, don't we need a compressed outermost?

A compressed sparse vector is basically in COO format ;-)
You cannot have a singleton as first entry

So I think you can simply to the test that compressed is outermost, followed by singletons.
For vectors, the folllowed by part will be skipped

[Peiming]

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

555

Do you think it would be nicer if the body builder provide
(OpBuilder &builder, Location loc, Value v, ValueRange indices)? so that the value and indices are separated?

[Peiming]

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

499–511

I think these two can be implemented in loop emitter probably

499–511

Probably not sparse constant.

[Peiming]

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

499–511

I can extend foreach operator to accept these inputs if it is more desirable.

[bixia]

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

159

Good idea.

499–511

the sparse constant actually have two constants: one for indices and one for values.
We currently use an existing block of code genDenseTensorOrSparseConstantIterLoop for the conversion path.
We modify this if we will have better supporting functions.

555

I don't see a big difference, the current interface is fine as the value is in args.back().

[Peiming]

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

499–511

Okay, SG.

I extended the dense case. sparse constant probably need more work!

[bixia]

Rebase.

[bixia]

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

499–511

Thanks! I will see if I can modify genDenseTensorOrSparseConstantIterLoop for the codegen and conversion to take advantage of the fact that foreach can handle dense tensor.
I will do this in a separate PR.

[aartbik]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

275

very minor nit, I would use

for (uint64_t i = 1, rank = tp.getRank();  i < rank; ++i)

since it is only used in the for loop, but up to you

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

160

I think we should avoid returning vectors by value (even though the compiler may actually do some interesting optimizations). For such an output vector, simply pass in by reference

488

comment on dense2dense?

580

can be move that up into the actual codegen rewriting rule above, so it is now hidden inside this method?

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

1

add a FIXME or TODO

[bixia]

Address review comments.

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

275

right, done.

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

488

Add comment to stay that dense-to-dense is a nop handled by canonicalization.

580

Yes, moved this check up.

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

1

this was a mistake. Fixed.

[aartbik]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

265

do we care about ordered/unique here? [there are a few different flavors of COO]

given the usage later, I think we can accept sorted and unsorted, but probably not with duplicates?

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

159

document that dynSizes is the output (also, not sure if we have a convention on this, but I usually put input vectors before output parameters, especially now that they have the same reference type). WDYT?

[bixia]

Rename isCOOType to isUniqueCOOType and checks that the last dim is unique.
Reorder parameters for getDynamicSizes.

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

265

Add a check to make sure that the last dimension is unique.
Rename the func to isUniqueCOOType.

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

159

Add doc description for dynSizes.
C++ coding style says order input params before output params. Reorder the params.

[aartbik]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

278

make a note that rank == 1 works (unique the only compressed) as does rank > 1 (unique on the last singleton)

[bixia]

Add a comment for isUniqueCOOType.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 
 namespace mlir {
 namespace sparse_tensor {
 
 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
+/// Returns true iff the given type is a type for a COO tensor.
+bool isCOOType(RankedTensorType tp);
+
 //
 // Dimension level types.
 //
 
 // MSVC does not allow this function to be constexpr, because
 // `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
 // And therefore all functions calling it cannot be constexpr either.
 // TODO: since Clang does allow these to be constexpr, perhaps we should

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 SparseTensorEncodingAttr
 mlir::sparse_tensor::getSparseTensorEncoding(Type type) {
   if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
   return nullptr;
 }
 
+bool isCOOType(RankedTensorType tp) {
+  SparseTensorEncodingAttr enc = getSparseTensorEncoding(tp);
+
+  if (!enc)
+    return false;
+
+  if (!isCompressedDim(tp, 0))
+    return false;
+
+  uint64_t rank = tp.getRank();
+  for (uint64_t i = 1; i < rank; ++i)
+    if (!isSingletonDim(tp, i))
+      return false;
+
+  return true;
+}
+
 uint64_t mlir::sparse_tensor::toOrigDim(const SparseTensorEncodingAttr &enc,
                                         uint64_t d) {
   if (enc) {
     auto order = enc.getDimOrdering();
     if (order) {
       assert(order.isPermutation());
       return order.getDimPosition(d);
     }

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

static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
   // largest one among them) in the original operation instead of using the
   // default value.
   auto enc = SparseTensorEncodingAttr::get(
       ctx, dims, AffineMap::getMultiDimIdentityMap(rank, ctx), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
   return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
 }
 
+/// Collects the dynamic dimension sizes for `tp` with the assumption that
+/// `sizes` are the dimension sizes for the type.
+static SmallVector<Value, 4> getDynamicSizes(RankedTensorType tp,
+                                             SmallVectorImpl<Value> &sizes) {
+  SmallVector<Value, 4> dynSizes;
+  for (const auto &d : enumerate(tp.getShape())) {
+    if (d.value() == ShapedType::kDynamicSize)
+      dynSizes.push_back(sizes[d.index()]);
+  }
+  return dynSizes;
+}
+
 //===---------------------------------------------------------------------===//
 // The actual sparse tensor rewriting rules.
 //===---------------------------------------------------------------------===//
 
 namespace {
 
 /// Rewriting rule that converts direct yield of zero with initial allocation.
 struct FoldInvariantYield : public OpRewritePattern<GenericOp> {
for (Value input : op.getInputs()) {
       offset = rewriter.create<arith::AddIOp>(loc, offset,
                                               constantIndex(rewriter, loc, d));
     }
     rewriter.replaceOpWithNewOp<ConvertOp>(op, rtp, cooBuffer);
     return success();
   }
 };
 
+/// Sparse rewriting rule for the convert operator.
+struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(ConvertOp op,
+                                PatternRewriter &rewriter) const override {
+    auto encDst = getSparseTensorEncoding(op.getType());
+    auto encSrc = getSparseTensorEncoding(op.getSource().getType());
+    if (encDst && encSrc)
+      return sparse2SparseRewrite(op, rewriter);
+    if (encSrc && !encDst)
+      return sparse2DenseRewrite(op, rewriter);
+    if (!encSrc && encDst)
+      return dense2SparseRewrite(op, rewriter);
+
+    return failure();
+  }
+
+private:
+  // Handles sparse constant to sparse tensor or dense tensor to sparse tensor
+  // conversion as follows:
+  //   t = new sparse COO tensor
+  //   fill t using src
+  //   dst = convert t
+  //
+  // To fill the COO tensor from a dense tensor:
+  //   for i1 in dim1
+  //    ..
+  //     for ik in dimk
+  //       val = a[i1,..,ik]
+  //       if val != 0
+  //         t->add(val, [i1,..,ik], [p1,..,pk])
+  //
+  // To fill the COO tensor from a sparse constant in COO format:
+  //   for i in range(NNZ)
+  //     val = values[i]
+  //     [i1,..,ik] = indices[i]
+  //     t->add(val, [i1,..,ik], [p1,..,pk])
+  LogicalResult dense2SparseRewrite(ConvertOp op,
+                                    PatternRewriter &rewriter) const {
+    Location loc = op.getLoc();
+    Value src = op.getSource();
+    RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
+    SmallVector<Value, 4> sizes;
+    sizesFromSrc(rewriter, sizes, loc, src);
+    SmallVector<Value, 4> dynSizes = getDynamicSizes(dstTp, sizes);
+
+    RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
+    auto cooBuffer =
+        rewriter.create<AllocTensorOp>(loc, cooTp, dynSizes).getResult();
+    unsigned rank = dstTp.cast<ShapedType>().getRank();
+
+    genDenseTensorOrSparseConstantIterLoop(
+        rewriter, loc, src, rank,
+        [&](OpBuilder &builder, Location loc, Value val, ValueRange indices) {
+          builder.create<InsertOp>(loc, val, cooBuffer, indices);
+        });
+
+    rewriter.setInsertionPointAfter(op);
+    rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, cooBuffer);
+    rewriter.create<DeallocTensorOp>(loc, cooBuffer);
+
+    return success();
+  }
+
+  // Handles sparse tensor to dense tensor conversion as follows:
+  //   dst = new dense tensor;
+  //   foreach elemment in src
+  //     dst[elemment.indices] = element.value
+  LogicalResult sparse2DenseRewrite(ConvertOp op,
+                                    PatternRewriter &rewriter) const {
+    Location loc = op->getLoc();
+    RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
+    Value src = op.getSource();
+    RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
+
+    SmallVector<Value, 4> sizes;
+    sizesForTensor(rewriter, sizes, loc, srcTp, src);
+    Value dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
+
+    rewriter.create<ForeachOp>(
+        loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+          builder.create<memref::StoreOp>(loc, args.back(), dst,
+                                          args.drop_back());
+          builder.create<sparse_tensor::YieldOp>(loc);
+        });
+
+    rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
+    return success();
+  }
+
+  // Handles sparse tensor to sparse tensor conversion as follows:
+  //   if src is not COO
+  //       construct a COO to represent the src
+  //   sort the src COO
+  //   foreach elemment in the sorted src COO
+  //     insert element to dst
+  LogicalResult sparse2SparseRewrite(ConvertOp op,
+                                     PatternRewriter &rewriter) const {
+    Location loc = op->getLoc();
+    Value src = op.getSource();
+    RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
+    RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
+    SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
+    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
+
+    if (encSrc == encDst) {
+      // Trivial tensor conversion is handled in codegen.
+      return failure();
+    }
+
+    SmallVector<Value, 4> srcSizes;
+    sizesForTensor(rewriter, srcSizes, loc, srcTp, src);
+    Value tmpCoo = Value();
+    if (!isCOOType(srcTp)) {
+      // Construct a COO tensor from the src tensor.
+      // TODO: there may be cases for which more efficiently without
+      // going through an intermediate COO, such as cases that only change
+      // the overhead types.
+      SmallVector<Value, 4> dynSrcSizes = getDynamicSizes(srcTp, srcSizes);
+      srcTp = getUnorderedCOOFromType(srcTp);
+      tmpCoo =
+          rewriter.create<AllocTensorOp>(loc, srcTp, dynSrcSizes).getResult();
+      rewriter.create<ForeachOp>(
+          loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+            SmallVector<Value, 4> indices;
+            for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
+              uint64_t dim = toStoredDim(encSrc, i);
+              indices.push_back(args[dim]);
+            }
+            builder.create<InsertOp>(loc, args.back(), tmpCoo, indices);
+            builder.create<sparse_tensor::YieldOp>(loc);
+          });
+      src = tmpCoo;
+    }
+
+    // Sort the COO tensor so that its elements are ordered via increasing
+    // indices for the storage ordering of the dst tensor.
+    auto dynShape = {ShapedType::kDynamicSize};
+    auto indTp =
+        MemRefType::get(dynShape, getIndexOverheadType(rewriter, encSrc));
+    uint64_t rank = dstTp.getRank();
+    // Gather the indices-arrays in the dst tensor storage order.
+    SmallVector<Value, 4> xs(rank, Value());
+    for (int64_t i = 0; i < rank; i++) {
+      uint64_t orgDim = toOrigDim(encSrc, i);
+      xs[toStoredDim(encDst, orgDim)] = rewriter.create<ToIndicesOp>(
+          loc, indTp, src, rewriter.getIndexAttr(orgDim));
+    }
+
+    // Retrieve NNZ.
+    auto ptrTp =
+        MemRefType::get(dynShape, getPointerOverheadType(rewriter, encSrc));
+    Value p0 =
+        rewriter.create<ToIndicesOp>(loc, ptrTp, src, rewriter.getIndexAttr(0));
+    Value c1 = constantIndex(rewriter, loc, 1);
+    Value nnz = rewriter.create<memref::LoadOp>(loc, p0, c1);
+    nnz =
+        rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), nnz);
+
+    // Retrieve the values-array.
+    auto valTp = MemRefType::get(dynShape, srcTp.getElementType());
+    Value y = rewriter.create<ToValuesOp>(loc, valTp, src);
+
+    // Sort the COO tensor.
+    rewriter.create<SortOp>(loc, nnz, xs, ValueRange{y});
+
+    // For each element in the COO tensor, insert the element to the dst tensor.
+    SmallVector<Value, 4> dynDstSizes = getDynamicSizes(dstTp, srcSizes);
+    Value dst =
+        rewriter.create<AllocTensorOp>(loc, dstTp, dynDstSizes).getResult();
+    rewriter.create<ForeachOp>(
+        loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
+          SmallVector<Value, 4> indices;
+          for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
+            uint64_t dim = toStoredDim(encDst, i);
+            indices.push_back(args[dim]);
+          }
+          builder.create<InsertOp>(loc, args.back(), dst, indices);
+          builder.create<sparse_tensor::YieldOp>(loc);
+        });
+
+    // Release the temporary COO if it is created.
+    if (tmpCoo)
+      rewriter.create<DeallocTensorOp>(loc, tmpCoo);
+
+    // Directly replace op with dst results in bufferization error message
+    // "sparse tensor allocation should not escape function".
+    // As such, we insert a trivial tensor convert which will be removed by
+    // codegen.
+    rewriter.setInsertionPointAfter(op);
+    rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, dst);
+    return success();
+  }
+};
+
 /// Sparse rewriting rule for the foreach operator.
 struct ForeachRewriter : public OpRewritePattern<ForeachOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForeachOp op,
                                 PatternRewriter &rewriter) const override {
 

 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns,
                                          bool enableRT, bool enableForeach,
-                                         bool /*enableConvert*/) {
+                                         bool enableConvert) {
   patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
                ReshapeRewriter<tensor::ExpandShapeOp>,
                ReshapeRewriter<tensor::CollapseShapeOp>>(patterns.getContext());
   if (enableForeach)
     patterns.add<ForeachRewriter>(patterns.getContext());
-
   // TODO: If RT not enabled, rewrite concatenate ops, etc here.
-  if (!enableRT)
+  if (!enableRT) {
     patterns.add<ConcatenateRewriter, NewRewriter,
                  Sparse2SparseReshapeRewriter<tensor::ExpandShapeOp>,
                  Sparse2SparseReshapeRewriter<tensor::CollapseShapeOp>>(
         patterns.getContext());
+    if (enableConvert)
+      patterns.add<ConvertRewriter>(patterns.getContext());
+  }
 }

mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir

 // RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
+// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>

 //       CHECK:     memref.store %[[J]], %[[M]][%[[C1]]] : memref<2xindex>
 //       CHECK:     memref.store %[[E]], %[[BUF]][] : memref<f64>
 //       CHECK:     call @addEltF64(%[[C]], %[[BUF]], %[[T]], %[[Z]])
 //       CHECK:   }
 //       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]])
 //       CHECK: call @delSparseTensorCOOF64(%[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
+
+// CHECK-RWT-LABEL:   func.func @sparse_convert_2d(
+//  CHECK-RWT-SAME:   %[[A:.*]]: tensor<2x4xf64>) -> tensor<2x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
+//  CHECK-RWT-DAG:    %[[C0:.*]] = arith.constant 0 : index
+//  CHECK-RWT-DAG:    %[[C1:.*]] = arith.constant 1 : index
+//  CHECK-RWT-DAG:    %[[C2:.*]] = arith.constant 2 : index
+//  CHECK-RWT-DAG:    %[[C4:.*]] = arith.constant 4 : index
+//  CHECK-RWT-DAG:    %[[F0:.*]] = arith.constant 0.000000e+00 : f64
+//      CHECK-RWT:    %[[COO:.*]] = bufferization.alloc_tensor()
+//      CHECK-RWT:    scf.for %[[FI:.*]] = %[[C0]] to %[[C2]] step %[[C1]] {
+//      CHECK-RWT:      scf.for %[[FJ:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
+//      CHECK-RWT:        %[[V:.*]] = tensor.extract %[[A]]{{\[}}%[[FI]], %[[FJ]]] : tensor<2x4xf64>
+//      CHECK-RWT:        %[[NZ:.*]] = arith.cmpf une, %[[V]], %[[F0]] : f64
+//      CHECK-RWT:        scf.if %[[NZ]] {
+//      CHECK-RWT:          %{{.*}} = sparse_tensor.insert %[[V]] into %[[COO]]{{\[}}%[[FI]], %[[FJ]]]
+//      CHECK-RWT:        }
+//      CHECK-RWT:      }
+//      CHECK-RWT:    }
+//      CHECK-RWT:    %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
+//      CHECK-RWT:    %[[I1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
+//      CHECK-RWT:    %[[NNZ:.*]] = memref.load %[[I0]]{{\[}}%[[C1]]] : memref<?xindex>
+//      CHECK-RWT:    %[[V2:.*]] = sparse_tensor.values %[[COO]]
+//      CHECK-RWT:    sparse_tensor.sort %[[NNZ]], %[[I0]], %[[I1]] jointly %[[V2]]
+//      CHECK-RWT:    %[[DST:.*]] = bufferization.alloc_tensor()
+//      CHECK-RWT:    sparse_tensor.foreach in %[[COO]]
+//      CHECK-RWT:    ^bb0(%[[FI0:.*]]: index, %[[FI1:.*]]: index, %[[FV:.*]]: f64):
+//      CHECK-RWT:      sparse_tensor.insert %[[FV]] into %[[DST]]{{\[}}%[[FI0]], %[[FI1]]]
+//      CHECK-RWT:    }
+//      CHECK-RWT:    %[[R:.*]] = sparse_tensor.convert %[[DST]]
+//      CHECK-RWT:    bufferization.dealloc_tensor %[[COO]]
+//      CHECK-RWT:    return %[[R]] : tensor<2x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>
 func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
   %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #CSR>
   return %0 : tensor<2x4xf64, #CSR>
 }
 
 // CHECK-LABEL: func @sparse_constant() -> !llvm.ptr<i8> {
 //   CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32
 //   CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32

 //       CHECK-DAG:   memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex>
 //       CHECK-DAG:   %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32>
 //       CHECK:   memref.store %[[V]], %[[BUF]][] : memref<f32>
 //       CHECK:   call @addEltF32(%{{.*}}, %[[BUF]], %[[N]], %{{.*}})
 //       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]])
 //       CHECK: call @delSparseTensorCOOF32(%[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
+
+// CHECK-RWT-LABEL:  func.func @sparse_constant()
+//   CHECK-RWT-DAG:  %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-RWT-DAG:  %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-RWT-DAG:  %[[SI:.*]] = arith.constant dense<{{\[\[}}0, 0], [1, 6]]> : tensor<2x2xi64>
+//   CHECK-RWT-DAG:  %[[SV:.*]] = arith.constant dense<[1.000000e+00, 5.000000e+00]> : tensor<2xf32>
+//   CHECK-RWT-DAG:  %[[C2:.*]] = arith.constant 2 : index
+//       CHECK-RWT:  %[[COO:.*]] = bufferization.alloc_tensor()
+//       CHECK-RWT:  scf.for %[[FI:.*]] = %[[C0]] to %[[C2]] step %[[C1]] {
+//       CHECK-RWT:    %[[I0r:.*]] = tensor.extract %[[SI]]{{\[}}%[[FI]], %[[C0]]] : tensor<2x2xi64>
+//       CHECK-RWT:    %[[I0:.*]] = arith.index_cast %[[I0r]] : i64 to index
+//       CHECK-RWT:    %[[I1r:.*]] = tensor.extract %[[SI]]{{\[}}%[[FI]], %[[C1]]] : tensor<2x2xi64>
+//       CHECK-RWT:    %[[I1:.*]] = arith.index_cast %[[I1r]] : i64 to index
+//       CHECK-RWT:    %[[V:.*]] = tensor.extract %[[SV]]{{\[}}%[[FI]]] : tensor<2xf32>
+//       CHECK-RWT:    sparse_tensor.insert %[[V]] into %[[COO]]{{\[}}%[[I0]], %[[I1]]]
+//       CHECK-RWT:  }
+//       CHECK-RWT:  %[[TI0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
+//       CHECK-RWT:  %[[TI1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
+//       CHECK-RWT:  %[[NNZ:.*]] = memref.load %[[TI0]]{{\[}}%[[C1]]] : memref<?xindex>
+//       CHECK-RWT:  %[[TV:.*]] = sparse_tensor.values %[[COO]]
+//       CHECK-RWT:  sparse_tensor.sort %[[NNZ]], %[[TI0]], %[[TI1]] jointly %[[TV]]
+//       CHECK-RWT:  %[[DST:.*]] = bufferization.alloc_tensor()
+//       CHECK-RWT:  sparse_tensor.foreach in %[[COO]]
+//       CHECK-RWT:  ^bb0(%[[F2I0:.*]]: index, %[[F2I1:.*]]: index, %[[F2V:.*]]: f32):
+//       CHECK-RWT:    sparse_tensor.insert %[[F2V]] into %[[DST]]{{\[}}%[[F2I0]], %[[F2I1]]]
+//       CHECK-RWT:  }
+//       CHECK-RWT:  %[[R:.*]] = sparse_tensor.convert %[[DST]]
+//       CHECK-RWT:  bufferization.dealloc_tensor %[[COO]]
+//       CHECK-RWT:  return %[[R]] : tensor<8x7xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>
 func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{
   // Initialize a tensor.
   %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
   // Convert the tensor to a sparse tensor.
   %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>
   return %1 : tensor<8x7xf32, #CSR>
 }
 

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

-// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
+// RxUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
+
+// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>

 //       CHECK:   %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
 //       CHECK:   %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
 //       CHECK:   %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
 //       CHECK:   memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<2x4xf64>
 //       CHECK:   scf.yield
 //       CHECK: }
 //       CHECK: %[[T:.*]] = bufferization.to_tensor %[[M]] : memref<2x4xf64>
 //       CHECK: return %[[T]] : tensor<2x4xf64>
+
+// CHECK-RWT-LABEL: func.func @sparse_convert_2d(
+//  CHECK-RWT-SAME: %[[A:.*]]: tensor<2x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>) -> tensor<2x4xf64> {
+//       CHECK-RWT: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
+//       CHECK-RWT: %[[B:.*]] = memref.alloc() : memref<2x4xf64>
+//       CHECK-RWT: linalg.fill ins(%[[F0]] : f64) outs(%[[B]]
+//       CHECK-RWT: sparse_tensor.foreach in %[[A]]
+//       CHECK-RWT: ^bb0(%[[FI0:.*]]: index, %[[FI1:.*]]: index, %[[FV:.*]]: f64):
+//       CHECK-RWT:   memref.store %[[FV]], %[[B]]{{\[}}%[[FI0]], %[[FI1]]]
+//       CHECK-RWT: }
+//       CHECK-RWT: %[[T:.*]] = bufferization.to_tensor %[[B]]
+//       CHECK-RWT: return %[[T]] : tensor<2x4xf64>
 func.func @sparse_convert_2d(%arg0: tensor<2x4xf64, #SparseMatrix>) -> tensor<2x4xf64> {
   %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64, #SparseMatrix> to tensor<2x4xf64>
   return %0 : tensor<2x4xf64>
 }
 
 // CHECK-LABEL: func @sparse_convert_2d_dyn0(
 //  CHECK-SAME: %[[Arg:.*]]: !llvm.ptr<i8>) -> tensor<?x4xf64>
 //   CHECK-DAG: %[[I0:.*]] = arith.constant 0 : index

 //       CHECK:   %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
 //       CHECK:   %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
 //       CHECK:   %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
 //       CHECK:   memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<?x?xf64>
 //       CHECK:   scf.yield
 //       CHECK: }
 //       CHECK: %[[T:.*]] = bufferization.to_tensor %[[M]] : memref<?x?xf64>
 //       CHECK: return %[[T]] : tensor<?x?xf64>
+
+// CHECK-RWT-LABEL: func.func @sparse_convert_2d_dyn2(
+//  CHECK-RWT-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>) -> tensor<?x?xf64> {
+//   CHECK-RWT-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-RWT-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-RWT-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
+//       CHECK-RWT: %[[D0:.*]] = tensor.dim %[[A]], %[[C0]]
+//       CHECK-RWT: %[[D1:.*]] = tensor.dim %[[A]], %[[C1]]
+//       CHECK-RWT: %[[B:.*]] = memref.alloc(%[[D0]], %[[D1]])
+//       CHECK-RWT: linalg.fill ins(%[[F0]] : f64) outs(%[[B]]
+//       CHECK-RWT: sparse_tensor.foreach in %[[A]]
+//       CHECK-RWT: ^bb0(%[[FI0:.*]]: index, %[[FI1:.*]]: index, %[[FV:.*]]: f64):
+//       CHECK-RWT:   memref.store %[[FV]], %[[B]]{{\[}}%[[FI0]], %[[FI1]]]
+//       CHECK-RWT: }
+//       CHECK-RWT: %[[T:.*]] = bufferization.to_tensor %[[B]]
+//       CHECK-RWT: return %[[T]] : tensor<?x?xf64>
 func.func @sparse_convert_2d_dyn2(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x?xf64> {
   %0 = sparse_tensor.convert %arg0 : tensor<?x?xf64, #SparseMatrix> to tensor<?x?xf64>
   return %0 : tensor<?x?xf64>
 }
 
 // CHECK-LABEL: func @sparse_convert_3d(
 //  CHECK-SAME: %[[Arg:.*]]: !llvm.ptr<i8>) -> tensor<2x3x4xf64>
 //   CHECK-DAG: %[[I0:.*]] = arith.constant 0 : index

mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir

 // First use with `kViaCOO` for sparse2sparse conversion (the old way).
 // RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=1" \
 // RUN:    --canonicalize --cse | FileCheck %s -check-prefix=CHECK-COO
 //
 // Now again with `kAuto` (the new default).
 // RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=0" \
 // RUN:    --canonicalize --cse | FileCheck %s -check-prefixes=CHECK-AUTO,CHECK
 
+// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
+
 #SparseVector64 = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"],
   pointerBitWidth = 64,
   indexBitWidth = 64
 }>
 
 #SparseVector32 = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"],

 //   CHECK-AUTO-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8>
 //   CHECK-AUTO-DAG: %[[Q:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-AUTO-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-AUTO-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8>
 //   CHECK-AUTO-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex>
 //   CHECK-AUTO-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex>
 //       CHECK-AUTO: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[SparseToSparse]], %[[A]])
 //       CHECK-AUTO: return %[[T]] : !llvm.ptr<i8>
+
+// CHECK-RWT-LABEL: func.func @sparse_convert(
+//  CHECK-RWT-SAME: %[[A:.*]]: tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 64, indexBitWidth = 64 }>>)
+//  CHECK-RWT-DAG:  %[[C0:.*]] = arith.constant 0 : index
+//  CHECK-RWT-DAG:  %[[C1:.*]] = arith.constant 1 : index
+//      CHECK-RWT:  %[[D:.*]] = tensor.dim %[[A]], %[[C0]]
+//      CHECK-RWT:  %[[COO:.*]] = bufferization.alloc_tensor(%[[D]])
+//      CHECK-RWT:  sparse_tensor.foreach in %[[A]]
+//      CHECK-RWT:  ^bb0(%[[FI:.*]]: index, %[[FV:.*]]: f32):
+//      CHECK-RWT:    sparse_tensor.insert %[[FV]] into %[[COO]]{{\[}}%[[FI]]]
+//      CHECK-RWT:  }
+//      CHECK-RWT:  %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
+//      CHECK-RWT:  %[[NNZr:.*]] = memref.load %[[I0]]{{\[}}%[[C1]]] : memref<?xi64>
+//      CHECK-RWT:  %[[NNZ:.*]] = arith.index_cast %[[NNZr]] : i64 to index
+//      CHECK-RWT:  %[[V:.*]] = sparse_tensor.values %[[COO]]
+//      CHECK-RWT:  sparse_tensor.sort %[[NNZ]], %[[I0]] jointly %[[V]]
+//      CHECK-RWT:  %[[DST:.*]] = bufferization.alloc_tensor(%[[D]])
+//      HECK-RWT:  sparse_tensor.foreach in %[[COO]]
+//      CHECK-RWT:  ^bb0(%[[FI2:.*]]: index, %[[FV2:.*]]: f32):
+//      CHECK-RWT:    sparse_tensor.insert %[[FV2]] into %[[DST]]{{\[}}%[[FI2]]]
+//      CHECK-RWT:  }
+//      CHECK-RWT:  bufferization.dealloc_tensor %[[COO]]
+//      CHECK-RWT:  %[[R:.*]] = sparse_tensor.convert %[[DST]]
+//      CHECK-RWT:  return %[[R]] : tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 32, indexBitWidth = 32 }>>
 func.func @sparse_convert(%arg0: tensor<?xf32, #SparseVector64>) -> tensor<?xf32, #SparseVector32> {
   %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector64> to tensor<?xf32, #SparseVector32>
   return %0 : tensor<?xf32, #SparseVector32>
 }
 
 #SparseSingleton64 = #sparse_tensor.encoding<{
   dimLevelType = ["singleton"],
   pointerBitWidth = 64,

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

-// RUN: mlir-opt %s -sparse-tensor-rewrite=enable-runtime-library=false  | FileCheck %s
+// RUN: mlir-opt %s -sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" |\
+// RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
 // CHECK-LABEL:   func.func @sparse_new(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
 // CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index

mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir

-// RUN: mlir-opt %s --sparse-tensor-rewrite=enable-runtime-library=false --sparsification | FileCheck %s
+// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" \
+// RUN: --sparsification | FileCheck %s
 
 #DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 
 // CHECK-LABEL: @concat_sparse_sparse(
 //  CHECK-SAME:  %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor
 //  CHECK-SAME:  %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor
 //  CHECK-SAME:  %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor
 //       CHECK:  %[[TMP_c0:.*]] = arith.constant 0 : index

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
+// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" \
+// RUN: --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(