[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

514

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

603–613

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

516

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?

596

llvm style does not want underscores in local vars

603–613

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

516

This is done. Thanks for the suggestion!

603–613

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

520

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

496

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

496

Good point. Removed code for static shapes.

520

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

470–471

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!)

477

same, when assigning, only work on references

482

perhaps good to keep this assert since you use the literal

496

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.

517

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

527

assert on not dynamic?

539

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)

547

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

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

27 ↗

(On Diff #456805)

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

71 ↗

(On Diff #456805)

as here

93 ↗

(On Diff #456805)

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

605–612

srcTp or dstTp here?

[anlunx]

Address comments

[anlunx]

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

482

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.

496

Added the --canonicalize flag which folds constants.

527

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

547

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.

605–612

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

93 ↗

(On Diff #456805)

Added codegen tests for dynamic reshape

[anlunx]

fix format

[aartbik]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

25 ↗

(On Diff #460571)

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 ↗

(On Diff #460571)

same

84 ↗

(On Diff #460571)

extra above
roundtrip:

[anlunx]

Fix test format

[anlunx]

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

25 ↗

(On Diff #460571)

Removed hardcoded variable

[aartbik]

one last nit, but good to go

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

482

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

482

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, SmallVector<Value, 4> dstShape,
+                             SmallVector<Value, 4> 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();
+  SmallVector<Value, 4> 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.
+void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
+                        SmallVector<Value, 4> &dstShape,
+                        const SmallVector<Value, 4> &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;
+  for (unsigned i = 0; i < srcShape.size(); i++) {
+    auto map = reassociation[i];
+    auto srcDim = srcShape[i];
+    int64_t product = 1;
+    for (unsigned j = start; j < start + map.size(); j++) {
+      if (staticDstShape[j] == ShapedType::kDynamicSize) {
+        // Compute the product of the static sizes.
+        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 {
+        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(Operation *op, ConversionPatternRewriter &rewriter,
   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, noPerm, srcTp, src);
-  newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, sizes,
+  sizesFromPtr(rewriter, srcSizes, loc, noPerm, srcTp, src);
+  newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, srcSizes,
             src);
   Value iter = genNewCall(rewriter, loc, params);
   // Start a new COO for the destination tensor.
-  sizes.clear();
+  SmallVector<Value, 4> dstSizes;
   params.clear();
-  sizesFromPtr(rewriter, sizes, loc, encDst, dstTp, src);
-  newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, sizes);
+  if (srcTp.hasStaticShape()) {
+    sizesFromPtr(rewriter, dstSizes, loc, encDst, dstTp, src);
+  } else {
+    ArrayRef<int64_t> dstShape = dstTp.getShape();
+    genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
+                       reassociation);
+  }
+  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, reassociation, dstTp, srcTp, dstIdx, srcIdx);
+  translateIndices(loc, rewriter, reassociation, dstTp, srcTp, 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/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
   }
 }