[mlir][sparse] Support sparse2sparse collapse for dynamic sizes

[anlunx]

Summary

This patch implements sparse2sparse collapse for operands with dynamic shape.

[aartbik]

please use tags in the title, our team usually usesSparse reshaping for dynamic sizes

[mlir][sparse] sparse reshaping for dynamic sizes

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

525

we already have hasStaticShape() as a utility

[aartbik]

The "Index is too large for the dimension"' failure indicates that you are not passing the right (runtime) sizes for the COO data structure. I suspect you pass the -1 of the dynamic size to the kEmptyCOO call.

[anlunx]

use hasStaticShape

[anlunx]

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

626–636

In D131599#3713509, @aartbik wrote:

The "Index is too large for the dimension"' failure indicates that you are not passing the right (runtime) sizes for the COO data structure. I suspect you pass the -1 of the dynamic size to the kEmptyCOO call.

I'm using sizesFromPtr to compute the size of the destination tensor. Is it a valid thing to do? I was thinking that it should compute the runtime size because it generates the "sparseDimSize" function.

[aartbik]

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

527

This is really very similar to translateIndices, and contains way too much codedup. Can we merge the two and deal with static/dynamic in the logic?

619

llvm style does not want underscores in local vars

626–636

The problem is that you are using

sizesFromPtr(rewriter, dst_sizes, op, encDst, dstTp, src);

to fill the dimensions of the *destination*, but are you are querying the *src* to get the dimension size dynamically!
That way, you set the size of the destination to tensor<1xf64> and the second insertion goes out of bounds.

The problem, of course, is that you have no dst operand to query the sizes from ;-)
You will have to compute the destination format dynamically, using the dynamic sizes of the src to build it.
In this case you will need a computation that does dst.size[0] = src.size[0] * src.size[1]
for lower to higher, the computation will be more elaborate!

[anlunx]

Add function for generate collapse destination shape
Merge translateIndices with translateIndicesDyn

[anlunx]

Rename genDstShape to genReshapeDstShape

[anlunx]

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

527

This is done. Thanks for the suggestion!

626–636

Thank you for the explanation! Added a function genReshapeDstShape to compute the dynamic size of the destination.

[anlunx]

Fix style using clang-format

[aartbik]

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

531

Can we make expand of this revision too?

If not, then don't assert, but simply return failure() in the rewriter, i.e. pattern will not kick in.

[Peiming]

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

507

Is it more easy to follow to always generate the MLIR operation (same below) regardless of whether the size is statically know? We can rely on constant folding to eliminate unnecessary runtime computation.

I am not sure whether it is worth it though :-)

[anlunx]

Implement expand

[anlunx]

Fix typo

[anlunx]

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

507

Good point. Removed code for static shapes.

531

Added expand

[anlunx]

Fix sparse_reshape round-trip test

[aartbik]

apologies for the long review time this time around....

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

481–482

you have to use ArrayRef<> in the parameter; SmallVector<> is only used at the concrete implementation side (and as written, you copy contents into the call!)

488

same, when assigning, only work on references

493

perhaps good to keep this assert since you use the literal

507

Can you check that. In the CHECK test below, we just use the unrelated --cse as extra parameter, and you see that constants are not folded. It is okay to add more flags there, but I would like to see evidence that the code folds indeed.

528

document the "output" parameter, use ArrayRef for all others

538

assert on not dynamic?

550

I typically use the idiom

for (unsigned i = 0, size = srcShape.size(); i < size; i++)

to avoid the potential call in loop (depending on how smart compiler is)

558

Is this skip over intended? Your code can definitely use some examples for each case

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

27

see above, okay to add an extra flag, but I would like to keep what is on the left

66

as here

83–170

why did you not add any CHECK tests for the dynamic case

you only rely on integration test now, but I would like to see a complex example generating the right code.

[Peiming]

This comment has been deleted.

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

628–635

srcTp or dstTp here?

[anlunx]

Address comments

[anlunx]

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

493

I think I'm not using compile-time constants here. The shape vector here is computed at runtime, so we don't need this assert anymore.

507

Added the --canonicalize flag which folds constants.

538

IIUC srcShape is available no matter if the source shape is dynamic or static? It is computed by calling sizesFromPtr

558

