[mlir][sparse] Add layout to the memref for the indices buffers to prepare for the AOS storage optimization for COO regions.

[bixia]

Summary

Fix relevant FileCheck tests.

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

76 ↗

(On Diff #485544)

Why does it have to be unique to apply AOS?

[bixia]

Move getCooStart to SparseTensor.h.

[bixia]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

76 ↗

(On Diff #485544)

This ensures the the last dim is a unique singleton.
What does it mean if the last dim is a non-unique singleton? If that is also a meaningful case we want to support, we can relax the requirement here.

[aartbik]

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

48

Empty line around each individual method (style in this file)

54

this feels a bit too long for linline, also move to cpp

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

276

we probably need to assert s vs. rank, just to make sure this does not return the wrong value when passing in the wrong value for s

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

919

now only 2 out of 3 getters use the SparseGetter abstraction? Is it still worth it. Shall we just move that back into the other rewriting rules too and make it more regular?

[bixia]

Move getCooStart to cpp. Remove SparseGetter. Add an assertion to check the input dim.

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

276

getDimLevelType has such an assert, but we add an assert here also.

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

919

remove SparseGetter then.

[Peiming]

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

141

Why not create a new builder for ToPointer/IndiceOp which infers the type?

It would be cleaner.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

76 ↗

(On Diff #485544)

It would become a BCOO

[bixia]

Rename getCooStart to getCOOStart to be consistent with other routines. Change the routine to not require unique last dim.

[bixia]

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

141

I will do this in another CL, to add builders to ToPointersOp, ToValuesOp, and ToIndicesOp that infer the result types.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

76 ↗

(On Diff #485544)

Remove the requirement for unique.

[bixia]

Add codegen utils for inferring the result types and generate ToValues/Indices/PointersOp.

[bixia]

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

141

I realized adding builder methods to infer result type will create a dependence from SparseTensorDialect.cpp to CodeGenUtils.cpp/h, which is what we want to avoid. Add a codegenutil instead.

[bixia]

Replace "tailing" with "trailing".

[aartbik]

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

279

This comment is now outdated since you have the isUnique override parameter. You should really say something like, if unique is requested then ....

288

by scanning only up to rank-1, you ensure there is always a compressed + one or more singleton

this is getting subtle, so really needs proper doc (not just in API, but here in implementation)

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

187

I would not say that you add a layout when set but

When withLayout is set, the returned layout has unknown strides and offsets. Otherwise,
the standard unit stride zero offset layout is returned.

or something like that

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

74–75

why not just complete the type? it does not vary right?

Leaving an < unclosed feels a bit wrong

[bixia]

Fix/add documentation. Complete the strided memref in FileCheck tests.

[bixia]

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

288

The document for the routine also mentions this, add doc here:

We only consider COO region with at least two dimensions for the purpose
of AOS storage optimization.

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

74–75

Fix here an other places.

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

 
 /// 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 with the last
 /// dimension level type being unique.
 bool isUniqueCOOType(RankedTensorType tp);
+/// Returns true iff the given sparse tensor encoding attribute has a tailing
+/// COO region starting at the given dimension.
+bool isUniqueCOOType(SparseTensorEncodingAttr enc, uint64_t s);
+/// Returns the starting dimension for a tailing COO region that spans across
+/// at least two dimensions. If no such COO region is found, returns the rank
+/// of the tensor.
+inline unsigned getCooStart(SparseTensorEncodingAttr enc) {
+  unsigned rank = enc.getDimLevelType().size();
+  if (rank <= 1)
+    return rank;
+
+  for (unsigned r = 0; r < rank - 1; r++) {
+    if (isUniqueCOOType(enc, r))
+      return r;
+  }
+
+  return rank;
+}
 
 //
 // 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.

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

if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
   if (auto mdtp = type.dyn_cast<StorageSpecifierType>())
     return mdtp.getEncoding();
   return nullptr;
 }
 
 bool mlir::sparse_tensor::isUniqueCOOType(RankedTensorType tp) {
   SparseTensorEncodingAttr enc = getSparseTensorEncoding(tp);
-
   if (!enc)
     return false;
 
-  if (!isCompressedDim(tp, 0))
+  return isUniqueCOOType(enc, 0);
+}
+
+bool mlir::sparse_tensor::isUniqueCOOType(SparseTensorEncodingAttr enc,
+                                          uint64_t s) {
+  if (!isCompressedDim(enc, s))
     return false;
 
-  for (uint64_t i = 1, e = tp.getRank(); i < e; ++i)
-    if (!isSingletonDim(tp, i))
+  uint64_t rank = enc.getDimLevelType().size();
+  for (uint64_t i = s + 1; i < rank; ++i)
+    if (!isSingletonDim(enc, i))
       return false;
 
   // This works for rank == 1 (unique the only compressed) and rank > 1 (unique
   // on the last singleton).
-  return isUniqueDim(tp, tp.getRank() - 1);
+  return isUniqueDLT(getDimLevelType(enc, rank - 1));
 }
 
 uint64_t mlir::sparse_tensor::toOrigDim(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/CodegenUtils.h

 void genDenseTensorOrSparseConstantIterLoop(
     OpBuilder &builder, Location loc, Value src, unsigned rank,
     function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder);
 
 /// Populates given sizes array from dense tensor or sparse tensor constant.
 void sizesFromSrc(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
                   Location loc, Value src);
 
+/// Generates a 1D MemRefType with a dynamic size. When withLayout is set, add
+/// a layout with unknown strides and offset to the type.
+inline MemRefType get1DMemRefType(Type etp, bool withLayout) {
+  auto layout = withLayout ? StridedLayoutAttr::StridedLayoutAttr::get(
+                                 etp.getContext(), ShapedType::kDynamic,
+                                 {ShapedType::kDynamic})
+                           : StridedLayoutAttr();
+  return MemRefType::get(ShapedType::kDynamic, etp, layout);
+}
+
 /// Scans to top of generated loop.
 Operation *getTop(Operation *op);
 
 /// Iterate over a sparse constant, generates constantOp for value and indices.
 /// E.g.,
 /// sparse<[ [0], [28], [31] ],
 ///          [ (-5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] >
 /// =>

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

for (size_t t = 0, e = tensors.size(); t < e; t++) {
     auto rtp = tensor.getType().dyn_cast<RankedTensorType>();
     if (!rtp)
       // Skips only scalar, zero ranked tensor still need to be bufferized and
       // (probably) filled with zeros by users.
       continue;
     auto rank = rtp.getRank();
     auto shape = rtp.getShape();
     auto enc = getSparseTensorEncoding(rtp);
-    auto dynShape = {ShapedType::kDynamic};
+    uint64_t cooStart = enc ? getCooStart(enc) : rank;
     // Scan all dimensions of current tensor.
     for (int64_t d = 0; d < rank; d++) {
       // This should be called only once at beginning.
       assert(!ptrBuffer[t][d] && !idxBuffer[t][d] && !highs[t][d]);
       // Handle sparse storage schemes.
       if (isCompressedDLT(dimTypes[t][d])) {
-        auto ptrTp =
-            MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
-        auto indTp =
-            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        auto ptrTp = get1DMemRefType(getPointerOverheadType(builder, enc),
+                                     /*withLayout=*/false);
+        auto indTp = get1DMemRefType(getIndexOverheadType(builder, enc),
+                                     /*withLayout=*/d >= cooStart);
         auto dim = builder.getIndexAttr(d);
         // Generate sparse primitives to obtains pointer and indices.
         ptrBuffer[t][d] = builder.create<ToPointersOp>(loc, ptrTp, tensor, dim);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
       } else if (isSingletonDLT(dimTypes[t][d])) {
         // Singleton dimension, fetch indices.
-        auto indTp =
-            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        auto indTp = get1DMemRefType(getIndexOverheadType(builder, enc),
+                                     /*withLayout=*/d >= cooStart);
         auto dim = builder.getIndexAttr(d);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
       } else {
         // Dense dimension, nothing to fetch.
         assert(isDenseDLT(dimTypes[t][d]));
       }
 
       // Find upper bound in current dimension.

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

public:
   static Value getFieldForOp(const SparseTensorDescriptor &desc,
                              ToPointersOp op) {
     uint64_t dim = op.getDimension().getZExtValue();
     return desc.getPtrMemRef(dim);
   }
 };
 
 /// Sparse codegen rule for index accesses.
-class SparseToIndicesConverter
-    : public SparseGetterOpConverter<ToIndicesOp, SparseToIndicesConverter> {
+class SparseToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
-  using SparseGetterOpConverter::SparseGetterOpConverter;
-  // Callback for SparseGetterOpConverter.
-  static Value getFieldForOp(const SparseTensorDescriptor &desc,
-                             ToIndicesOp op) {
+  using OpAdaptor = typename ToIndicesOp::Adaptor;
+  using OpConversionPattern<ToIndicesOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Replace the requested pointer access with corresponding field.
+    // The cast_op is inserted by type converter to intermix 1:N type
+    // conversion.
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     uint64_t dim = op.getDimension().getZExtValue();
-    return desc.getIdxMemRef(dim);
+    Value field = desc.getIdxMemRef(dim);
+
+    // Insert a cast to bridge the actual type to the user expected type. If the
+    // actual type and the user expected type aren't compatible, the compiler or
+    // the runtime will issue an error.
+    Type resType = op.getResult().getType();
+    if (resType != field.getType())
+      field = rewriter.create<memref::CastOp>(op.getLoc(), resType, field);
+    rewriter.replaceOp(op, field);
+
+    return success();
   }
 };
 
 /// Sparse codegen rule for value accesses.
 class SparseToValuesConverter
     : public SparseGetterOpConverter<ToValuesOp, SparseToValuesConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;

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

 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type indType = resType.cast<ShapedType>().getElementType();
     SmallString<15> name{"sparseIndices", overheadTypeFunctionSuffix(indType)};
-    Value dim =
-        constantIndex(rewriter, op->getLoc(), op.getDimension().getZExtValue());
-    replaceOpWithFuncCall(rewriter, op, name, resType,
-                          {adaptor.getTensor(), dim}, EmitCInterface::On);
+    Location loc = op->getLoc();
+    Value dim = constantIndex(rewriter, loc, op.getDimension().getZExtValue());
+
+    // The function returns a MemRef without a layout.
+    MemRefType callRetType = get1DMemRefType(indType, false);
+    SmallVector<Value> operands{adaptor.getTensor(), dim};
+    auto fn = getFunc(op->getParentOfType<ModuleOp>(), name, callRetType,
+                      operands, EmitCInterface::On);
+    Value callRet =
+        rewriter.create<func::CallOp>(loc, callRetType, fn, operands)
+            .getResult(0);
+
+    // Cast the MemRef type to the type expected by the users, though these
+    // two types should be compatible at runtime.
+    if (resType != callRetType)
+      callRet = rewriter.create<memref::CastOp>(loc, resType, callRet);
+    rewriter.replaceOp(op, callRet);
+
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
 class SparseTensorToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;

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

LogicalResult sparse2SparseRewrite(ConvertOp op,
     }
 
     // Only need to sort if the srcTp is not already sorted (we faithfully take
     // the guarantee from the sparse tensor encoding).
     if (!isAllDimOrdered(srcTp)) {
       // Sort the COO tensor so that its elements are ordered via increasing
       // indices for the storage ordering of the dst tensor.
       SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
-      auto dynShape = {ShapedType::kDynamic};
-      auto indTp =
-          MemRefType::get(dynShape, getIndexOverheadType(rewriter, encSrc));
       uint64_t rank = dstTp.getRank();
+      uint64_t cooStart = getCooStart(encSrc);
       // Gather the indices-arrays in the dst tensor storage order.
       SmallVector<Value> xs(rank, Value());
       for (uint64_t i = 0; i < rank; i++) {
         uint64_t orgDim = toOrigDim(encSrc, i);
+        MemRefType indTp = get1DMemRefType(
+            getIndexOverheadType(rewriter, encSrc), i >= cooStart);
         xs[toStoredDim(encDst, orgDim)] = rewriter.create<ToIndicesOp>(
             loc, indTp, src, rewriter.getIndexAttr(i));
       }
 
       // Retrieve NNZ.
       Value nnz = rewriter.create<NumberOfEntriesOp>(loc, src);
       nnz = rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(),
                                                 nnz);
 
       // Retrieve the values-array.
-      auto valTp = MemRefType::get(dynShape, srcTp.getElementType());
+      MemRefType valTp = get1DMemRefType(srcTp.getElementType(), false);
       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> dynDstSizes;

mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir

 //       CHECK-RWT:        }
 //       CHECK-RWT:        sparse_tensor.yield %[[IFR]]
 //       CHECK-RWT:     }
 //       CHECK-RWT:     %[[COO:.*]] = sparse_tensor.load %[[T2]] hasInserts
 //       CHECK-RWT:     %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
 //       CHECK-RWT:     %[[I1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
 //       CHECK-RWT:     %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
 //       CHECK-RWT:     %[[V:.*]] = sparse_tensor.values %[[COO]]
-//       CHECK-RWT:     sparse_tensor.sort %[[NNZ]], %[[I0]], %[[I1]] jointly %[[V]] : memref<?xindex>, memref<?xindex> jointly memref<?xf64>
+//       CHECK-RWT:     sparse_tensor.sort %[[NNZ]], %[[I0]], %[[I1]] jointly %[[V]]
 //       CHECK-RWT:     %[[T3:.*]] = bufferization.alloc_tensor()
 //       CHECK-RWT:     %[[T4:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T3]])
 //       CHECK-RWT:     ^bb0(%[[L1I0:.*]]: index, %[[L1I1:.*]]: index, %[[L1V:.*]]: f64, %[[L1T:.*]]: tensor
 //       CHECK-RWT:       %[[L1T2:.*]] = sparse_tensor.insert %[[L1V]] into %[[L1T]]{{\[}}%[[L1I0]], %[[L1I1]]]
 //       CHECK-RWT:       sparse_tensor.yield %[[L1T2]]
 //       CHECK-RWT:     }
 //       CHECK-RWT:     %[[T5:.*]] = sparse_tensor.load %[[T4]] hasInserts
 //       CHECK-RWT:     %[[T6:.*]] = sparse_tensor.convert %[[T5]]

 //       CHECK-RWT:       %[[L0T2:.*]] = sparse_tensor.insert %[[L0V]] into %[[L0T]]{{\[}}%[[L0I0]], %[[L0I1]]]
 //       CHECK-RWT:       sparse_tensor.yield %[[L0T2]]
 //       CHECK-RWT:     }
 //       CHECK-RWT:     %[[COO:.*]] = sparse_tensor.load %[[T1]] hasInserts
 //       CHECK-RWT:     %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
 //       CHECK-RWT:     %[[I1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
 //       CHECK-RWT:     %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
 //       CHECK-RWT:     %[[V:.*]] = sparse_tensor.values %[[COO]]
-//       CHECK-RWT:     sparse_tensor.sort %[[NNZ]], %[[I0]], %[[I1]] jointly %[[V]] : memref<?xindex>, memref<?xindex> jointly memref<?xf32>
+//       CHECK-RWT:     sparse_tensor.sort %[[NNZ]], %[[I0]], %[[I1]] jointly %[[V]]
 //       CHECK-RWT:     %[[T3:.*]] = bufferization.alloc_tensor()
 //       CHECK-RWT:     %[[T4:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T3]])
 //       CHECK-RWT:     ^bb0(%[[L1I0:.*]]: index, %[[L1I1:.*]]: index, %[[L1V:.*]]: f32, %[[L1T:.*]]: tensor
 //       CHECK-RWT:       %[[L1T2:.*]] = sparse_tensor.insert %[[L1V]] into %[[L1T]]{{\[}}%[[L1I0]], %[[L1I1]]]
 //       CHECK-RWT:       sparse_tensor.yield %[[L1T2]]
 //       CHECK-RWT:     }
 //       CHECK-RWT:     %[[T5:.*]] = sparse_tensor.load %[[T4]] hasInserts
 //       CHECK-RWT:     %[[T6:.*]] = sparse_tensor.convert %[[T5]]

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

 
 // CHECK-LABEL: func.func @matvec(
 // CHECK-SAME:    %[[VAL_0:.*]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: tensor<64xf64>,
 // CHECK-SAME:    %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
 // CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
+// CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
 // CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
 // CHECK-DAG:     %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<64xf64>
 // CHECK-DAG:     %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xf64>
 // CHECK-DAG:     %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
 // CHECK-DAG:     %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:         scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
-// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex
 // CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf64>
-// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex
 // CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf64>
 // CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_16]]] : memref<64xf64>
 // CHECK:           %[[VAL_19:.*]] = arith.mulf %[[VAL_17]], %[[VAL_18]] : f64
 // CHECK:           %[[VAL_20:.*]] = arith.addf %[[VAL_15]], %[[VAL_19]] : f64
 // CHECK:           memref.store %[[VAL_20]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf64>
 // CHECK:         }
 // CHECK:         %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf64>
 // CHECK:         return %[[VAL_21]] : tensor<32xf64>

 // CHECK-LABEL: func.func @mateltmul(
 // CHECK-SAME:    %[[VAL_0:.*0]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
 // CHECK-SAME:    %[[VAL_1:.*1]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
 // CHECK-SAME:    %[[VAL_2:.*2]]: tensor<32x64xf64>) -> tensor<32x64xf64> {
 // CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f64
 // CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
+// CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
 // CHECK-DAG:     %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
 // CHECK-DAG:     %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
-// CHECK-DAG:     %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
+// CHECK-DAG:     %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex
 // CHECK-DAG:     %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
 // CHECK-DAG:     %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x64xf64>
 // CHECK-DAG:     linalg.fill ins(%[[VAL_3]] : f64) outs(%[[VAL_14]] : memref<32x64xf64>)
 // CHECK-DAG:     %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK-DAG:     %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK-DAG:     %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK-DAG:     %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:         %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_15]], %[[VAL_21:.*]] = %[[VAL_17]]) : (index, index) -> (index, index) {
 // CHECK:           %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_16]] : index
 // CHECK:           %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_18]] : index
 // CHECK:           %[[VAL_24:.*]] = arith.andi %[[VAL_22]], %[[VAL_23]] : i1
 // CHECK:           scf.condition(%[[VAL_24]]) %[[VAL_20]], %[[VAL_21]] : index, index
 // CHECK:         } do {
 // CHECK:         ^bb0(%[[VAL_25:.*]]: index, %[[VAL_26:.*]]: index):
-// CHECK:           %[[VAL_27:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_25]]] : memref<?xindex>
-// CHECK:           %[[VAL_28:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_26]]] : memref<?xindex>
+// CHECK:           %[[VAL_27:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_25]]] : memref<?xindex
+// CHECK:           %[[VAL_28:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_26]]] : memref<?xindex
 // CHECK:           %[[VAL_29:.*]] = arith.cmpi ult, %[[VAL_28]], %[[VAL_27]] : index
 // CHECK:           %[[VAL_30:.*]] = arith.select %[[VAL_29]], %[[VAL_28]], %[[VAL_27]] : index
 // CHECK:           %[[VAL_31:.*]] = arith.cmpi eq, %[[VAL_27]], %[[VAL_30]] : index
 // CHECK:           %[[VAL_32:.*]] = arith.cmpi eq, %[[VAL_28]], %[[VAL_30]] : index
 // CHECK:           %[[VAL_33:.*]] = arith.andi %[[VAL_31]], %[[VAL_32]] : i1
 // CHECK:           scf.if %[[VAL_33]] {
 // CHECK:             %[[VAL_34:.*]] = arith.addi %[[VAL_25]], %[[VAL_5]] : index
 // CHECK:             %[[VAL_35:.*]] = arith.addi %[[VAL_26]], %[[VAL_5]] : index
 // CHECK:             %[[VAL_36:.*]]:2 = scf.while (%[[VAL_37:.*]] = %[[VAL_25]], %[[VAL_38:.*]] = %[[VAL_26]]) : (index, index) -> (index, index) {
 // CHECK:               %[[VAL_39:.*]] = arith.cmpi ult, %[[VAL_37]], %[[VAL_34]] : index
 // CHECK:               %[[VAL_40:.*]] = arith.cmpi ult, %[[VAL_38]], %[[VAL_35]] : index
 // CHECK:               %[[VAL_41:.*]] = arith.andi %[[VAL_39]], %[[VAL_40]] : i1
 // CHECK:               scf.condition(%[[VAL_41]]) %[[VAL_37]], %[[VAL_38]] : index, index
 // CHECK:             } do {
 // CHECK:             ^bb0(%[[VAL_42:.*]]: index, %[[VAL_43:.*]]: index):
-// CHECK:               %[[VAL_44:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_42]]] : memref<?xindex>
-// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_43]]] : memref<?xindex>
+// CHECK:               %[[VAL_44:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_42]]] : memref<?xindex
+// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_43]]] : memref<?xindex
 // CHECK:               %[[VAL_46:.*]] = arith.cmpi ult, %[[VAL_45]], %[[VAL_44]] : index
 // CHECK:               %[[VAL_47:.*]] = arith.select %[[VAL_46]], %[[VAL_45]], %[[VAL_44]] : index
 // CHECK:               %[[VAL_48:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_47]] : index
 // CHECK:               %[[VAL_49:.*]] = arith.cmpi eq, %[[VAL_45]], %[[VAL_47]] : index
 // CHECK:               %[[VAL_50:.*]] = arith.andi %[[VAL_48]], %[[VAL_49]] : i1
 // CHECK:               scf.if %[[VAL_50]] {
 // CHECK:                 %[[VAL_51:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_42]]] : memref<?xf64>
 // CHECK:                 %[[VAL_52:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_43]]] : memref<?xf64>