Yes, it is intended. For example, if the srcShape is [8], and the staticDstShape is [2 x ? x 2], then we compute the unknown dimension by dividing 8 by 2 x 2. Therefore we skip the unknown dimension when computing the denominator. Added detailed comments for the algorithm.

628–635

I think they are equivalent since dstTp has static shape if and only if srcTp has static shape. But I agree that dstTp makes more sense here. Changed it to dstTp.

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

83–170

Added codegen tests for dynamic reshape

[anlunx]

fix format

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

25

hardcoding %21 is way too brittle (and granted, it was wrong in the original test

this should be either a captured variable, as in

%[[C:.*]] = ....
scf.condition(%[[C]])

or simply

scf.condition

64

same

84

extra above
roundtrip:

[anlunx]

Fix test format

[anlunx]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

25

Removed hardcoded variable

[aartbik]

one last nit, but good to go

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

493

Yeah, but since you use the shape[j], I would like to make sure we never multiply it with -1 ;-)

I realize we just constructed it, but defensive asserts like this ensure that if later somebody changes that logic to allow for dynamic, we will not forget about this one.

[anlunx]

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

493

I still don't understand why the assert is necessary. This change already allows reshaping with dynamic sizes, so the shape variable here should contain actual dynamic sizes, and can never contain -1?

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

static bool canUseDirectConversion(
   return true;
 }
 
 /// Helper method to translate indices during a reshaping operation.
 /// TODO: provide as general utility to MLIR at large?
 static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
                              ArrayRef<ReassociationIndices> reassociation,
                              TensorType dstTp, TensorType srcTp, Value dstIdx,
-                             Value srcIdx) {
+                             Value srcIdx, ArrayRef<Value> dstShape,
+                             ArrayRef<Value> srcShape) {
   unsigned dstRank = dstTp.getRank();
   unsigned srcRank = srcTp.getRank();
   unsigned start = 0;
   unsigned i = 0;
   bool isExpand = srcRank > dstRank;
-  ArrayRef<int64_t> shape = isExpand ? srcTp.getShape() : dstTp.getShape();
+  ArrayRef<Value> shape = isExpand ? srcShape : dstShape;
   // Iterate over reassociation map.
   for (const auto &map : llvm::enumerate(reassociation)) {
     // Prepare strides information in dimension slice.
-    uint64_t linear = 1;
+    Value linear = constantIndex(rewriter, loc, 1);
     for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
-      assert(!ShapedType::isDynamic(shape[j]));
-      linear *= shape[j];
+      linear = rewriter.create<arith::MulIOp>(loc, linear, shape[j]);
     }
     // Start collapse.
     Value idx = constantIndex(rewriter, loc, i++);
     Value val;
     if (!isExpand)
       val = rewriter.create<memref::LoadOp>(loc, srcIdx, idx);
     // Iterate over dimension slice.
     for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
-      linear /= shape[j];
-      Value stride = constantIndex(rewriter, loc, linear);
+      linear = rewriter.create<arith::DivUIOp>(loc, linear, shape[j]);
       Value jdx = constantIndex(rewriter, loc, j);
       if (isExpand) {
         Value old = rewriter.create<memref::LoadOp>(loc, srcIdx, jdx);
-        Value mul = linear == 1
-                        ? old
-                        : rewriter.create<arith::MulIOp>(loc, old, stride);
+        Value mul = rewriter.create<arith::MulIOp>(loc, old, linear);
         val = val ? rewriter.create<arith::AddIOp>(loc, val, mul) : mul;
       } else {
         Value old = val;
-        if (linear != 1)
-          val = rewriter.create<arith::DivUIOp>(loc, val, stride);
+        val = rewriter.create<arith::DivUIOp>(loc, val, linear);
         rewriter.create<memref::StoreOp>(loc, val, dstIdx, jdx);
-        if (linear != 1)
-          val = rewriter.create<arith::RemUIOp>(loc, old, stride);
+        val = rewriter.create<arith::RemUIOp>(loc, old, linear);
       }
     }
     // Finalize expansion.
     if (isExpand)
       rewriter.create<memref::StoreOp>(loc, val, dstIdx, idx);
     start += map.value().size();
   }
   // Sanity.
   assert((isExpand && i == dstRank) || (!isExpand && i == srcRank));
 }
 
+/// Helper method to compute the shape of destination tensor of a reshape
+/// operator. This is only used when operands have dynamic shape. The shape of
+/// the destination is stored into dstShape.
+void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
+                        SmallVector<Value, 4> &dstShape,
+                        ArrayRef<Value> srcShape,
+                        ArrayRef<int64_t> staticDstShape,
+                        ArrayRef<ReassociationIndices> reassociation) {
+  // Collapse shape.
+  if (reassociation.size() < srcShape.size()) {
+    unsigned start = 0;
+    for (const auto &map : llvm::enumerate(reassociation)) {
+      auto dstDim = constantIndex(rewriter, loc, 1);
+      for (unsigned i = start; i < start + map.value().size(); i++) {
+        dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
+      }
+      dstShape.push_back(dstDim);
+      start = start + map.value().size();
+    }
+    assert(start == srcShape.size());
+    return;
+  }
+
+  // Expand shape.
+  assert(reassociation.size() == srcShape.size());
+  unsigned start = 0;
+  // Expand the i-th dimension in srcShape.
+  for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
+    auto map = reassociation[i];
+    auto srcDim = srcShape[i];
+    // Iterate through dimensions expanded from the i-th dimension.
+    for (unsigned j = start; j < start + map.size(); j++) {
+      // There can be only one dynamic sized dimension among dimensions expanded
+      // from the i-th dimension in srcShape. For example, if srcDim = 8, then
+      // the expanded shape could be <2x?x2>, but not <2x?x?>.
+      if (staticDstShape[j] == ShapedType::kDynamicSize) {
+        // The expanded dimension has dynamic size. We compute the dimension
+        // by dividing srcDim by the product of the static dimensions.
+        int64_t product = 1;
+        for (unsigned k = start; k < start + map.size(); k++) {
+          if (staticDstShape[k] != ShapedType::kDynamicSize) {
+            product *= staticDstShape[k];
+          }
+        }
+        // Compute the dynamic dimension size.
+        Value productVal = constantIndex(rewriter, loc, product);
+        Value dynamicSize =
+            rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
+        dstShape.push_back(dynamicSize);
+      } else {
+        // The expanded dimension is statically known.
+        dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
+      }
+    }
+    start = start + map.size();
+  }
+  assert(start == staticDstShape.size());
+}
+
 /// Generate code for a general sparse to sparse reshaping operation.
 /// Note that unlike dense reshaping (which can be done with a "cheap"
 /// change of view), sparse reshaping is currently done with actual
 /// data shuffling.
 ///
 /// TODO: proportional to nnz, but still a lot of data movement
 ///       https://github.com/llvm/llvm-project/issues/56477
 ///
genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
   unsigned dstRank = dstTp.getRank();
   Type elemTp = srcTp.getElementType();
   assert(elemTp == dstTp.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
       op->getContext(), encSrc.getDimLevelType(), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-  SmallVector<Value, 4> sizes;
+  SmallVector<Value, 4> srcSizes;
   SmallVector<Value, 8> params;
-  sizesFromPtr(rewriter, sizes, loc, encSrc, srcTp, adaptor.getSrc());
-  newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, sizes,
+  sizesFromPtr(rewriter, srcSizes, loc, encSrc, srcTp, adaptor.getSrc());
+  newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, srcSizes,
             adaptor.getSrc());
   Value iter = genNewCall(rewriter, loc, params);
   // Start a new COO for the destination tensor.
-  sizes.clear();
+  SmallVector<Value, 4> dstSizes;
   params.clear();
-  // Fills sizes array using the sizes from destination type.
-  assert(dstTp.hasStaticShape());
-  sizesFromType(rewriter, sizes, loc, dstTp);
-  newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, sizes);
+  if (dstTp.hasStaticShape()) {
+    sizesFromType(rewriter, dstSizes, loc, dstTp);
+  } else {
+    ArrayRef<int64_t> dstShape = dstTp.getShape();
+    genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
+                       op.getReassociationIndices());
+  }
+  newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, dstSizes);
   Value coo = genNewCall(rewriter, loc, params);
   Value dstPerm = params[2];
   // Construct a while loop over the iterator.
   Value srcIdx = genAlloca(rewriter, loc, srcRank, rewriter.getIndexType());
   Value dstIdx = genAlloca(rewriter, loc, dstRank, rewriter.getIndexType());
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
   Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
   rewriter.setInsertionPointToEnd(before);
   Value cond = genGetNextCall(rewriter, loc, iter, srcIdx, elemPtr);
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   // Translate indices from source to target and insert. Note that we do
   // not need to store the value in elemPtr, as the value is still there.
   Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
   rewriter.setInsertionPointToStart(after);
   translateIndices(loc, rewriter, op.getReassociationIndices(), dstTp, srcTp,
-                   dstIdx, srcIdx);
+                   dstIdx, srcIdx, dstSizes, srcSizes);
   genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
   rewriter.create<scf::YieldOp>(loc);
   // Final call to construct sparse tensor storage and free temporary resources.
   rewriter.setInsertionPointAfter(whileOp);
   params[6] = constantAction(rewriter, loc, Action::kFromCOO);
   params[7] = coo;
   Value dst = genNewCall(rewriter, loc, params);
   genDelCOOCall(rewriter, loc, elemTp, coo);

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 | FileCheck %s --check-prefix=CHECK-CONV
+// RUN: mlir-opt %s --sparse-tensor-conversion --cse --canonicalize | FileCheck %s --check-prefix=CHECK-CONV
 
 #SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 #SparseMatrix = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 //
 // roundtrip:
 //
 // CHECK-ROUND-LABEL: func.func @sparse_expand(
 // CHECK-ROUND-SAME:  %[[A:.*]]: tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
 //      CHECK-ROUND:  %[[E:.*]] = tensor.expand_shape %[[A]] {{\[\[}}0, 1]] : tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
 //      CHECK-ROUND:  return %[[E]] : tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
 //
 // conversion:
 //
 // CHECK-CONV-LABEL: func.func @sparse_expand(
 // CHECK-CONV-DAG:  %[[C0:.*]] = arith.constant 0 : index
 // CHECK-CONV-DAG:  %[[C1:.*]] = arith.constant 1 : index
 // CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
-// CHECK-CONV:        scf.condition(%13)
+// CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
-// CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<?xindex>
+// CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        %[[D:.*]] = arith.divui %[[X]], %[[C10]] : index
-// CHECK-CONV:        memref.store %[[D]], %{{.*}}[%[[C0]]] : memref<?xindex>
+// CHECK-CONV:        memref.store %[[D]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        %[[R:.*]] = arith.remui %[[X]], %[[C10]] : index
-// CHECK-CONV:        memref.store %[[R]], %{{.*}}[%[[C1]]] : memref<?xindex>
+// CHECK-CONV:        memref.store %[[R]], %{{.*}}[%[[C1]]] : memref<2xindex>
 // 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>
 //

 // CHECK-CONV-LABEL: func.func @sparse_collapse(
 // CHECK-CONV-DAG:  %[[C0:.*]] = arith.constant 0 : index
 // CHECK-CONV-DAG:  %[[C1:.*]] = arith.constant 1 : index
 // CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV-DAG:  call @newSparseTensor
 // CHECK-CONV:      scf.while : () -> () {
 // CHECK-CONV:        call @getNextF64
-// CHECK-CONV:        scf.condition(%13)
+// CHECK-CONV:        scf.condition
 // CHECK-CONV:      } do {
-// CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<?xindex>
+// CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        %[[M:.*]] = arith.muli %[[X]], %[[C10]] : index
-// CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<?xindex>
+// CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        %[[A:.*]] = arith.addi %[[M]], %[[Y]] : index
-// CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<?xindex>
+// CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
 // 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>
 //
 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-ROUND-LABEL: func.func @dynamic_sparse_expand(
+// CHECK-ROUND-SAME:  %[[A:.*]]: tensor<?xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<?x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//      CHECK-ROUND:  %[[E:.*]] = tensor.expand_shape %[[A]] {{\[\[}}0, 1]] : tensor<?xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<?x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//      CHECK-ROUND:  return %[[E]] : tensor<?x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//
+// conversion:
+//
+// CHECK-CONV-LABEL: func.func @dynamic_sparse_expand(
+// CHECK-CONV-DAG:  %[[C0:.*]] = arith.constant 0 : index
+// CHECK-CONV-DAG:  %[[C1:.*]] = arith.constant 1 : index
+// CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
+// CHECK-CONV-DAG:  %[[D1:.*]] = arith.divui %{{.*}}, %[[C10]] : index
+// CHECK-CONV-DAG:  call @newSparseTensor
+// CHECK-CONV-DAG:  call @newSparseTensor
+// CHECK-CONV:      scf.while : () -> () {
+// CHECK-CONV:        call @getNextF64
+// CHECK-CONV:        scf.condition
+// CHECK-CONV:      } do {
+// CHECK-CONV:        %[[M:.*]] = arith.muli %[[D1]], %[[C10]] : index
+// CHECK-CONV:        %[[L:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<1xindex>
+// CHECK-CONV:        %[[D2:.*]] = arith.divui %[[M]], %[[D1]] : index
+// CHECK-CONV:        %[[D3:.*]] = arith.divui %[[L]], %[[D2]] : index
+// CHECK-CONV:        memref.store %[[D3]], %{{.*}}[%[[C0]]] : memref<2xindex>
+// CHECK-CONV:        %[[R:.*]] = arith.remui %[[L]], %[[D2]] : index
+// CHECK-CONV:        %[[D4:.*]] = arith.divui %[[D2]], %[[C10]] : index
+// CHECK-CONV:        %[[D5:.*]] = arith.divui %[[R]], %[[D4]] : index
+// CHECK-CONV:        memref.store %[[D5]], %{{.*}}[%[[C1]]] : memref<2xindex>
+// 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>
+//
+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-ROUND-LABEL: func.func @dynamic_sparse_collapse(
+// CHECK-ROUND-SAME:  %[[A:.*]]: tensor<10x?xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<?xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//      CHECK-ROUND:  %[[C:.*]] = tensor.collapse_shape %[[A]] {{\[\[}}0, 1]] : tensor<10x?xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<?xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//      CHECK-ROUND:  return %[[C]] : tensor<?xf64, #sparse_tensor.encoding<{{{.*}}}>>
+//
+// conversion:
+//
+// CHECK-CONV-LABEL: func.func @dynamic_sparse_collapse(
+// CHECK-CONV-DAG:  %[[C0:.*]] = arith.constant 0 : index
+// CHECK-CONV-DAG:  %[[C1:.*]] = arith.constant 1 : index
+// CHECK-CONV-DAG:  %[[C10:.*]] = arith.constant 10 : index
+// CHECK-CONV-DAG:  %[[M1:.*]] = arith.muli %{{.*}}, %[[C10]] : index
+// CHECK-CONV-DAG:  call @newSparseTensor
+// CHECK-CONV-DAG:  call @newSparseTensor
+// CHECK-CONV:      scf.while : () -> () {
+// CHECK-CONV:        call @getNextF64
+// CHECK-CONV:        scf.condition
+// CHECK-CONV:      } do {
+// CHECK-CONV:        %[[D1:.*]] = arith.divui %[[M1]], %[[C10]] : index
+// CHECK-CONV:        %[[X:.*]] = memref.load %{{.*}}[%[[C0]]] : memref<2xindex>
+// CHECK-CONV:        %[[M2:.*]] = arith.muli %[[X]], %[[D1]] : index
+// CHECK-CONV:        %[[D2:.*]] = arith.divui %[[D1]], %{{.*}} : index
+// CHECK-CONV:        %[[Y:.*]] = memref.load %{{.*}}[%[[C1]]] : memref<2xindex>
+// CHECK-CONV:        %[[M3:.*]] = arith.muli %[[Y]], %[[D2]] : index
+// CHECK-CONV:        %[[A:.*]] = arith.addi %[[M2]], %[[M3]] : index
+// CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
+// 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>
+//
+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>
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir

func.func @expand_from_sparse_dyn(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64> {
     return %0 : tensor<?x2x?xf64>
   }
 
   func.func @expand_to_sparse_dyn(%arg0: tensor<?x?xf64>) -> tensor<?x2x?xf64, #Sparse3dTensor> {
     %0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64> into tensor<?x2x?xf64, #Sparse3dTensor>
     return %0 : tensor<?x2x?xf64, #Sparse3dTensor>
   }
 
+  func.func @expand_sparse2sparse_dyn(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64, #Sparse3dTensor> {
+    %0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64, #SparseMatrix> into tensor<?x2x?xf64, #Sparse3dTensor>
+    return %0 : tensor<?x2x?xf64, #Sparse3dTensor>
+  }
+
   func.func @collapse_dense_dyn(%arg0: tensor<?x?x?x?xf64>) -> tensor<?x?xf64> {
     %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?x?x?xf64> into tensor<?x?xf64>
     return %0 : tensor<?x?xf64>
   }
 
   func.func @collapse_from_sparse_dyn(%arg0: tensor<?x?x?x?xf64, #Sparse4dTensor>) -> tensor<?x?xf64> {
     %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?x?x?xf64, #Sparse4dTensor> into tensor<?x?xf64>
     return %0 : tensor<?x?xf64>
   }
 
   func.func @collapse_to_sparse_dyn(%arg0: tensor<?x?x?x?xf64>) -> tensor<?x?xf64, #SparseMatrix> {
     %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?x?x?xf64> into tensor<?x?xf64, #SparseMatrix>
     return %0 : tensor<?x?xf64, #SparseMatrix>
   }
 
+  func.func @collapse_sparse2sparse_dyn(%arg0: tensor<?x?x?x?xf64, #Sparse4dTensor>) -> tensor<?x?xf64, #SparseMatrix> {
+    %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?x?x?xf64, #Sparse4dTensor> into tensor<?x?xf64, #SparseMatrix>
+    return %0 : tensor<?x?xf64, #SparseMatrix>
+  }
+
   //
   // Main driver.
   //
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %df = arith.constant -1.0 : f64
 
     // Setup test vectors and matrices..
func.func @entry() {
     %expand3 = call @expand_sparse2sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix>
     %expand4 = call @expand_dense_3x2x2(%m) : (tensor<3x4xf64>) -> tensor<3x2x2xf64>
     %expand5 = call @expand_from_sparse_3x2x2(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<3x2x2xf64>
     %expand6 = call @expand_to_sparse_3x2x2(%m) : (tensor<3x4xf64>) -> tensor<3x2x2xf64, #Sparse3dTensor>
     %expand7 = call @expand_sparse2sparse_3x2x2(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<3x2x2xf64, #Sparse3dTensor>
     %expand8 = call @expand_dense_dyn(%dm) : (tensor<?x?xf64>) -> tensor<?x2x?xf64>
     %expand9 = call @expand_from_sparse_dyn(%sdm) : (tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64>
     %expand10 = call @expand_to_sparse_dyn(%dm) : (tensor<?x?xf64>) -> tensor<?x2x?xf64, #Sparse3dTensor>
+    %expand11 = call @expand_sparse2sparse_dyn(%sdm) : (tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64, #Sparse3dTensor>
 
     %collapse0 = call @collapse_dense(%m) : (tensor<3x4xf64>) -> tensor<12xf64>
     %collapse1 = call @collapse_from_sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64>
     %collapse2 = call @collapse_to_sparse(%m) : (tensor<3x4xf64>) -> tensor<12xf64, #SparseVector>
     %collapse3 = call @collapse_sparse2sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector>
     %collapse4 = call @collapse_dense_6x10(%n) : (tensor<2x3x5x2xf64>) -> tensor<6x10xf64>
     %collapse5 = call @collapse_from_sparse_6x10(%sn) : (tensor<2x3x5x2xf64, #Sparse4dTensor>) -> tensor<6x10xf64>
     %collapse6 = call @collapse_to_sparse_6x10(%n) : (tensor<2x3x5x2xf64>) -> tensor<6x10xf64, #SparseMatrix>
     %collapse7 = call @collapse_sparse2sparse_6x10(%sn) : (tensor<2x3x5x2xf64, #Sparse4dTensor>) -> tensor<6x10xf64, #SparseMatrix>
     %collapse8 = call @collapse_dense_dyn(%dn) : (tensor<?x?x?x?xf64>) -> tensor<?x?xf64>
     %collapse9 = call @collapse_from_sparse_dyn(%sdn) : (tensor<?x?x?x?xf64, #Sparse4dTensor>) -> tensor<?x?xf64>
     %collapse10 = call @collapse_to_sparse_dyn(%dn) : (tensor<?x?x?x?xf64>) -> tensor<?x?xf64, #SparseMatrix>
+    %collapse11 = call @collapse_sparse2sparse_dyn(%sdn) : (tensor<?x?x?x?xf64, #Sparse4dTensor>) -> tensor<?x?xf64, #SparseMatrix>
 
     //
     // Verify results of expand
     //
     // CHECK:      ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) )
     // CHECK-NEXT: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) )
     // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 )
     // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 )
     // CHECK-NEXT: ( ( ( 1.1, 1.2 ), ( 1.3, 1.4 ) ), ( ( 2.1, 2.2 ), ( 2.3, 2.4 ) ), ( ( 3.1, 3.2 ), ( 3.3, 3.4 ) ) )
     // CHECK-NEXT: ( ( ( 1.1, 1.2 ), ( 1.3, 1.4 ) ), ( ( 2.1, 2.2 ), ( 2.3, 2.4 ) ), ( ( 3.1, 3.2 ), ( 3.3, 3.4 ) ) )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
     // CHECK-NEXT: ( ( ( 1.1, 1.2 ), ( 1.3, 1.4 ) ), ( ( 2.1, 2.2 ), ( 2.3, 2.4 ) ), ( ( 3.1, 3.2 ), ( 3.3, 3.4 ) ) )
     // CHECK-NEXT: ( ( ( 1.1, 1.2 ), ( 1.3, 1.4 ) ), ( ( 2.1, 2.2 ), ( 2.3, 2.4 ) ), ( ( 3.1, 3.2 ), ( 3.3, 3.4 ) ) )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
     //
 
     %m0 = vector.transfer_read %expand0[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64>
     vector.print %m0 : vector<3x4xf64>
     %m1 = vector.transfer_read %expand1[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64>
     vector.print %m1 : vector<3x4xf64>
     %a2 = sparse_tensor.values %expand2 : tensor<3x4xf64, #SparseMatrix> to memref<?xf64>
     %m2 = vector.transfer_read %a2[%c0], %df: memref<?xf64>, vector<16xf64>
func.func @entry() {
 
     %m8 = vector.transfer_read %expand8[%c0, %c0, %c0], %df: tensor<?x2x?xf64>, vector<3x2x2xf64>
     vector.print %m8 : vector<3x2x2xf64>
     %m9 = vector.transfer_read %expand9[%c0, %c0, %c0], %df: tensor<?x2x?xf64>, vector<3x2x2xf64>
     vector.print %m9 : vector<3x2x2xf64>
     %a10 = sparse_tensor.values %expand10 : tensor<?x2x?xf64, #Sparse3dTensor> to memref<?xf64>
     %m10 = vector.transfer_read %a10[%c0], %df: memref<?xf64>, vector<16xf64>
     vector.print %m10 : vector<16xf64>
+    %a11 = sparse_tensor.values %expand11 : tensor<?x2x?xf64, #Sparse3dTensor> to memref<?xf64>
+    %m11 = vector.transfer_read %a11[%c0], %df: memref<?xf64>, vector<16xf64>
+    vector.print %m11 : vector<16xf64>
+
 
     // 
     // Verify results of collapse
     // 
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
     // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 )
     // CHECK-NEXT: ( ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ), ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ), ( 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ), ( 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 ), ( 41, 42, 43, 44, 45, 26, 47, 48, 49, 50 ), ( 51, 52, 53, 54, 55, 56, 57, 58, 59, 60 ) )
     // CHECK-NEXT: ( ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ), ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ), ( 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ), ( 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 ), ( 41, 42, 43, 44, 45, 26, 47, 48, 49, 50 ), ( 51, 52, 53, 54, 55, 56, 57, 58, 59, 60 ) )
     // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 26, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, -1, -1, -1, -1 )
     // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 26, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, -1, -1, -1, -1 )
     // CHECK-NEXT: ( ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ), ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ), ( 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ), ( 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 ), ( 41, 42, 43, 44, 45, 26, 47, 48, 49, 50 ), ( 51, 52, 53, 54, 55, 56, 57, 58, 59, 60 ) )
     // CHECK-NEXT: ( ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ), ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ), ( 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 ), ( 31, 32, 33, 34, 35, 36, 37, 38, 39, 40 ), ( 41, 42, 43, 44, 45, 26, 47, 48, 49, 50 ), ( 51, 52, 53, 54, 55, 56, 57, 58, 59, 60 ) )
     // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 26, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 26, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, -1, -1, -1, -1 )
     //
 
     %v0 = vector.transfer_read %collapse0[%c0], %df: tensor<12xf64>, vector<12xf64>
     vector.print %v0 : vector<12xf64>
     %v1 = vector.transfer_read %collapse1[%c0], %df: tensor<12xf64>, vector<12xf64>
     vector.print %v1 : vector<12xf64>
     %b2 = sparse_tensor.values %collapse2 : tensor<12xf64, #SparseVector> to memref<?xf64>
     %v2 = vector.transfer_read %b2[%c0], %df: memref<?xf64>, vector<16xf64>
func.func @entry() {
 
     %v8 = vector.transfer_read %collapse8[%c0, %c0], %df: tensor<?x?xf64>, vector<6x10xf64>
     vector.print %v8 : vector<6x10xf64>
     %v9 = vector.transfer_read %collapse9[%c0, %c0], %df: tensor<?x?xf64>, vector<6x10xf64>
     vector.print %v9 : vector<6x10xf64>
     %b10 = sparse_tensor.values %collapse10 : tensor<?x?xf64, #SparseMatrix> to memref<?xf64>
     %v10 = vector.transfer_read %b10[%c0], %df: memref<?xf64>, vector<64xf64>
     vector.print %v10 : vector<64xf64>
+    %b11 = sparse_tensor.values %collapse11 : tensor<?x?xf64, #SparseMatrix> to memref<?xf64>
+    %v11 = vector.transfer_read %b11[%c0], %df: memref<?xf64>, vector<64xf64>
+    vector.print %v11 : vector<64xf64>
+
 
     // Release sparse resources.
     bufferization.dealloc_tensor %sv : tensor<12xf64, #SparseVector>
     bufferization.dealloc_tensor %sm : tensor<3x4xf64, #SparseMatrix>
     bufferization.dealloc_tensor %sn : tensor<2x3x5x2xf64, #Sparse4dTensor>
     bufferization.dealloc_tensor %sdm : tensor<?x?xf64, #SparseMatrix>
     bufferization.dealloc_tensor %sdn : tensor<?x?x?x?xf64, #Sparse4dTensor>
     bufferization.dealloc_tensor %expand2 : tensor<3x4xf64, #SparseMatrix>
     bufferization.dealloc_tensor %expand3 : tensor<3x4xf64, #SparseMatrix>
     bufferization.dealloc_tensor %expand6 : tensor<3x2x2xf64, #Sparse3dTensor>
     bufferization.dealloc_tensor %expand7 : tensor<3x2x2xf64, #Sparse3dTensor>
     bufferization.dealloc_tensor %expand10 : tensor<?x2x?xf64, #Sparse3dTensor>
+    bufferization.dealloc_tensor %expand11 : tensor<?x2x?xf64, #Sparse3dTensor>
     bufferization.dealloc_tensor %collapse2 : tensor<12xf64, #SparseVector>
     bufferization.dealloc_tensor %collapse3 : tensor<12xf64, #SparseVector>
     bufferization.dealloc_tensor %collapse6 : tensor<6x10xf64, #SparseMatrix>
     bufferization.dealloc_tensor %collapse7 : tensor<6x10xf64, #SparseMatrix>
     bufferization.dealloc_tensor %collapse10 : tensor<?x?xf64, #SparseMatrix>
+    bufferization.dealloc_tensor %collapse11 : tensor<?x?xf64, #SparseMatrix>
 
     // Release dense resources.
     bufferization.dealloc_tensor %expand1 : tensor<3x4xf64>
     bufferization.dealloc_tensor %collapse1 : tensor<12xf64>
     bufferization.dealloc_tensor %expand5 : tensor<3x2x2xf64>
     bufferization.dealloc_tensor %collapse5 : tensor<6x10xf64>
     bufferization.dealloc_tensor %expand9 : tensor<?x2x?xf64>
     bufferization.dealloc_tensor %collapse9: tensor<?x?xf64>
 
     return
   }
 }