[mlir][sparse] packs sparse tensor sizes using llvm.struct instead of memref

[Peiming]

Summary

The sparse compiler used to stored the dimension sizes with a MemRef<rank x index>. This patch change the memory scheme by storing the dimension size using llvm.struct<array<rank x int>>.

There are multiple reasons behind the change

1. MemRef, in addition to pointers, stores extra information for strides and offsets, which are redundant for our case.
2. MemRef requires deallocation to avoid memory leak, which brings us extra challenges to implement sparse tensor slice in the future as sparse tensor slices will likely to hold a different MemRef for dimension sizes, which is required to be freed after the slices is no longer needed.

Note that this scheme is also closer to how MemRef stores its dimension sizes.

Future work:
Move the memSizes MemRef into the same llvm.struct together with dimSizes, it is currently blocked by sparse_tensor.push_back operation, which only takes a MemRef memSizes as inputs.

[Peiming]

rebase

[wrengr]

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

2

The name of the file should be right-aligned.

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

2

The name of the file should be left-aligned, and the end of the comment line should include -*- C++ -*- (as per the style guide).

[Peiming]

fix style.

[Peiming]

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

191

Is there any way that I can get the same bitwidth for index type as used by lowerToLLVM passes?

[Peiming]

rebase

[Peiming]

Will use sparse_tensor::StorageSpecifier for the same purpose

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace bufferization {
 struct OneShotBufferizationOptions;
 } // namespace bufferization
-
+namespace sparse_tensor {
+class SparseTensorTypeToBufferConverter;
+} // namespace sparse_tensor
 //===----------------------------------------------------------------------===//
 // The Sparsification pass.
 //===----------------------------------------------------------------------===//
 
 /// Defines a parallelization strategy. Any independent loop is a candidate
 /// for parallelization. The loop is made parallel if (1) allowed by the
 /// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse
 /// outermost loop only), and (2) the generated code is an actual for-loop

 std::unique_ptr<Pass> createSparseTensorConversionPass();
 std::unique_ptr<Pass>
 createSparseTensorConversionPass(const SparseTensorConversionOptions &options);
 
 //===----------------------------------------------------------------------===//
 // The SparseTensorCodegen pass.
 //===----------------------------------------------------------------------===//
 
-/// Sparse tensor type converter into an actual buffer.
-class SparseTensorTypeToBufferConverter : public TypeConverter {
-public:
-  SparseTensorTypeToBufferConverter();
-};
-
 /// Sets up sparse tensor conversion rules.
-void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
-                                         RewritePatternSet &patterns,
-                                         bool enableBufferInitialization);
+void populateSparseTensorCodegenPatterns(
+    sparse_tensor::SparseTensorTypeToBufferConverter &typeConverter,
+    RewritePatternSet &patterns, bool enableBufferInitialization);
 
 std::unique_ptr<Pass> createSparseTensorCodegenPass();
 std::unique_ptr<Pass>
 createSparseTensorCodegenPass(bool enableBufferInitialization);
 
 //===----------------------------------------------------------------------===//
 // The PreSparsificationRewriting pass.
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

```mlir
           return %arg2 : memref<?xindex>
         }
     ```
   }];
   let constructor = "mlir::createSparseTensorCodegenPass()";
   let dependentDialects = [
     "arith::ArithDialect",
     "bufferization::BufferizationDialect",
+    "LLVM::LLVMDialect",
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
   let options = [
     Option<"enableBufferInitialization", "enable-buffer-initialization", "bool",
            "false", "Enable zero-initialization of the memory buffers">,

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenUtils.cpp
   SparseBufferRewriting.cpp
+  SparseTensorBuilder.cpp
   SparseTensorCodegen.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
   SparseTensorRewriting.cpp
   SparseVectorization.cpp
   Sparsification.cpp
   SparsificationAndBufferizationPass.cpp
 

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

 }
 
 inline bool isZeroRankedTensorOrScalar(Type type) {
   auto rtp = type.dyn_cast<RankedTensorType>();
   return !rtp || rtp.getRank() == 0;
 }
 
 //===----------------------------------------------------------------------===//
-// SparseTensorDescriptor and helpers, manage the sparse tensor memory layout
-// scheme.
-//
-// Sparse tensor storage scheme for rank-dimensional tensor is organized
-// as a single compound type with the following fields. Note that every
-// memref with ? size actually behaves as a "vector", i.e. the stored
-// size is the capacity and the used size resides in the memSizes array.
-//
-// struct {
-//   memref<rank x index> dimSizes     ; size in each dimension
-//   memref<n x index> memSizes        ; sizes of ptrs/inds/values
-//   ; per-dimension d:
-//   ;  if dense:
-//        <nothing>
-//   ;  if compresed:
-//        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
-//        memref<? x idx>  indices-d   ; indices for sparse dim d
-//   ;  if singleton:
-//        memref<? x idx>  indices-d   ; indices for singleton dim d
-//   memref<? x eltType> values        ; values
-// };
-//
-//===----------------------------------------------------------------------===//
-enum class SparseTensorFieldKind {
-  DimSizes,
-  MemSizes,
-  PtrMemRef,
-  IdxMemRef,
-  ValMemRef
-};
-
-constexpr uint64_t dimSizesIdx = 0;
-constexpr uint64_t memSizesIdx = dimSizesIdx + 1;
-constexpr uint64_t dataFieldIdx = memSizesIdx + 1;
-
-/// For each field that will be allocated for the given sparse tensor encoding,
-/// calls the callback with the corresponding field index, field kind, dimension
-/// (for sparse tensor level memrefs) and dimlevelType.
-/// The field index always starts with zero and increments by one between two
-/// callback invocations.
-/// Ideally, all other methods should rely on this function to query a sparse
-/// tensor fields instead of relying on ad-hoc index computation.
-void foreachFieldInSparseTensor(
-    SparseTensorEncodingAttr,
-    llvm::function_ref<bool(unsigned /*fieldIdx*/,
-                            SparseTensorFieldKind /*fieldKind*/,
-                            unsigned /*dim (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>);
-
-/// Same as above, except that it also builds the Type for the corresponding
-/// field.
-void foreachFieldAndTypeInSparseTensor(
-    RankedTensorType,
-    llvm::function_ref<bool(Type /*fieldType*/, unsigned /*fieldIdx*/,
-                            SparseTensorFieldKind /*fieldKind*/,
-                            unsigned /*dim (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>);
-
-/// Gets the total number of fields for the given sparse tensor encoding.
-unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-/// Gets the total number of data fields (index arrays, pointer arrays, and a
-/// value array) for the given sparse tensor encoding.
-unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-/// Get the index of the field in memSizes (only valid for data fields).
-inline unsigned getFieldMemSizesIndex(unsigned fid) {
-  assert(fid >= dataFieldIdx);
-  return fid - dataFieldIdx;
-}
-
-template <bool>
-struct SparseTensorValueArrayRef;
-
-// Uses ValueRange for immuatable descriptors; uses SmallVectorImpl<Value> &
-// for mutable descriptors.
-template <>
-struct SparseTensorValueArrayRef<false> {
-  using ValueArray = ValueRange;
-};
-
-// Using SmallVector for mutable descriptor allows users to reuse it as a tmp
-// buffers to append value for some special cases, though users should be
-// responsible to restore the buffer to legal states after their use. It is
-// probably not a clean way, but it is the most efficient way to avoid copying
-// the fields into another SmallVector. If a more clear way is wanted, we
-// should change it to MutableArrayRef instead.
-template <>
-struct SparseTensorValueArrayRef<true> {
-  using ValueArray = SmallVectorImpl<Value> &;
-};
-
-/// A helper class around an array of values that corresponding to a sparse
-/// tensor, provides a set of meaningful APIs to query and update a particular
-/// field in a consistent way.
-/// Users should not make assumption on how a sparse tensor is laid out but
-/// instead relies on this class to access the right value for the right field.
-template <bool mut>
-class SparseTensorDescriptorImpl {
-private:
-  using Storage = typename SparseTensorValueArrayRef<mut>::ValueArray;
-
-public:
-  SparseTensorDescriptorImpl(Type tp, Storage fields)
-      : rType(tp.cast<RankedTensorType>()), fields(fields) {
-    assert(getSparseTensorEncoding(tp) &&
-           getNumFieldsFromEncoding(getSparseTensorEncoding(tp)) ==
-               fields.size());
-    // We should make sure the class is trivially copyable (and should be small
-    // enough) such that we can pass it by value.
-    static_assert(
-        std::is_trivially_copyable_v<SparseTensorDescriptorImpl<mut>>);
-  }
-
-  // Implicit (and cheap) type conversion from MutSparseTensorDescriptor to
-  // SparseTensorDescriptor.
-  template <typename T = SparseTensorDescriptorImpl<true>>
-  /*implicit*/ SparseTensorDescriptorImpl(std::enable_if_t<!mut, T> &mDesc)
-      : rType(mDesc.getTensorType()), fields(mDesc.getFields()) {}
-
-  ///
-  /// Getters: get the field index for required field.
-  ///
-
-  unsigned getPtrMemRefIndex(unsigned ptrDim) const {
-    return getFieldIndex(ptrDim, SparseTensorFieldKind::PtrMemRef);
-  }
-
-  unsigned getIdxMemRefIndex(unsigned idxDim) const {
-    return getFieldIndex(idxDim, SparseTensorFieldKind::IdxMemRef);
-  }
-
-  unsigned getValMemRefIndex() const { return fields.size() - 1; }
-
-  unsigned getPtrMemSizesIndex(unsigned dim) const {
-    return getPtrMemRefIndex(dim) - dataFieldIdx;
-  }
-
-  unsigned getIdxMemSizesIndex(unsigned dim) const {
-    return getIdxMemRefIndex(dim) - dataFieldIdx;
-  }
-
-  unsigned getValMemSizesIndex() const {
-    return getValMemRefIndex() - dataFieldIdx;
-  }
-
-  unsigned getNumFields() const { return fields.size(); }
-
-  ///
-  /// Getters: get the value for required field.
-  ///
-
-  Value getDimSizesMemRef() const { return fields[dimSizesIdx]; }
-  Value getMemSizesMemRef() const { return fields[memSizesIdx]; }
-
-  Value getPtrMemRef(unsigned ptrDim) const {
-    return fields[getPtrMemRefIndex(ptrDim)];
-  }
-
-  Value getIdxMemRef(unsigned idxDim) const {
-    return fields[getIdxMemRefIndex(idxDim)];
-  }
-
-  Value getValMemRef() const { return fields[getValMemRefIndex()]; }
-
-  Value getField(unsigned fid) const {
-    assert(fid < fields.size());
-    return fields[fid];
-  }
-
-  ///
-  /// Setters: update the value for required field (only enabled for
-  /// MutSparseTensorDescriptor).
-  ///
-
-  template <typename T = Value>
-  void setField(unsigned fid, std::enable_if_t<mut, T> v) {
-    assert(fid < fields.size());
-    fields[fid] = v;
-  }
-
-  RankedTensorType getTensorType() const { return rType; }
-  Storage getFields() const { return fields; }
-
-  Type getElementType(unsigned fidx) const {
-    return fields[fidx].getType().template cast<MemRefType>().getElementType();
-  }
-
-private:
-  unsigned getFieldIndex(unsigned dim, SparseTensorFieldKind kind) const {
-    unsigned fieldIdx = -1u;
-    foreachFieldInSparseTensor(
-        getSparseTensorEncoding(rType),
-        [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
-                               unsigned fDim, DimLevelType dlt) -> bool {
-          if (fDim == dim && kind == fKind) {
-            fieldIdx = fIdx;
-            // Returns false to break the iteration.
-            return false;
-          }
-          return true;
-        });
-    assert(fieldIdx != -1u);
-    return fieldIdx;
-  }
-
-  RankedTensorType rType;
-  Storage fields;
-};
-
-using SparseTensorDescriptor = SparseTensorDescriptorImpl<false>;
-using MutSparseTensorDescriptor = SparseTensorDescriptorImpl<true>;
-
-//===----------------------------------------------------------------------===//
 // SparseTensorLoopEmiter class, manages sparse tensors and helps to
 // generate loop structure to (co)-iterate sparse tensors.
 //
 // An example usage:
 // To generate the following loops over T1<?x?> and T2<?x?>
 //
 // for i in TENSOR_1_0 {
 //   for j : TENSOR_2_0 {

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

static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
                                       Value tensor,
                                       SmallVectorImpl<Value> &indicesArray,
                                       ValueRange ivs) {
   Value val = genValueForDense(builder, loc, tensor, ivs);
   indicesArray.append(ivs.begin(), ivs.end());
   return val;
 }
 
-void sparse_tensor::foreachFieldInSparseTensor(
-    const SparseTensorEncodingAttr enc,
-    llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
-                            DimLevelType)>
-        callback) {
-  assert(enc);
-
-#define RETURN_ON_FALSE(idx, kind, dim, dlt)                                   \
-  if (!(callback(idx, kind, dim, dlt)))                                        \
-    return;
-
-  RETURN_ON_FALSE(dimSizesIdx, SparseTensorFieldKind::DimSizes, -1u,
-                  DimLevelType::Undef);
-  RETURN_ON_FALSE(memSizesIdx, SparseTensorFieldKind::MemSizes, -1u,
-                  DimLevelType::Undef);
-
-  static_assert(dataFieldIdx == memSizesIdx + 1);
-  unsigned fieldIdx = dataFieldIdx;
-  // Per-dimension storage.
-  for (unsigned r = 0, rank = enc.getDimLevelType().size(); r < rank; r++) {
-    // Dimension level types apply in order to the reordered dimension.
-    // As a result, the compound type can be constructed directly in the given
-    // order.
-    auto dlt = getDimLevelType(enc, r);
-    if (isCompressedDLT(dlt)) {
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PtrMemRef, r, dlt);
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
-    } else if (isSingletonDLT(dlt)) {
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
-    } else {
-      assert(isDenseDLT(dlt)); // no fields
-    }
-  }
-  // The values array.
-  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::ValMemRef, -1u,
-                  DimLevelType::Undef);
-
-#undef RETURN_ON_FALSE
-}
-
-void sparse_tensor::foreachFieldAndTypeInSparseTensor(
-    RankedTensorType rType,
-    llvm::function_ref<bool(Type, unsigned, SparseTensorFieldKind, unsigned,
-                            DimLevelType)>
-        callback) {
-  auto enc = getSparseTensorEncoding(rType);
-  assert(enc);
-  // Construct the basic types.
-  Type indexType = IndexType::get(enc.getContext());
-  Type idxType = enc.getIndexType();
-  Type ptrType = enc.getPointerType();
-  Type eltType = rType.getElementType();
-  unsigned rank = rType.getShape().size();
-  // memref<rank x index> dimSizes
-  Type dimSizeType = MemRefType::get({rank}, indexType);
-  // memref<n x index> memSizes
-  Type memSizeType =
-      MemRefType::get({getNumDataFieldsFromEncoding(enc)}, indexType);
-  // memref<? x ptr>  pointers
-  Type ptrMemType = MemRefType::get({ShapedType::kDynamic}, ptrType);
-  // memref<? x idx>  indices
-  Type idxMemType = MemRefType::get({ShapedType::kDynamic}, idxType);
-  // memref<? x eltType> values
-  Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
-
-  foreachFieldInSparseTensor(
-      enc,
-      [dimSizeType, memSizeType, ptrMemType, idxMemType, valMemType,
-       callback](unsigned fieldIdx, SparseTensorFieldKind fieldKind,
-                 unsigned dim, DimLevelType dlt) -> bool {
-        switch (fieldKind) {
-        case SparseTensorFieldKind::DimSizes:
-          return callback(dimSizeType, fieldIdx, fieldKind, dim, dlt);
-        case SparseTensorFieldKind::MemSizes:
-          return callback(memSizeType, fieldIdx, fieldKind, dim, dlt);
-        case SparseTensorFieldKind::PtrMemRef:
-          return callback(ptrMemType, fieldIdx, fieldKind, dim, dlt);
-        case SparseTensorFieldKind::IdxMemRef:
-          return callback(idxMemType, fieldIdx, fieldKind, dim, dlt);
-        case SparseTensorFieldKind::ValMemRef:
-          return callback(valMemType, fieldIdx, fieldKind, dim, dlt);
-        };
-        llvm_unreachable("unrecognized field kind");
-      });
-}
-
-unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-  unsigned numFields = 0;
-  foreachFieldInSparseTensor(enc,
-                             [&numFields](unsigned, SparseTensorFieldKind,
-                                          unsigned, DimLevelType) -> bool {
-                               numFields++;
-                               return true;
-                             });
-  return numFields;
-}
-
-unsigned
-sparse_tensor::getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-  unsigned numFields = 0; // one value memref
-  foreachFieldInSparseTensor(enc,
-                             [&numFields](unsigned fidx, SparseTensorFieldKind,
-                                          unsigned, DimLevelType) -> bool {
-                               if (fidx >= dataFieldIdx)
-                                 numFields++;
-                               return true;
-                             });
-  assert(numFields == getNumFieldsFromEncoding(enc) - dataFieldIdx);
-  return numFields;
-}
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  StringAttr loopTag,
                                                  bool hasOutput,
                                                  bool isSparseOut,

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

+//===- SparseTensorBuilder.h ------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines utilities for lowering and access sparse tensor
+// types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
+
+#include "mlir/Conversion/LLVMCommon/StructBuilder.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace sparse_tensor {
+
+class SparseTensorTypeToBufferConverter : public TypeConverter {
+public:
+  SparseTensorTypeToBufferConverter(unsigned indexWidth);
+
+  /// Gets the bitwidth of the index type when converted to LLVM.
+  unsigned getIndexTypeBitwidth() const { return indexWidth; }
+
+  Optional<LogicalResult>
+  convertSparseTensorType(RankedTensorType rtp, SmallVectorImpl<Type> &fields);
+
+private:
+  unsigned indexWidth;
+};
+
+class SparseTensorMetaData : public StructBuilder {
+public:
+  explicit SparseTensorMetaData(Value metadata) : StructBuilder(metadata) {
+    assert(value); // value inherented from StructBuilder must have been set.
+  }
+
+  static SparseTensorMetaData undef(OpBuilder &builder, Location loc,
+                                    Type metaType);
+  static Type getMetaDataIndexType(Value data);
+
+  Value dimSize(OpBuilder &builder, Location loc, unsigned dim);
+  void setDimSize(OpBuilder &builder, Location loc, unsigned dim, Value size);
+};
+
+//===----------------------------------------------------------------------===//
+// SparseTensorDescriptor and helpers, manage the sparse tensor memory layout
+// scheme.
+//
+// Sparse tensor storage scheme for rank-dimensional tensor is organized
+// as a single compound type with the following fields. Note that every
+// memref with ? size actually behaves as a "vector", i.e. the stored
+// size is the capacity and the used size resides in the memSizes array.
+//
+// struct {
+//   ; sparse tensor metadata
+//   struct {
+//     array<rank x int> dimSizes    ; size in each dimension
+//   }
+//   memref<n x index> memSizes        ; sizes of ptrs/inds/values
+//   ; per-dimension d:
+//   ;  if dense:
+//        <nothing>
+//   ;  if compresed:
+//        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
+//        memref<? x idx>  indices-d   ; indices for sparse dim d
+//   ;  if singleton:
+//        memref<? x idx>  indices-d   ; indices for singleton dim d
+//   memref<? x eltType> values        ; values
+// };
+//
+//===----------------------------------------------------------------------===//
+enum class SparseTensorFieldKind {
+  MetaData,
+  MemSizes,
+  PtrMemRef,
+  IdxMemRef,
+  ValMemRef
+};
+
+constexpr uint64_t kMemSizesIdx = 0;
+constexpr uint64_t kDataFieldIdx = kMemSizesIdx + 1;
+
+/// For each field that will be allocated for the given sparse tensor encoding,
+/// calls the callback with the corresponding field index, field kind, dimension
+/// (for sparse tensor level memrefs) and dimlevelType.
+/// The field index always starts with zero and increments by one between two
+/// callback invocations.
+/// Ideally, all other methods should rely on this function to query a sparse
+/// tensor fields instead of relying on ad-hoc index computation.
+void foreachFieldInSparseTensor(
+    SparseTensorEncodingAttr,
+    llvm::function_ref<bool(unsigned /*fieldIdx*/,
+                            SparseTensorFieldKind /*fieldKind*/,
+                            unsigned /*dim (if applicable)*/,
+                            DimLevelType /*DLT (if applicable)*/)>);
+
+/// Same as above, except that it also builds the Type for the corresponding
+/// field.
+void foreachFieldAndTypeInSparseTensor(
+    RankedTensorType, unsigned,
+    llvm::function_ref<bool(Type /*fieldType*/, unsigned /*fieldIdx*/,
+                            SparseTensorFieldKind /*fieldKind*/,
+                            unsigned /*dim (if applicable)*/,
+                            DimLevelType /*DLT (if applicable)*/)>);
+
+/// Gets the total number of fields for the given sparse tensor encoding.
+unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
+
+/// Gets the total number of data fields (index arrays, pointer arrays, and a
+/// value array) for the given sparse tensor encoding.
+unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc);
+
+/// Get the index of the field in memSizes (only valid for data fields).
+inline unsigned getFieldMemSizesIndex(unsigned fid) {
+  assert(fid >= kDataFieldIdx);
+  return fid - kDataFieldIdx;
+}
+
+/// A helper class around an array of values that corresponding to a sparse
+/// tensor, provides a set of meaningful APIs to query and update a particular
+/// field in a consistent way.
+/// Users should not make assumption on how a sparse tensor is laid out but
+/// instead relies on this class to access the right value for the right field.
+template <bool mut>
+class SparseTensorDescriptorImpl {
+private:
+  // Uses ValueRange for immuatable descriptors; uses SmallVectorImpl<Value> &
+  // for mutable descriptors.
+  // Using SmallVector for mutable descriptor allows users to reuse it as a tmp
+  // buffers to append value for some special cases, though users should be
+  // responsible to restore the buffer to legal states after their use. It is
+  // probably not a clean way, but it is the most efficient way to avoid copying
+  // the fields into another SmallVector. If a more clear way is wanted, we
+  // should change it to MutableArrayRef instead.
+  using ValueArrayRef = typename std::conditional<mut, SmallVectorImpl<Value> &,
+                                                  ValueRange>::type;
+
+public:
+  SparseTensorDescriptorImpl(Type tp, ValueArrayRef fields)
+      : indexType(SparseTensorMetaData::getMetaDataIndexType(fields.back())),
+        rType(tp.cast<RankedTensorType>()), fields(fields) {
+    assert(getSparseTensorEncoding(tp) &&
+           getNumFieldsFromEncoding(getSparseTensorEncoding(tp)) ==
+               fields.size());
+    // We should make sure the class is trivially copyable (and should be small
+    // enough) such that we can pass it by value.
+    static_assert(
+        std::is_trivially_copyable_v<SparseTensorDescriptorImpl<mut>>);
+  }
+
+  // Implicit (and cheap) type conversion from MutSparseTensorDescriptor to
+  // SparseTensorDescriptor.
+  template <typename T = SparseTensorDescriptorImpl<true>>
+  /*implicit*/ SparseTensorDescriptorImpl(std::enable_if_t<!mut, T> &mDesc)
+      : indexType(mDesc.getIndexType()), rType(mDesc.getTensorType()),
+        fields(mDesc.getFields()) {}
+
+  ///
+  /// Getters: get the field index for required field.
+  ///
+
+  unsigned getPtrMemRefIndex(unsigned ptrDim) const {
+    return getFieldIndex(ptrDim, SparseTensorFieldKind::PtrMemRef);
+  }
+
+  unsigned getIdxMemRefIndex(unsigned idxDim) const {
+    return getFieldIndex(idxDim, SparseTensorFieldKind::IdxMemRef);
+  }
+
+  unsigned getValMemRefIndex() const { return fields.size() - 2; }
+
+  unsigned getPtrMemSizesIndex(unsigned dim) const {
+    return getPtrMemRefIndex(dim) - kDataFieldIdx;
+  }
+
+  unsigned getIdxMemSizesIndex(unsigned dim) const {
+    return getIdxMemRefIndex(dim) - kDataFieldIdx;
+  }
+
+  unsigned getValMemSizesIndex() const {
+    return getValMemRefIndex() - kDataFieldIdx;
+  }
+
+  unsigned getNumFields() const { return fields.size(); }
+
+  ///
+  /// Getters: get the value for required field.
+  ///
+
+  Value getDimSize(OpBuilder &builder, Location loc, unsigned dim) const {
+    SparseTensorMetaData md(fields.back());
+    return md.dimSize(builder, loc, dim);
+  }
+
+  Value getMemSizesMemRef() const { return fields[kMemSizesIdx]; }
+
+  Value getPtrMemRef(unsigned ptrDim) const {
+    return fields[getPtrMemRefIndex(ptrDim)];
+  }
+
+  Value getIdxMemRef(unsigned idxDim) const {
+    return fields[getIdxMemRefIndex(idxDim)];
+  }
+
+  Value getValMemRef() const { return fields[getValMemRefIndex()]; }
+
+  Value getField(unsigned fid) const {
+    assert(fid < fields.size());
+    return fields[fid];
+  }
+
+  ///
+  /// Setters: update the value for required field (only enabled for
+  /// MutSparseTensorDescriptor).
+  ///
+
+  template <typename T = Value>
+  void setField(unsigned fid, std::enable_if_t<mut, T> v) {
+    assert(fid < fields.size());
+    fields[fid] = v;
+  }
+
+  template <typename T = Value>
+  void setDimSize(OpBuilder &builder, Location loc, unsigned dim,
+                  std::enable_if_t<mut, T> v) {
+    assert(v.getType() == getIndexType());
+    SparseTensorMetaData md(fields.back());
+    md.setDimSize(builder, loc, dim, v);
+    fields.back() = md;
+  }
+
+  ValueRange getMemRefFields() const {
+    ValueRange ret = fields;
+    // drop the last metadata fields
+    return ret.slice(0, fields.size() - 1);
+  }
+
+  RankedTensorType getTensorType() const { return rType; }
+  ValueArrayRef getFields() const { return fields; }
+  Type getIndexType() const { return indexType; }
+  Type getElementType(unsigned fidx) const {
+    return fields[fidx].getType().template cast<MemRefType>().getElementType();
+  }
+
+private:
+  unsigned getFieldIndex(unsigned dim, SparseTensorFieldKind kind) const {
+    unsigned fieldIdx = -1u;
+    foreachFieldInSparseTensor(
+        getSparseTensorEncoding(rType),
+        [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
+                               unsigned fDim, DimLevelType dlt) -> bool {
+          if (fDim == dim && kind == fKind) {
+            fieldIdx = fIdx;
+            // Returns false to break the iteration.
+            return false;
+          }
+          return true;
+        });
+    assert(fieldIdx != -1u);
+    return fieldIdx;
+  }
+
+  Type indexType;
+  RankedTensorType rType;
+  ValueArrayRef fields;
+};
+
+using SparseTensorDescriptor = SparseTensorDescriptorImpl<false>;
+using MutSparseTensorDescriptor = SparseTensorDescriptorImpl<true>;
+
+/// Returns the "tuple" value of the adapted tensor.
+inline UnrealizedConversionCastOp getTuple(Value tensor) {
+  return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
+}
+
+/// Packs the given values as a "tuple" value.
+inline Value genTuple(OpBuilder &builder, Location loc, Type tp,
+                      ValueRange values) {
+  return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
+      .getResult(0);
+}
+
+inline Value genTuple(OpBuilder &builder, Location loc,
+                      SparseTensorDescriptor desc) {
+  return genTuple(builder, loc, desc.getTensorType(), desc.getFields());
+}
+
+inline SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
+  auto tuple = getTuple(tensor);
+  return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
+}
+
+inline MutSparseTensorDescriptor
+getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl<Value> &fields) {
+  auto tuple = getTuple(tensor);
+  fields.assign(tuple.getInputs().begin(), tuple.getInputs().end());
+  return MutSparseTensorDescriptor(tuple.getResultTypes()[0], fields);
+}
+
+} // namespace sparse_tensor
+} // namespace mlir
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_

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

+//===- SparseTensorBuilder.cpp --------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "SparseTensorBuilder.h"
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+using namespace sparse_tensor;
+
+static SmallVector<Type, 1>
+getSparseTensorMetaDataFields(RankedTensorType rtp, unsigned indexBitwidth) {
+  MLIRContext *ctx = rtp.getContext();
+  SmallVector<Type, 1> result;
+
+  auto dimSizes = LLVM::LLVMArrayType::get(
+      ctx, IntegerType::get(ctx, indexBitwidth), rtp.getRank());
+  result.push_back(dimSizes);
+  return result;
+}
+
+static Type convertSparseTensorMetaData(RankedTensorType rtp,
+                                        unsigned indexBitwidth) {
+  return LLVM::LLVMStructType::getLiteral(
+      rtp.getContext(), getSparseTensorMetaDataFields(rtp, indexBitwidth));
+}
+
+Optional<LogicalResult>
+SparseTensorTypeToBufferConverter::convertSparseTensorType(
+    RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
+  auto enc = getSparseTensorEncoding(rtp);
+  if (!enc)
+    return std::nullopt;
+
+  foreachFieldAndTypeInSparseTensor(
+      rtp, getIndexTypeBitwidth(),
+      [&fields](Type fieldType, unsigned fieldIdx,
+                SparseTensorFieldKind /*fieldKind*/, unsigned /*dim*/,
+                DimLevelType /*dlt*/) -> bool {
+        assert(fieldIdx == fields.size());
+        fields.push_back(fieldType);
+        return true;
+      });
+  return success();
+}
+
+SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter(
+    unsigned indexWidth)
+    : indexWidth(indexWidth) {
+  addConversion([](Type type) { return type; });
+  addConversion([&](RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
+    return convertSparseTensorType(rtp, fields);
+  });
+
+  // Required by scf.for 1:N type conversion.
+  addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp,
+                              ValueRange inputs,
+                              Location loc) -> Optional<Value> {
+    if (!getSparseTensorEncoding(tp))
+      // Not a sparse tensor.
+      return std::nullopt;
+    // Sparse compiler knows how to cancel out these casts.
+    return genTuple(builder, loc, tp, inputs);
+  });
+}
+
+// MetaData
+constexpr uint64_t kDimSizePosInMetaData = 0;
+
+SparseTensorMetaData SparseTensorMetaData::undef(OpBuilder &builder,
+                                                 Location loc, Type metaType) {
+  Value metaData = builder.create<LLVM::UndefOp>(loc, metaType);
+  return SparseTensorMetaData(metaData);
+}
+
+Type SparseTensorMetaData::getMetaDataIndexType(Value data) {
+  return data.getType()
+      .cast<LLVM::LLVMStructType>()
+      .getBody()[kDimSizePosInMetaData]
+      .cast<LLVM::LLVMArrayType>()
+      .getElementType();
+}
+
+/// Builds IR inserting the pos-th size into the descriptor
+void SparseTensorMetaData::setDimSize(OpBuilder &builder, Location loc,
+                                      unsigned dim, Value size) {
+  value = builder.create<LLVM::InsertValueOp>(
+      loc, value, size, ArrayRef<int64_t>({kDimSizePosInMetaData, dim}));
+}
+
+/// Builds IR inserting the pos-th size into the descriptor
+Value SparseTensorMetaData::dimSize(OpBuilder &builder, Location loc,
+                                    unsigned dim) {
+  return builder.create<LLVM::ExtractValueOp>(
+      loc, value, ArrayRef<int64_t>({kDimSizePosInMetaData, dim}));
+}
+
+void sparse_tensor::foreachFieldInSparseTensor(
+    const SparseTensorEncodingAttr enc,
+    llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
+                            DimLevelType)>
+        callback) {
+  assert(enc);
+
+#define RETURN_ON_FALSE(idx, kind, dim, dlt)                                   \
+  if (!(callback(idx, kind, dim, dlt)))                                        \
+    return;
+
+  RETURN_ON_FALSE(kMemSizesIdx, SparseTensorFieldKind::MemSizes, -1u,
+                  DimLevelType::Undef);
+
+  static_assert(kDataFieldIdx == kMemSizesIdx + 1);
+  unsigned fieldIdx = kDataFieldIdx;
+  // Per-dimension storage.
+  for (unsigned r = 0, rank = enc.getDimLevelType().size(); r < rank; r++) {
+    // Dimension level types apply in order to the reordered dimension.
+    // As a result, the compound type can be constructed directly in the given
+    // order.
+    auto dlt = getDimLevelType(enc, r);
+    if (isCompressedDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PtrMemRef, r, dlt);
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
+    } else if (isSingletonDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
+    } else {
+      assert(isDenseDLT(dlt)); // no fields
+    }
+  }
+  // The values array.
+  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::ValMemRef, -1u,
+                  DimLevelType::Undef);
+
+  // Put metadata at the end.
+  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::MetaData, -1u,
+                  DimLevelType::Undef);
+
+#undef RETURN_ON_FALSE
+}
+
+void sparse_tensor::foreachFieldAndTypeInSparseTensor(
+    RankedTensorType rType, unsigned indexBitwidth,
+    llvm::function_ref<bool(Type, unsigned, SparseTensorFieldKind, unsigned,
+                            DimLevelType)>
+        callback) {
+  auto enc = getSparseTensorEncoding(rType);
+  assert(enc);
+  // Construct the basic types.
+  Type indexType = IndexType::get(enc.getContext());
+  Type idxType = enc.getIndexType();
+  Type ptrType = enc.getPointerType();
+  Type eltType = rType.getElementType();
+
+  Type metaDataType = convertSparseTensorMetaData(rType, indexBitwidth);
+  // memref<n x index> memSizes
+  Type memSizeType =
+      MemRefType::get({getNumDataFieldsFromEncoding(enc)}, indexType);
+  // memref<? x ptr>  pointers
+  Type ptrMemType = MemRefType::get({ShapedType::kDynamic}, ptrType);
+  // memref<? x idx>  indices
+  Type idxMemType = MemRefType::get({ShapedType::kDynamic}, idxType);
+  // memref<? x eltType> values
+  Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
+
+  foreachFieldInSparseTensor(
+      enc,
+      [metaDataType, memSizeType, ptrMemType, idxMemType, valMemType,
+       callback](unsigned fieldIdx, SparseTensorFieldKind fieldKind,
+                 unsigned dim, DimLevelType dlt) -> bool {
+        switch (fieldKind) {
+        case SparseTensorFieldKind::MetaData:
+          return callback(metaDataType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::MemSizes:
+          return callback(memSizeType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::PtrMemRef:
+          return callback(ptrMemType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::IdxMemRef:
+          return callback(idxMemType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::ValMemRef:
+          return callback(valMemType, fieldIdx, fieldKind, dim, dlt);
+        };
+        llvm_unreachable("unrecognized field kind");
+      });
+}
+
+unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  unsigned numFields = 0;
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
+                               numFields++;
+                               return true;
+                             });
+  return numFields;
+}
+
+unsigned
+sparse_tensor::getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  unsigned numFields = 0; // one value memref
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned fidx, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
+                               if (fidx >= kDataFieldIdx)
+                                 numFields++;
+                               return true;
+                             });
+  numFields -= 1; // the last field is MetaData field
+  assert(numFields == getNumFieldsFromEncoding(enc) - kDataFieldIdx - 1);
+  return numFields;
+}

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

 // the selected sparse tensor storage schemes. This pass provides an alternative
 // to the SparseTensorConversion pass, eliminating the dependence on a runtime
 // support library, and providing much more opportunities for subsequent
 // compiler optimization of the generated code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
+#include "SparseTensorBuilder.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
 using FuncGeneratorType =
     function_ref<void(OpBuilder &, ModuleOp, func::FuncOp, RankedTensorType)>;
 
 static constexpr const char kInsertFuncNamePrefix[] = "_insert_";
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
-/// Returns the "tuple" value of the adapted tensor.
-static UnrealizedConversionCastOp getTuple(Value tensor) {
-  return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
-}
-
-static SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
-  auto tuple = getTuple(tensor);
-  return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
-}
-
-static MutSparseTensorDescriptor
-getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl<Value> &fields) {
-  auto tuple = getTuple(tensor);
-  fields.assign(tuple.getInputs().begin(), tuple.getInputs().end());
-  return MutSparseTensorDescriptor(tuple.getResultTypes()[0], fields);
-}
-
-/// Packs the given values as a "tuple" value.
-static Value genTuple(OpBuilder &builder, Location loc, Type tp,
-                      ValueRange values) {
-  return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
-      .getResult(0);
-}
-
-static Value genTuple(OpBuilder &builder, Location loc,
-                      SparseTensorDescriptor desc) {
-  return builder
-      .create<UnrealizedConversionCastOp>(loc, desc.getTensorType(),
-                                          desc.getFields())
-      .getResult(0);
-}
-
 /// Flatten a list of operands that may contain sparse tensors.
 static void flattenOperands(ValueRange operands,
                             SmallVectorImpl<Value> &flattened) {
   // In case of
   // sparse_tensor, c, sparse_tensor
   // ==>
   // memref ..., c, memref ...
   for (auto operand : operands) {
static Optional<Value> sizeFromTensorAtDim(OpBuilder &builder, Location loc,
   // Access into static dimension can query original type directly.
   // Note that this is typically already done by DimOp's folding.
   auto shape = rtp.getShape();
   if (!ShapedType::isDynamic(shape[dim]))
     return constantIndex(builder, loc, shape[dim]);
 
   // Any other query can consult the dimSizes array at field DimSizesIdx,
   // accounting for the reordering applied to the sparse storage.
-  Value idx = constantIndex(builder, loc, toStoredDim(rtp, dim));
-  return builder.create<memref::LoadOp>(loc, desc.getDimSizesMemRef(), idx)
-      .getResult();
+  return toType(builder, loc,
+                desc.getDimSize(builder, loc, toStoredDim(rtp, dim)),
+                builder.getIndexType());
 }
 
 // Gets the dimension size at the given stored dimension 'd', either as a
 // constant for a static size, or otherwise dynamically through memSizes.
 Value sizeAtStoredDim(OpBuilder &builder, Location loc,
                       SparseTensorDescriptor desc, unsigned d) {
   RankedTensorType rtp = desc.getTensorType();
   unsigned dim = toOrigDim(rtp, d);
   auto shape = rtp.getShape();
   if (!ShapedType::isDynamic(shape[dim]))
     return constantIndex(builder, loc, shape[dim]);
 
-  return genLoad(builder, loc, desc.getDimSizesMemRef(),
-                 constantIndex(builder, loc, d));
+  return toType(builder, loc, desc.getDimSize(builder, loc, d),
+                builder.getIndexType());
 }
 
 static void createPushback(OpBuilder &builder, Location loc,
                            MutSparseTensorDescriptor desc, unsigned fidx,
                            Value value, Value repeat = Value()) {
   Type etp = desc.getElementType(fidx);
   Value field = desc.getField(fidx);
   Value newField = builder.create<PushBackOp>(
       loc, field.getType(), desc.getMemSizesMemRef(), field,
       toType(builder, loc, value, etp), APInt(64, getFieldMemSizesIndex(fidx)),
       repeat);
   desc.setField(fidx, newField);
 }
 
-/// Maps a sparse tensor type to the appropriate compounded buffers.
-static Optional<LogicalResult>
-convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
-  auto enc = getSparseTensorEncoding(type);
-  if (!enc)
-    return std::nullopt;
-
-  RankedTensorType rType = type.cast<RankedTensorType>();
-  foreachFieldAndTypeInSparseTensor(
-      rType,
-      [&fields](Type fieldType, unsigned fieldIdx,
-                SparseTensorFieldKind /*fieldKind*/, unsigned /*dim*/,
-                DimLevelType /*dlt*/) -> bool {
-        assert(fieldIdx == fields.size());
-        fields.push_back(fieldType);
-        return true;
-      });
-  return success();
-}
-
 /// Generates code that allocates a sparse storage scheme for given rank.
 static void allocSchemeForRank(OpBuilder &builder, Location loc,
                                MutSparseTensorDescriptor desc, unsigned r0) {
   RankedTensorType rtp = desc.getTensorType();
   unsigned rank = rtp.getShape().size();
   Value linear = constantIndex(builder, loc, 1);
   for (unsigned r = r0; r < rank; r++) {
     if (isCompressedDim(rtp, r)) {

 /// for all dynamic memrefs, the memory size is really the capacity of
 /// the "vector", while the actual size resides in the sizes array.
 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
 ///       on the required capacities (see heuristic variable).
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                               ValueRange dynSizes, bool enableInit,
-                              SmallVectorImpl<Value> &fields) {
+                              SmallVectorImpl<Value> &fields,
+                              unsigned indexBitwidth) {
   RankedTensorType rtp = type.cast<RankedTensorType>();
   Value heuristic = constantIndex(builder, loc, 16);
 
   foreachFieldAndTypeInSparseTensor(
-      rtp,
+      rtp, indexBitwidth,
       [&builder, &fields, loc, heuristic,
        enableInit](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
                    unsigned /*dim*/, DimLevelType /*dlt*/) -> bool {
         assert(fields.size() == fIdx);
-        auto memRefTp = fType.cast<MemRefType>();
         Value field;
         switch (fKind) {
-        case SparseTensorFieldKind::DimSizes:
+        case SparseTensorFieldKind::MetaData:
+          field = SparseTensorMetaData::undef(builder, loc, fType);
+          break;
         case SparseTensorFieldKind::MemSizes:
-          field = builder.create<memref::AllocOp>(loc, memRefTp);
+          field =
+              builder.create<memref::AllocOp>(loc, fType.cast<MemRefType>());
           break;
         case SparseTensorFieldKind::PtrMemRef:
         case SparseTensorFieldKind::IdxMemRef:
         case SparseTensorFieldKind::ValMemRef:
-          field =
-              createAllocation(builder, loc, memRefTp, heuristic, enableInit);
+          field = createAllocation(builder, loc, fType.cast<MemRefType>(),
+                                   heuristic, enableInit);
           break;
         }
         assert(field);
         fields.push_back(field);
         // Returns true to continue the iteration.
         return true;
       });
 
builder.create<linalg::FillOp>(
       loc, constantZero(builder, loc, builder.getIndexType()),
       desc.getMemSizesMemRef()); // zero memSizes
 
   Value ptrZero =
       constantZero(builder, loc, getSparseTensorEncoding(rtp).getPointerType());
   for (unsigned r = 0; r < rank; r++) {
     unsigned ro = toOrigDim(rtp, r);
     // Fills dim sizes array.
-    genStore(builder, loc, sizes[ro], desc.getDimSizesMemRef(),
-             constantIndex(builder, loc, r));
+    desc.setDimSize(builder, loc, r,
+                    toType(builder, loc, sizes[ro], desc.getIndexType()));
 
     // Pushes a leading zero to pointers memref.
     if (isCompressedDim(rtp, r))
       createPushback(builder, loc, desc, desc.getPtrMemRefIndex(r), ptrZero);
   }
   allocSchemeForRank(builder, loc, desc, /*rank=*/0);
 }
 
if (!index || !getSparseTensorEncoding(adaptor.getSource().getType()))
       return failure();
 
     auto desc = getDescriptorFromTensorTuple(adaptor.getSource());
     auto sz = sizeFromTensorAtDim(rewriter, op.getLoc(), desc, *index);
 
     if (!sz)
       return failure();
 
-    rewriter.replaceOp(op, *sz);
+    rewriter.replaceOp(
+        op, toType(rewriter, op.getLoc(), *sz, rewriter.getIndexType()));
     return success();
   }
 };
 
 /// Sparse codegen rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;

 };
 
 /// Sparse codgen rule for the alloc operator.
 class SparseTensorAllocConverter
     : public OpConversionPattern<bufferization::AllocTensorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   SparseTensorAllocConverter(TypeConverter &typeConverter, MLIRContext *context,
-                             bool enableInit)
+                             bool enableInit, unsigned indexBitwidth)
       : OpConversionPattern(typeConverter, context),
-        enableBufferInitialization(enableInit) {}
+        enableBufferInitialization(enableInit), indexBitwidth(indexBitwidth) {}
+
   LogicalResult
   matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     RankedTensorType resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     if (op.getCopy())
       return rewriter.notifyMatchFailure(op, "tensor copy not implemented");
 
     // Construct allocation for each field.
     Location loc = op.getLoc();
     SmallVector<Value> fields;
     createAllocFields(rewriter, loc, resType, adaptor.getOperands(),
-                      enableBufferInitialization, fields);
+                      enableBufferInitialization, fields, indexBitwidth);
     // Replace operation with resulting memrefs.
     rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields));
     return success();
   }
 
 private:
   bool enableBufferInitialization;
+  unsigned indexBitwidth;
 };
 
 /// Sparse codegen rule for the dealloc operator.
 class SparseTensorDeallocConverter
     : public OpConversionPattern<bufferization::DeallocTensorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto enc = getSparseTensorEncoding(op.getTensor().getType());
     if (!enc)
       return failure();
 
     // Replace the sparse tensor deallocation with field deallocations.
     Location loc = op.getLoc();
-    auto tuple = getTuple(adaptor.getTensor());
-    for (auto input : tuple.getInputs())
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    for (auto input : desc.getMemRefFields())
       // Deallocate every buffer used to store the sparse tensor handler.
       rewriter.create<memref::DeallocOp>(loc, input);
 
     rewriter.eraseOp(op);
     return success();
   }
 };
 
rewriter.replaceOpWithNewOp<memref::LoadOp>(op, desc.getMemSizesMemRef(),
                                                 field);
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Sparse tensor type conversion into an actual buffer.
-//===----------------------------------------------------------------------===//
-
-mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
-  addConversion([](Type type) { return type; });
-  addConversion(convertSparseTensorType);
-
-  // Required by scf.for 1:N type conversion.
-  addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp,
-                              ValueRange inputs,
-                              Location loc) -> Optional<Value> {
-    if (!getSparseTensorEncoding(tp))
-      // Not a sparse tensor.
-      return std::nullopt;
-    // Sparse compiler knows how to cancel out these casts.
-    return genTuple(builder, loc, tp, inputs);
-  });
-}
-
-//===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
-    TypeConverter &typeConverter, RewritePatternSet &patterns,
-    bool enableBufferInitialization) {
+    SparseTensorTypeToBufferConverter &typeConverter,
+    RewritePatternSet &patterns, bool enableBufferInitialization) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorDeallocConverter,
                SparseTensorLoadConverter, SparseExpandConverter,
                SparseCompressConverter, SparseInsertConverter,
                SparseToPointersConverter, SparseToIndicesConverter,
                SparseToValuesConverter, SparseConvertConverter,
                SparseNumberOfEntriesConverter>(typeConverter,
                                                patterns.getContext());
-  patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
-                                           enableBufferInitialization);
+
+  patterns.add<SparseTensorAllocConverter>(
+      typeConverter, patterns.getContext(), enableBufferInitialization,
+      typeConverter.getIndexTypeBitwidth());
 }

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

 //===- SparseTensorPasses.cpp - Pass for autogen sparse tensor code -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+#include "SparseTensorBuilder.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
struct SparseTensorCodegenPass
   SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;
   SparseTensorCodegenPass(bool enableInit) {
     enableBufferInitialization = enableInit;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    SparseTensorTypeToBufferConverter converter;
+    // TODO: make the index type bitwidth an option.
+    SparseTensorTypeToBufferConverter converter(64u);
     ConversionTarget target(*ctx);
     // Most ops in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
     target.addLegalOp<SortOp>();
     target.addLegalOp<SortCooOp>();
     target.addLegalOp<PushBackOp>();
     // All dynamic rules below accept new function, call, return, and various
     // tensor and bufferization operations as legal output of the rewriting
target.addDynamicallyLegalOp<bufferization::AllocTensorOp>(
         });
     target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(
         [&](bufferization::DeallocTensorOp op) {
           return converter.isLegal(op.getTensor().getType());
         });
     // The following operations and dialects may be introduced by the
     // codegen rules, and are therefore marked as legal.
     target.addLegalOp<linalg::FillOp>();
-    target.addLegalDialect<
-        arith::ArithDialect, bufferization::BufferizationDialect,
-        complex::ComplexDialect, memref::MemRefDialect, scf::SCFDialect>();
+    target.addLegalDialect<arith::ArithDialect,
+                           bufferization::BufferizationDialect,
+                           complex::ComplexDialect, memref::MemRefDialect,
+                           scf::SCFDialect, LLVM::LLVMDialect>();
     target.addLegalOp<UnrealizedConversionCastOp>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
     populateSparseTensorCodegenPatterns(converter, patterns,
                                         enableBufferInitialization);

mlir/test/Dialect/SparseTensor/codegen.mlir

 }>
 
 #Dense3D = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "dense" ],
   dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
 }>
 
 // CHECK-LABEL: func @sparse_nop(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*]]: !llvm.struct<(array<1 x i64>)>) 
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
   return %arg0 : tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_multi_ret(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<1xindex>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<3xindex>,
-//  CHECK-SAME: %[[A7:.*7]]: memref<?xi32>,
-//  CHECK-SAME: %[[A8:.*8]]: memref<?xi64>,
-//  CHECK-SAME: %[[A9:.*9]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<3xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xi32>,
+//  CHECK-SAME: %[[A7:.*7]]: memref<?xi64>,
+//  CHECK-SAME: %[[A8:.*8]]: memref<?xf64>,
+//  CHECK-SAME: %[[A9:.*9]]: !llvm.struct<(array<1 x i64>)>) 
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]], %[[A8]], %[[A9]] :
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>,
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>,
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 func.func @sparse_nop_multi_ret(%arg0: tensor<?xf64, #SparseVector>,
                                 %arg1: tensor<?xf64, #SparseVector>) ->
                                 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
   return %arg0, %arg1 : tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_call(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<1xindex>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<3xindex>,
-//  CHECK-SAME: %[[A7:.*7]]: memref<?xi32>,
-//  CHECK-SAME: %[[A8:.*8]]: memref<?xi64>,
-//  CHECK-SAME: %[[A9:.*9]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<3xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xi32>,
+//  CHECK-SAME: %[[A7:.*7]]: memref<?xi64>,
+//  CHECK-SAME: %[[A8:.*8]]: memref<?xf64>,
+//  CHECK-SAME: %[[A9:.*9]]: !llvm.struct<(array<1 x i64>)>) 
 //       CHECK: %[[T:.*]]:10 = call @sparse_nop_multi_ret(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]], %[[A8]], %[[A9]])
 //       CHECK: return %[[T]]#0, %[[T]]#1, %[[T]]#2, %[[T]]#3, %[[T]]#4, %[[T]]#5, %[[T]]#6, %[[T]]#7, %[[T]]#8, %[[T]]#9 :
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>,
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>,
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 func.func @sparse_nop_call(%arg0: tensor<?xf64, #SparseVector>,
                            %arg1: tensor<?xf64, #SparseVector>) ->
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
   %1, %2 = call @sparse_nop_multi_ret(%arg0, %arg1) :
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ->
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
   return %1, %2: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_cast(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf32>)
+//  CHECK-SAME: %[[A0:.*]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf32>,
+//  CHECK-SAME: %[[A4:.*]]: !llvm.struct<(array<1 x i64>)>) 
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
 func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_cast_3d(
-//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xf32>)
+//  CHECK-SAME: %[[A0:.*]]: memref<1xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xf32>,
+//  CHECK-SAME: %[[A2:.*]]: !llvm.struct<(array<3 x i64>)>) 
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]] :
-//  CHECK-SAME:   memref<3xindex>, memref<1xindex>, memref<?xf32>
+//  CHECK-SAME:   memref<1xindex>, memref<?xf32>, !llvm.struct<(array<3 x i64>)>
 func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor<?x?x?xf32, #Dense3D> {
   %0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor<?x?x?xf32, #Dense3D>
   return %0 : tensor<?x?x?xf32, #Dense3D>
 }
 
 // CHECK-LABEL: func @sparse_dense_2d(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*]]: memref<1xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A2:.*]]: !llvm.struct<(array<2 x i64>)>) {
 //       CHECK: return
 func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) {
   return
 }
 
 // CHECK-LABEL: func @sparse_row(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>) {
 //       CHECK: return
 func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {
   return
 }
 
 // CHECK-LABEL: func @sparse_csr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*]]: !llvm.struct<(array<2 x i64>)>) {
 //       CHECK: return
 func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {
   return
 }
 
 // CHECK-LABEL: func @sparse_dcsr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*0]]: memref<5xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: !llvm.struct<(array<2 x i64>)>)
 //       CHECK: return
 func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {
   return
 }
 
 //
 // Querying for dimension 1 in the tensor type can immediately
 // fold using the original static dimension sizes.
 //
 // CHECK-LABEL: func @sparse_dense_3d(
-//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*]]: memref<1xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A2:.*]]: !llvm.struct<(array<3 x i64>)>)
 //       CHECK: %[[C:.*]] = arith.constant 20 : index
 //       CHECK: return %[[C]] : index
 func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index {
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D>
   return %0 : index
 }
 
 //
 // Querying for dimension 1 in the tensor type needs to be permuted
 // into querying for dimension 2 in the stored sparse tensor scheme,
 // since the latter honors the dimOrdering.
 //
 // CHECK-LABEL: func @sparse_dense_3d_dyn(
-//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
-//       CHECK: %[[C:.*]] = arith.constant 2 : index
-//       CHECK: %[[L:.*]] = memref.load %[[A0]][%[[C]]] : memref<3xindex>
-//       CHECK: return %[[L]] : index
+//  CHECK-SAME: %[[A0:.*]]: memref<1xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A2:.*]]: !llvm.struct<(array<3 x i64>)>)
+//       CHECK: %[[A3:.*]] = llvm.extractvalue %[[A2]][0, 2] 
+//       CHECK: %[[A4:.*]] = arith.index_cast %[[A3]] : i64 to index
+//       CHECK: return %[[A4]] : index
 func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index {
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #Dense3D>
   return %0 : index
 }
 
 // CHECK-LABEL: func @sparse_pointers_dcsr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
-//       CHECK: return %[[A4]] : memref<?xi32>
+//  CHECK-SAME: %[[A0:.*0]]: memref<5xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: !llvm.struct<(array<2 x i64>)>
+//       CHECK: return %[[A3]] : memref<?xi32>
 func.func @sparse_pointers_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> {
   %0 = sparse_tensor.pointers %arg0 { dimension = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi32>
   return %0 : memref<?xi32>
 }
 
 // CHECK-LABEL: func @sparse_indices_dcsr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
-//       CHECK: return %[[A5]] : memref<?xi64>
+//  CHECK-SAME: %[[A0:.*0]]: memref<5xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: !llvm.struct<(array<2 x i64>)>
+//       CHECK: return %[[A4]] : memref<?xi64>
 func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {
   %0 = sparse_tensor.indices %arg0 { dimension = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi64>
   return %0 : memref<?xi64>
 }
 
 // CHECK-LABEL: func @sparse_values_dcsr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
-//       CHECK: return %[[A6]] : memref<?xf64>
+//  CHECK-SAME: %[[A0:.*0]]: memref<5xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: !llvm.struct<(array<2 x i64>)>)
+//       CHECK: return %[[A5]] : memref<?xf64>
 func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // CHECK-LABEL: func @sparse_noe(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
+//  CHECK-SAME: %[[A0:.*]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*]]: !llvm.struct<(array<1 x i64>)>)
 //       CHECK: %[[C2:.*]] = arith.constant 2 : index
-//       CHECK: %[[NOE:.*]] = memref.load %[[A1]][%[[C2]]] : memref<3xindex>
+//       CHECK: %[[NOE:.*]] = memref.load %[[A0]][%[[C2]]] : memref<3xindex>
 //       CHECK: return %[[NOE]] : index
 func.func @sparse_noe(%arg0: tensor<128xf64, #SparseVector>) -> index {
   %0 = sparse_tensor.number_of_entries %arg0 : tensor<128xf64, #SparseVector>
   return %0 : index
 }
 
 // CHECK-LABEL: func @sparse_dealloc_csr(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
-//       CHECK: memref.dealloc %[[A0]] : memref<2xindex>
-//       CHECK: memref.dealloc %[[A1]] : memref<3xindex>
-//       CHECK: memref.dealloc %[[A2]] : memref<?xi32>
-//       CHECK: memref.dealloc %[[A3]] : memref<?xi64>
-//       CHECK: memref.dealloc %[[A4]] : memref<?xf64>
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>)
+//       CHECK: memref.dealloc %[[A0]] : memref<3xindex>
+//       CHECK: memref.dealloc %[[A1]] : memref<?xi32>
+//       CHECK: memref.dealloc %[[A2]] : memref<?xi64>
+//       CHECK: memref.dealloc %[[A3]] : memref<?xf64>
 //       CHECK: return
 func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {
   bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>
   return
 }
 
 // CHECK-LABEL: func @sparse_alloc_csc(
-//  CHECK-SAME: %[[A:.*]]: index) ->
-//  CHECK-SAME: memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
-//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
-//       CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>
-//       CHECK: %[[T1:.*]] = memref.alloc() : memref<3xindex>
-//       CHECK: %[[T2:.*]] = memref.alloc() : memref<16xindex>
-//       CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<16xindex> to memref<?xindex>
-//       CHECK: %[[T4:.*]] = memref.alloc() : memref<16xindex>
-//       CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<16xindex> to memref<?xindex>
-//       CHECK: %[[T6:.*]] = memref.alloc() : memref<16xf64>
-//       CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<16xf64> to memref<?xf64>
-//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T1]] : memref<3xindex>)
-//       CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
-//       CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
-//       CHECK: %[[P0:.*]] = sparse_tensor.push_back %[[T1]], %[[T3]]
-//       CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[T1]], %[[P0]]
-//       CHECK: return %[[T0]], %[[T1]], %[[P1]], %[[T5]], %[[T7]] :
-//  CHECK-SAME:   memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//  CHECK-SAME: %[[A0:.*]]: index)
+//   CHECK-DAG: %[[A1:.*]] = arith.constant 10 : i64
+//   CHECK-DAG: %[[A2:.*]] = arith.constant 0 : index
+//       CHECK: %[[A3:.*]] = memref.alloc() : memref<3xindex>
+//       CHECK: %[[A4:.*]] = memref.alloc() : memref<16xindex>
+//       CHECK: %[[A5:.*]] = memref.cast %[[A4]] : memref<16xindex> to memref<?xindex>
+//       CHECK: %[[A6:.*]] = memref.alloc() : memref<16xindex>
+//       CHECK: %[[A7:.*]] = memref.cast %[[A6]] : memref<16xindex> to memref<?xindex>
+//       CHECK: %[[A8:.*]] = memref.alloc() : memref<16xf64>
+//       CHECK: %[[A9:.*]] = memref.cast %[[A8]] : memref<16xf64> to memref<?xf64>
+//       CHECK: %[[A10:.*]] = llvm.mlir.undef : !llvm.struct<(array<2 x i64>)>
+//       CHECK: linalg.fill ins(%[[A2]] : index) outs(%[[A3]] : memref<3xindex>)
+//       CHECK: %[[A11:.*]] = arith.index_cast %[[A0]] : index to i64
+//       CHECK: %[[A12:.*]] = llvm.insertvalue %[[A11]], %[[A10]][0, 0]
+//       CHECK: %[[A13:.*]] = llvm.insertvalue %[[A1]], %[[A12]][0, 1]
+//       CHECK: %[[A14:.*]] = sparse_tensor.push_back %[[A3]], %[[A5]], %[[A2]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index
+//       CHECK: %[[A15:.*]] = sparse_tensor.push_back %[[A3]], %[[A14]], %[[A2]], %[[A0]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index, index
+//       CHECK: return %[[A3]], %[[A15]], %[[A7]], %[[A9]], %[[A13]] : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {
   %0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>
   %1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>
   return %1 : tensor<10x?xf64, #CSC>
 }
 
-// CHECK-LABEL: func @sparse_alloc_3d() ->
-//  CHECK-SAME: memref<3xindex>, memref<1xindex>, memref<?xf64>
-//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
-//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
-//   CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
-//   CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index
-//   CHECK-DAG: %[[C6000:.*]] = arith.constant 6000 : index
-//       CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>
-//       CHECK: %[[A1:.*]] = memref.alloc() : memref<1xindex>
-//       CHECK: %[[AV:.*]] = memref.alloc() : memref<16xf64>
-//       CHECK: %[[A2:.*]] = memref.cast %[[AV]] : memref<16xf64> to memref<?xf64>
-//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[A1]] : memref<1xindex>)
-//       CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
-//       CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
-//       CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
-//       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A1]], %[[A2]], %[[F0]], %[[C6000]]
-//       CHECK: return %[[A0]], %[[A1]], %[[P]] :
-//  CHECK-SAME:   memref<3xindex>, memref<1xindex>, memref<?xf64>
+// CHECK-LABEL: func @sparse_alloc_3d()
+//   CHECK-DAG: %[[A0:.*]] = arith.constant 6000 : index
+//   CHECK-DAG: %[[A1:.*]] = arith.constant 20 : i64
+//   CHECK-DAG: %[[A2:.*]] = arith.constant 10 : i64
+//   CHECK-DAG: %[[A3:.*]] = arith.constant 30 : i64
+//   CHECK-DAG: %[[A4:.*]] = arith.constant 0.000000e+00 : f64
+//   CHECK-DAG: %[[A5:.*]] = arith.constant 0 : index
+//       CHECK: %[[A6:.*]] = memref.alloc() : memref<1xindex>
+//       CHECK: %[[A7:.*]] = memref.alloc() : memref<16xf64>
+//       CHECK: %[[A8:.*]] = memref.cast %[[A7]] : memref<16xf64> to memref<?xf64>
+//       CHECK: %[[A9:.*]] = llvm.mlir.undef : !llvm.struct<(array<3 x i64>)>
+//       CHECK: linalg.fill ins(%[[A5]] : index) outs(%[[A6]] : memref<1xindex>)
+//       CHECK: %[[A10:.*]] = llvm.insertvalue %[[A3]], %[[A9]][0, 0]
+//       CHECK: %[[A11:.*]] = llvm.insertvalue %[[A2]], %[[A10]][0, 1]
+//       CHECK: %[[A12:.*]] = llvm.insertvalue %[[A1]], %[[A11]][0, 2]
+//       CHECK: %[[A13:.*]] = sparse_tensor.push_back %[[A6]], %[[A8]], %[[A4]], %[[A0]] {idx = 0 : index} : memref<1xindex>, memref<?xf64>, f64, index
+//       CHECK: return %[[A6]], %[[A13]], %[[A12]] : memref<1xindex>, memref<?xf64>, !llvm.struct<(array<3 x i64>)>
 func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
   %0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>
   %1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>
   return %1 : tensor<10x20x30xf64, #Dense3D>
 }
 
 // CHECK-LABEL: func.func @sparse_expansion1()
 //       CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
func.func @sparse_expansion2() -> memref<?xindex> {
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
   return %added : memref<?xindex>
 }
 
 // CHECK-LABEL: func.func @sparse_expansion3(
 //  CHECK-SAME: %[[D0:.*]]: index,
 //  CHECK-SAME: %{{.*}}: index) -> memref<?xindex> {
-//       CHECK: %[[C1:.*]] = arith.constant 1 : index
-//       CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex>
-//       CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
-//       CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
-//       CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64>
-//       CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1>
-//       CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex>
+//       CHECK: %[[V:.*]] = memref.alloc(%[[D0]]) : memref<?xf64>
+//       CHECK: %[[B:.*]] = memref.alloc(%[[D0]]) : memref<?xi1>
+//       CHECK: %[[D:.*]] = memref.alloc(%[[D0]]) : memref<?xindex>
 //       CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>)
 //       CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
 //       CHECK: return %[[D]] : memref<?xindex>
 func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
   %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
   return %added : memref<?xindex>
 }
 
 // CHECK-LABEL: func.func private @_insert_C_100_f64_0_0(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[PV]]
+//       CHECK: %[[A22:.*]] = sparse_tensor.push_back %[[A0]], %[[A3]], %[[A6]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
+//       CHECK: return %[[A0]], %[[A1]], %{{.*}}, %[[A22]], %[[A4]] : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 //
 // CHECK-LABEL: func @sparse_compression_1d(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
 //  CHECK-SAME: %[[A6:.*6]]: memref<?xi1>,
 //  CHECK-SAME: %[[A7:.*7]]: memref<?xindex>,
 //  CHECK-SAME: %[[A8:.*8]]: index)
-//   CHECK-DAG: %[[B0:.*]] = arith.constant false
-//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
-//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A9:.*]] = arith.constant false
+//   CHECK-DAG: %[[A10:.*]] = arith.constant 0.000000e+00 : f64
+//   CHECK-DAG: %[[A11:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A12:.*]] = arith.constant 0 : index
 //       CHECK: sparse_tensor.sort %[[A8]], %[[A7]] : memref<?xindex>
-//       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
-//  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-//       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
-//       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_C_100_f64_0_0(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[INDEX]], %[[VAL]])
-//       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
-//       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//       CHECK: %[[A13:.*]]:5 = scf.for %[[A14:.*]] = %[[A12]] to %[[A8]] step %[[A11]] iter_args(%[[A15:.*]] = %[[A0]], %[[A16:.*]] = %[[A1]], %[[A17:.*]] = %[[A2]], %[[A18:.*]] = %[[A3]], %[[A19:.*]] = %[[A4]]) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>) {
+//       CHECK:   %[[A20:.*]] = memref.load %[[A7]]{{\[}}%[[A14]]] : memref<?xindex>
+//       CHECK:   %[[A21:.*]] = memref.load %[[A5]]{{\[}}%[[A20]]] : memref<?xf64>
+//       CHECK:   %[[A22:.*]]:5 = func.call @_insert_C_100_f64_0_0(%[[A15]], %[[A16]], %[[A17]], %[[A18]], %[[A19]], %[[A20]], %[[A21]]) : (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>, index, f64) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>)
+//       CHECK:   memref.store %[[A10]], %[[A5]]{{\[}}%[[A20]]] : memref<?xf64>
+//       CHECK:   memref.store %[[A9]], %[[A6]]{{\[}}%[[A20]]] : memref<?xi1>
+//       CHECK:   scf.yield %[[A22]]#0, %[[A22]]#1, %[[A22]]#2, %[[A22]]#3, %[[A22]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
-//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//       CHECK: return %[[A23:.*]]#0, %[[A23]]#1, %[[A23]]#2, %[[A23]]#3, %[[A23]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 func.func @sparse_compression_1d(%tensor: tensor<100xf64, #SV>,
                                  %values: memref<?xf64>,
                                  %filled: memref<?xi1>,
                                  %added: memref<?xindex>,
                                  %count: index) -> tensor<100xf64, #SV> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<100xf64, #SV>
   %1 = sparse_tensor.load %0 hasInserts : tensor<100xf64, #SV>
   return %1 : tensor<100xf64, #SV>
 }
 
 // CHECK-LABEL: func.func private @_insert_D_C_8_8_f64_64_32(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: index,
 //  CHECK-SAME: %[[A7:.*7]]: f64)
-//       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[PV]]
-//
+//       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A0]], %[[A3]], %[[A7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
+//       CHECK: return %[[A0]], %[[A1]], %{{.*}}, %[[PV]], %[[A4]]
+///
 // CHECK-LABEL: func @sparse_compression(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
 //  CHECK-SAME: %[[A6:.*6]]: memref<?xi1>,
 //  CHECK-SAME: %[[A7:.*7]]: memref<?xindex>,
 //  CHECK-SAME: %[[A8:.*8]]: index,
 //  CHECK-SAME: %[[A9:.*9]]: index)
-//   CHECK-DAG: %[[B0:.*]] = arith.constant false
-//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
-//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A10:.*]] = arith.constant 0 : i32
+//   CHECK-DAG: %[[A11:.*]] = arith.constant false
+//   CHECK-DAG: %[[A12:.*]] = arith.constant 0.000000e+00 : f64
+//   CHECK-DAG: %[[A13:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A14:.*]] = arith.constant 0 : index
 //       CHECK: sparse_tensor.sort %[[A8]], %[[A7]] : memref<?xindex>
-//       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
-//  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>) {
-//       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
-//       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_D_C_8_8_f64_64_32(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
-//       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
-//       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//       CHECK: %[[R:.*]]:5 = scf.for %[[A16:.*]] = %[[A14]] to %[[A8]] step %[[A13]] iter_args(%[[A17:.*]] = %[[A0]], %[[A18:.*]] = %[[A1]], %[[A19:.*]] = %[[A2]], %[[A20:.*]] = %[[A3]], %[[A21:.*]] = %[[A4]]) -> (memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+//       CHECK:   %[[A22:.*]] = memref.load %[[A7]]{{\[}}%[[A16]]] : memref<?xindex>
+//       CHECK:   %[[A23:.*]] = memref.load %[[A5]]{{\[}}%[[A22]]] : memref<?xf64>
+//       CHECK:   %[[A24:.*]]:5 = func.call @_insert_D_C_8_8_f64_64_32(%[[A17]], %[[A18]], %[[A19]], %[[A20]], %[[A21]], %[[A9]], %[[A22]], %[[A23]])
+//       CHECK:   memref.store %[[A12]], %[[A5]]{{\[}}%[[A22]]] : memref<?xf64>
+//       CHECK:   memref.store %[[A11]], %[[A6]]{{\[}}%[[A22]]] : memref<?xi1>
+//       CHECK:   scf.yield %[[A24]]#0, %[[A24]]#1, %[[A24]]#2, %[[A24]]#3, %[[A24]]#4 : memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
-//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
-//  CHECK-SAME:   memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4 : memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 func.func @sparse_compression(%tensor: tensor<8x8xf64, #CSR>,
                               %values: memref<?xf64>,
                               %filled: memref<?xi1>,
                               %added: memref<?xindex>,
                               %count: index,
                               %i: index) -> tensor<8x8xf64, #CSR> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #CSR>
   %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #CSR>
   return %1 : tensor<8x8xf64, #CSR>
 }
 
 // CHECK-LABEL: func.func private @_insert_D_C_8_8_f64_0_0(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: index,
 //  CHECK-SAME: %[[A7:.*7]]: f64)
-//       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[PV]]
+//       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A0]], %[[A3]], %[[A7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
+//       CHECK: return %[[A0]], %[[A1]], %{{.*}}, %[[PV]], %[[A4]]
 //
 // CHECK-LABEL: func @sparse_compression_unordered(
-//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<2 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
 //  CHECK-SAME: %[[A6:.*6]]: memref<?xi1>,
 //  CHECK-SAME: %[[A7:.*7]]: memref<?xindex>,
 //  CHECK-SAME: %[[A8:.*8]]: index,
 //  CHECK-SAME: %[[A9:.*9]]: index)
-//   CHECK-DAG: %[[B0:.*]] = arith.constant false
-//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
-//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A10:.*]] = arith.constant false
+//   CHECK-DAG: %[[A11:.*]] = arith.constant 0.000000e+00 : f64
+//   CHECK-DAG: %[[A12:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[A13:.*]] = arith.constant 1 : index
 //   CHECK-NOT: sparse_tensor.sort
-//       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
-//  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-//       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
-//       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_D_C_8_8_f64_0_0(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
-//       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
-//       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//       CHECK: %[[R:.*]]:5 = scf.for %[[A15:.*]] = %[[A12]] to %[[A8]] step %[[A13]] iter_args(%[[A16:.*]] = %[[A0]], %[[A17:.*]] = %[[A1]], %[[A18:.*]] = %[[A2]], %[[A19:.*]] = %[[A3]], %[[A20:.*]] = %[[A4]]) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+//       CHECK:   %[[A21:.*]] = memref.load %[[A7]]{{\[}}%[[A15]]] : memref<?xindex>
+//       CHECK:   %[[A22:.*]] = memref.load %[[A5]]{{\[}}%[[A21]]] : memref<?xf64>
+//       CHECK:   %[[A23:.*]]:5 = func.call @_insert_D_C_8_8_f64_0_0(%[[A16]], %[[A17]], %[[A18]], %[[A19]], %[[A20]], %[[A9]], %[[A21]], %[[A22]]) : (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>, index, index, f64) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>)
+//       CHECK:   memref.store %[[A11]], %[[A5]]{{\[}}%[[A21]]] : memref<?xf64>
+//       CHECK:   memref.store %[[A10]], %[[A6]]{{\[}}%[[A21]]] : memref<?xi1>
+//       CHECK:   scf.yield %[[A23]]#0, %[[A23]]#1, %[[A23]]#2, %[[A23]]#3, %[[A23]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
-//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
-//  CHECK-SAME:   memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 func.func @sparse_compression_unordered(%tensor: tensor<8x8xf64, #UCSR>,
                                         %values: memref<?xf64>,
                                         %filled: memref<?xi1>,
                                         %added: memref<?xindex>,
                                         %count: index,
                                         %i: index) -> tensor<8x8xf64, #UCSR> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #UCSR>
   %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #UCSR>
   return %1 : tensor<8x8xf64, #UCSR>
 }
 
 // CHECK-LABEL: func.func private @_insert_C_128_f64_0_0(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
-//  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]]  {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[P]] :
-//       CHECK: func @sparse_insert(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: f64) 
+//       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A0]], %[[A3]], %[[A6]]  {idx = 2 : index}
+//       CHECK: return %[[A0]], %[[A1]], %{{.*}}, %[[P]], %[[A4]]
+
+// CHECK-LABEL: func @sparse_insert(
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
-//  CHECK-SAME: %[[A6:.*6]]: f64)
+//  CHECK-SAME: %[[A6:.*6]]: f64) 
 //       CHECK: %[[R:.*]]:5 = call @_insert_C_128_f64_0_0(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 func.func @sparse_insert(%arg0: tensor<128xf64, #SV>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SV> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SV>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SV>
   return %1 : tensor<128xf64, #SV>
 }
 
 // CHECK-LABEL: func.func private @_insert_C_128_f64_64_32(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]]  {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[P]] :
-//       CHECK: func @sparse_insert_typed(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A0]], %[[A3]], %[[A6]]  {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
+//       CHECK: return %[[A0]], %[[A1]], %{{.*}}, %[[P]], %[[A4]]
+
+// CHECK-LABEL: func @sparse_insert_typed(
+//  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*4]]: !llvm.struct<(array<1 x i64>)>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
-//  CHECK-SAME: %[[A6:.*6]]: f64)
+//  CHECK-SAME: %[[A6:.*6]]: f64) 
 //       CHECK: %[[R:.*]]:5 = call @_insert_C_128_f64_64_32(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
-//  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 func.func @sparse_insert_typed(%arg0: tensor<128xf64, #SparseVector>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SparseVector>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SparseVector>
   return %1 : tensor<128xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func.func @sparse_nop_convert(
-//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xf32>)
+//  CHECK-SAME: %[[A0:.*]]: memref<3xindex>,
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf32>,
+//  CHECK-SAME: %[[A4:.*]]: !llvm.struct<(array<1 x i64>)>) 
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
-//  CHECK-SAME: memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
+//  CHECK-SAME:   memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>, !llvm.struct<(array<1 x i64>)>
 func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
-}
+}
 

mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir

 // RUN: mlir-opt %s --sparse-tensor-codegen=enable-buffer-initialization=true  --canonicalize --cse | FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 
 // CHECK-LABEL: func @sparse_alloc_sparse_vector(
 //  CHECK-SAME: %[[A:.*]]: index) ->
-//  CHECK-SAME: memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//  CHECK-SAME: memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<1 x i64>)>
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
-//       CHECK: %[[T0:.*]] = memref.alloc() : memref<1xindex>
 //       CHECK: %[[T1:.*]] = memref.alloc() : memref<3xindex>
 //       CHECK: %[[T2:.*]] = memref.alloc() : memref<16xindex>
 //       CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<16xindex> to memref<?xindex>
 //       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T2]] : memref<16xindex>)
 //       CHECK: %[[T4:.*]] = memref.alloc() : memref<16xindex>
 //       CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<16xindex> to memref<?xindex>
 //       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T4]] : memref<16xindex>)
 //       CHECK: %[[T6:.*]] = memref.alloc() : memref<16xf64>
 //       CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<16xf64> to memref<?xf64>
 //       CHECK: linalg.fill ins(%[[F0]] : f64) outs(%[[T6]] : memref<16xf64>)
+//       CHECK: %[[T11:.*]] = llvm.mlir.undef : !llvm.struct<(array<1 x i64>)>
 //       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T1]] : memref<3xindex>)
-//       CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<1xindex>
+//       CHECK: %[[T12:.*]] = arith.index_cast %[[A]] : index to i64
+//       CHECK: %[[MD:.*]] = llvm.insertvalue %[[T12]], %[[T11]][0, 0] 
 //       CHECK: %[[P0:.*]] = sparse_tensor.push_back %[[T1]], %[[T3]]
 //       CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[T1]], %[[P0]]
-//       CHECK: return %[[T0]], %[[T1]], %[[P1]], %[[T5]], %[[T7]] :
+//       CHECK: return %[[T1]], %[[P1]], %[[T5]], %[[T7]], %[[MD]]
 func.func @sparse_alloc_sparse_vector(%arg0: index) -> tensor<?xf64, #SV> {
   %0 = bufferization.alloc_tensor(%arg0) : tensor<?xf64, #SV>
   %1 = sparse_tensor.load %0 : tensor<?xf64, #SV>
   return %1 : tensor<?xf64, #SV>
 }

mlir/test/Dialect/SparseTensor/scf_1_N_conversion.mlir

 // RUN: mlir-opt %s -sparse-tensor-codegen -cse | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
-// CHECK-LABEL:  func @for(
-// CHECK-SAME:             %[[DIM_SIZE:.*0]]: memref<1xindex>,
-// CHECK-SAME:             %[[MEM_SIZE:.*1]]: memref<3xindex>,
-// CHECK-SAME:             %[[POINTER:.*2]]: memref<?xindex>,
-// CHECK-SAME:             %[[INDICES:.*3]]: memref<?xindex>,
-// CHECK-SAME:             %[[VALUE:.*4]]: memref<?xf32>,
-// CHECK-SAME:             %[[LB:.*5]]: index,
-// CHECK-SAME:             %[[UB:.*6]]: index,
-// CHECK-SAME:             %[[STEP:.*7]]: index)
-// CHECK:          %[[OUT:.*]]:5 = scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] iter_args(
-// CHECK-SAME:       %[[SIZE:.*]] = %[[DIM_SIZE]],
-// CHECK-SAME:       %[[MEM:.*]] = %[[MEM_SIZE]],
-// CHECK-SAME:       %[[PTR:.*]] = %[[POINTER]],
-// CHECK-SAME:       %[[IDX:.*]] = %[[INDICES]],
-// CHECK-SAME:       %[[VAL:.*]] = %[[VALUE]])
-// CHECK:            scf.yield %[[SIZE]], %[[MEM]], %[[PTR]], %[[IDX]], %[[VAL]] : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
-// CHECK:          }
-// CHECK:          return %[[OUT]]#0, %[[OUT]]#1, %[[OUT]]#2, %[[OUT]]#3, %[[OUT]]#4 : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
+
+// CHECK-LABEL:   func.func @for(
+// CHECK-SAME:                   %[[VAL_0:.*0]]: memref<3xindex>,
+// CHECK-SAME:                   %[[VAL_1:.*1]]: memref<?xindex>,
+// CHECK-SAME:                   %[[VAL_2:.*2]]: memref<?xindex>,
+// CHECK-SAME:                   %[[VAL_3:.*3]]: memref<?xf32>,
+// CHECK-SAME:                   %[[VAL_4:.*4]]: !llvm.struct<(array<1 x i64>)>,
+// CHECK-SAME:                   %[[VAL_5:.*5]]: index,
+// CHECK-SAME:                   %[[VAL_6:.*6]]: index,
+// CHECK-SAME:                   %[[VAL_7:.*7]]: index) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>, !llvm.struct<(array<1 x i64>)>) {
+// CHECK:           %[[VAL_8:.*]]:5 = scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_6]] step %[[VAL_7]] iter_args(
+// CHECK-SAME:        %[[VAL_10:.*]] = %[[VAL_0]],
+// CHECK-SAME:        %[[VAL_11:.*]] = %[[VAL_1]],
+// CHECK-SAME:        %[[VAL_12:.*]] = %[[VAL_2]],
+// CHECK-SAME:        %[[VAL_13:.*]] = %[[VAL_3]],
+// CHECK-SAME:        %[[VAL_14:.*]] = %[[VAL_4]])
+// CHECK:             scf.yield %[[VAL_10]], %[[VAL_11]], %[[VAL_12]], %[[VAL_13]], %[[VAL_14]] :
+// CHECK:           }
+// CHECK:           return %[[VAL_8]]#0, %[[VAL_8]]#1, %[[VAL_8]]#2, %[[VAL_8]]#3, %[[VAL_8]]#4 :
+// CHECK:         }
 func.func @for(%in: tensor<1024xf32, #SparseVector>,
                %lb: index, %ub: index, %step: index) -> tensor<1024xf32, #SparseVector> {
   %1 = scf.for %i = %lb to %ub step %step iter_args(%vin = %in)
      -> tensor<1024xf32, #SparseVector> {
     scf.yield %vin : tensor<1024xf32, #SparseVector>
   }
   return %1 : tensor<1024xf32, #SparseVector>
 }
 
-
-// CHECK-LABEL:  func @if(
-//  CHECK-SAME:          %[[DIM_SIZE:.*0]]: memref<1xindex>,
-//  CHECK-SAME:          %[[MEM_SIZE:.*1]]: memref<3xindex>,
-//  CHECK-SAME:          %[[POINTER:.*2]]: memref<?xindex>,
-//  CHECK-SAME:          %[[INDICES:.*3]]: memref<?xindex>,
-//  CHECK-SAME:          %[[VALUE:.*4]]: memref<?xf32>,
-//  CHECK-SAME:          %[[DIM_SIZE_1:.*5]]: memref<1xindex>,
-//  CHECK-SAME:          %[[MEM_SIZE_1:.*6]]: memref<3xindex>,
-//  CHECK-SAME:          %[[POINTER_1:.*7]]: memref<?xindex>,
-//  CHECK-SAME:          %[[INDICES_1:.*8]]: memref<?xindex>,
-//  CHECK-SAME:          %[[VALUE_1:.*9]]: memref<?xf32>,
-//  CHECK-SAME:          %[[I1:.*10]]: i1) ->
-//  CHECK-SAME:          (memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) {
-//       CHECK:  %[[SV:.*]]:5 = scf.if %[[I1]] -> (memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) {
-//       CHECK:    scf.yield %[[DIM_SIZE]], %[[MEM_SIZE]], %[[POINTER]], %[[INDICES]], %[[VALUE]] : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
-//       CHECK:  } else {
-//       CHECK:    scf.yield %[[DIM_SIZE_1]], %[[MEM_SIZE_1]], %[[POINTER_1]], %[[INDICES_1]], %[[VALUE_1]] : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
-//       CHECK:  }
-//       CHECK:  return %[[SV]]#0, %[[SV]]#1, %[[SV]]#2, %[[SV]]#3, %[[SV]]#4 : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
+// CHECK-LABEL:   func.func @if(
+// CHECK-SAME:                  %[[VAL_0:.*0]]: memref<3xindex>,
+// CHECK-SAME:                  %[[VAL_1:.*1]]: memref<?xindex>,
+// CHECK-SAME:                  %[[VAL_2:.*2]]: memref<?xindex>,
+// CHECK-SAME:                  %[[VAL_3:.*3]]: memref<?xf32>,
+// CHECK-SAME:                  %[[VAL_4:.*4]]: !llvm.struct<(array<1 x i64>)>,
+// CHECK-SAME:                  %[[VAL_5:.*5]]: memref<3xindex>,
+// CHECK-SAME:                  %[[VAL_6:.*6]]: memref<?xindex>,
+// CHECK-SAME:                  %[[VAL_7:.*7]]: memref<?xindex>,
+// CHECK-SAME:                  %[[VAL_8:.*8]]: memref<?xf32>,
+// CHECK-SAME:                  %[[VAL_9:.*9]]: !llvm.struct<(array<1 x i64>)>,
+// CHECK-SAME:                  %[[VAL_10:.*]]: i1)
+// CHECK:           %[[VAL_11:.*]]:5 = scf.if %[[VAL_10]]
+// CHECK:             scf.yield %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]], %[[VAL_4]]
+// CHECK:           } else {
+// CHECK:             scf.yield %[[VAL_5]], %[[VAL_6]], %[[VAL_7]], %[[VAL_8]], %[[VAL_9]]
+// CHECK:           }
+// CHECK:           return %[[VAL_11]]#0, %[[VAL_11]]#1, %[[VAL_11]]#2, %[[VAL_11]]#3, %[[VAL_11]]#4 :
+// CHECK-SAME:        memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>, !llvm.struct<(array<1 x i64>)>
 func.func @if(%t: tensor<1024xf32, #SparseVector>,
               %f: tensor<1024xf32, #SparseVector>,
               %c: i1) -> tensor<1024xf32, #SparseVector> {
   %1 = scf.if %c -> tensor<1024xf32, #SparseVector> {
     scf.yield %t : tensor<1024xf32, #SparseVector>
   } else {
     scf.yield %f : tensor<1024xf32, #SparseVector>
   }
   return %1 : tensor<1024xf32, #SparseVector>
 }
 
-// CHECK-LABEL:  func @while(
-//  CHECK-SAME:              %[[DIM_SIZE:.*0]]: memref<1xindex>,
-//  CHECK-SAME:              %[[MEM_SIZE:.*1]]: memref<3xindex>,
-//  CHECK-SAME:              %[[POINTER:.*2]]: memref<?xindex>,
-//  CHECK-SAME:              %[[INDICES:.*3]]: memref<?xindex>,
-//  CHECK-SAME:              %[[VALUE:.*4]]: memref<?xf32>,
-//  CHECK-SAME:              %[[I1:.*5]]: i1) ->
-//  CHECK-SAME:              (memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) {
-//       CHECK:  %[[SV:.*]]:5 = scf.while (
-//  CHECK-SAME:              %[[TMP_DIM:.*]] = %[[DIM_SIZE]],
-//  CHECK-SAME:              %[[TMP_MEM:.*]] = %[[MEM_SIZE]],
-//  CHECK-SAME:              %[[TMP_PTR:.*]] = %[[POINTER]],
-//  CHECK-SAME:              %[[TMP_IND:.*]] = %[[INDICES]],
-//  CHECK-SAME:              %[[TMP_VAL:.*]] = %[[VALUE]])
-//       CHECK:    scf.condition(%[[I1]]) %[[TMP_DIM]], %[[TMP_MEM]], %[[TMP_PTR]], %[[TMP_IND]], %[[TMP_VAL]] : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
-//       CHECK:  } do {
-//       CHECK:  ^bb0(%[[TMP_DIM]]: memref<1xindex>, %[[TMP_MEM]]: memref<3xindex>, %[[TMP_PTR]]: memref<?xindex>, %[[TMP_IND]]: memref<?xindex>, %[[TMP_VAL]]: memref<?xf32>):
-//       CHECK:    scf.yield %[[TMP_DIM]], %[[TMP_MEM]], %[[TMP_PTR]], %[[TMP_IND]], %[[TMP_VAL]] : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
-//       CHECK:  }
-//       CHECK:  return %[[SV]]#0, %[[SV]]#1, %[[SV]]#2, %[[SV]]#3, %[[SV]]#4 : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>
+
+// CHECK-LABEL:   func.func @while(
+// CHECK-SAME:                     %[[VAL_0:.*0]]: memref<3xindex>,
+// CHECK-SAME:                     %[[VAL_1:.*1]]: memref<?xindex>,
+// CHECK-SAME:                     %[[VAL_2:.*2]]: memref<?xindex>,
+// CHECK-SAME:                     %[[VAL_3:.*3]]: memref<?xf32>,
+// CHECK-SAME:                     %[[VAL_4:.*4]]: !llvm.struct<(array<1 x i64>)>,
+// CHECK-SAME:                     %[[VAL_5:.*5]]: i1)
+// CHECK:           %[[VAL_6:.*]]:5 = scf.while (
+// CHECK-SAME:        %[[VAL_7:.*]] = %[[VAL_0]],
+// CHECK-SAME:        %[[VAL_8:.*]] = %[[VAL_1]],
+// CHECK-SAME:        %[[VAL_9:.*]] = %[[VAL_2]],
+// CHECK-SAME:        %[[VAL_10:.*]] = %[[VAL_3]],
+// CHECK-SAME:        %[[VAL_11:.*]] = %[[VAL_4]])
+// CHECK:             scf.condition(%[[VAL_5]]) %[[VAL_7]], %[[VAL_8]], %[[VAL_9]], %[[VAL_10]], %[[VAL_11]]
+// CHECK:           } do {
+// CHECK:           ^bb0(%[[VAL_12:.*6]]: memref<3xindex>,
+// CHECK-SAME:           %[[VAL_13:.*7]]: memref<?xindex>,
+// CHECK-SAME:           %[[VAL_14:.*8]]: memref<?xindex>,
+// CHECK-SAME:           %[[VAL_15:.*9]]: memref<?xf32>,
+// CHECK-SAME:           %[[VAL_16:.*10]]: !llvm.struct<(array<1 x i64>)>):
+// CHECK:             scf.yield %[[VAL_12]], %[[VAL_13]], %[[VAL_14]], %[[VAL_15]], %[[VAL_16]]
+// CHECK:           }
+// CHECK:           return %[[VAL_6]]#0, %[[VAL_6]]#1, %[[VAL_6]]#2, %[[VAL_6]]#3, %[[VAL_6]]#4 :
+// CHECK-SAME:        memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>, !llvm.struct<(array<1 x i64>)>
 func.func @while(%arg0: tensor<1024xf32, #SparseVector>, %c: i1) -> tensor<1024xf32, #SparseVector> {
   %0 = scf.while (%in = %arg0) : (tensor<1024xf32, #SparseVector>) -> tensor<1024xf32, #SparseVector> {
     scf.condition(%c) %in : tensor<1024xf32, #SparseVector>
   } do {
   ^bb0(%arg1: tensor<1024xf32, #SparseVector>):
     scf.yield %arg1 : tensor<1024xf32, #SparseVector>
   }
   return %0: tensor<1024xf32, #SparseVector>
 }

mlir/test/Dialect/SparseTensor/sparse_matmul_codegen.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 
 // RUN: mlir-opt %s --linalg-generalize-named-ops \
 // RUN:  --sparsification --sparse-tensor-codegen \
 // RUN:  --canonicalize --cse | FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 
 //
 // Computes C = A x B with all matrices sparse (SpMSpM) in CSR.
 //
 // CHECK-LABEL:   func.func private @_insert_D_C_4_4_f64_0_0(
-// CHECK-SAME:      %[[VAL_0:.*]]: memref<2xindex>,
-// CHECK-SAME:      %[[VAL_1:.*]]: memref<3xindex>,
-// CHECK-SAME:      %[[VAL_2:[^ ]+]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_3:.*]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_4:.*]]: memref<?xf64>,
-// CHECK-SAME:      %[[VAL_5:[^ ]+]]: index,
-// CHECK-SAME:      %[[VAL_6:.*]]: index,
-// CHECK-SAME:      %[[VAL_7:.*]]: f64) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-// CHECK-DAG:       %[[VAL_8:.*]] = arith.constant false
-// CHECK-DAG:       %[[VAL_9:.*]] = arith.constant 1 : index
+// CHECK-SAME:      %[[VAL_0:.*0]]: memref<3xindex>,
+// CHECK-SAME:      %[[VAL_1:.*1]]: memref<?xindex>,
+// CHECK-SAME:      %[[VAL_2:.*2]]: memref<?xindex>,
+// CHECK-SAME:      %[[VAL_3:.*3]]: memref<?xf64>,
+// CHECK-SAME:      %[[VAL_4:.*4]]: !llvm.struct<(array<2 x i64>)>,
+// CHECK-SAME:      %[[VAL_5:.*5]]: index,
+// CHECK-SAME:      %[[VAL_6:.*6]]: index,
+// CHECK-SAME:      %[[VAL_7:.*7]]: f64) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+// CHECK:           %[[VAL_8:.*]] = arith.constant false
+// CHECK:           %[[VAL_9:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_10:.*]] = arith.addi %[[VAL_5]], %[[VAL_9]] : index
-// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]] : memref<?xindex>
-// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_9]]] : memref<3xindex>
+// CHECK:           %[[VAL_11:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK:           %[[VAL_12:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_10]]] : memref<?xindex>
+// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_9]]] : memref<3xindex>
 // CHECK:           %[[VAL_14:.*]] = arith.subi %[[VAL_12]], %[[VAL_9]] : index
 // CHECK:           %[[VAL_15:.*]] = arith.cmpi ult, %[[VAL_11]], %[[VAL_12]] : index
 // CHECK:           %[[VAL_16:.*]] = scf.if %[[VAL_15]] -> (i1) {
-// CHECK:             %[[VAL_17:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_14]]] : memref<?xindex>
+// CHECK:             %[[VAL_17:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_14]]] : memref<?xindex>
 // CHECK:             %[[VAL_18:.*]] = arith.cmpi eq, %[[VAL_17]], %[[VAL_6]] : index
 // CHECK:             scf.yield %[[VAL_18]] : i1
 // CHECK:           } else {
-// CHECK:             memref.store %[[VAL_13]], %[[VAL_2]]{{\[}}%[[VAL_5]]] : memref<?xindex>
+// CHECK:             memref.store %[[VAL_13]], %[[VAL_1]]{{\[}}%[[VAL_5]]] : memref<?xindex>
 // CHECK:             scf.yield %[[VAL_8]] : i1
 // CHECK:           }
 // CHECK:           %[[VAL_19:.*]] = scf.if %[[VAL_20:.*]] -> (memref<?xindex>) {
-// CHECK:             scf.yield %[[VAL_3]] : memref<?xindex>
+// CHECK:             scf.yield %[[VAL_2]] : memref<?xindex>
 // CHECK:           } else {
 // CHECK:             %[[VAL_21:.*]] = arith.addi %[[VAL_13]], %[[VAL_9]] : index
-// CHECK:             memref.store %[[VAL_21]], %[[VAL_2]]{{\[}}%[[VAL_10]]] : memref<?xindex>
-// CHECK:             %[[VAL_22:.*]] = sparse_tensor.push_back %[[VAL_1]], %[[VAL_3]], %[[VAL_6]] {idx = 1 : index} : memref<3xindex>, memref<?xindex>, index
+// CHECK:             memref.store %[[VAL_21]], %[[VAL_1]]{{\[}}%[[VAL_10]]] : memref<?xindex>
+// CHECK:             %[[VAL_22:.*]] = sparse_tensor.push_back %[[VAL_0]], %[[VAL_2]], %[[VAL_6]] {idx = 1 : index} : memref<3xindex>, memref<?xindex>, index
 // CHECK:             scf.yield %[[VAL_22]] : memref<?xindex>
 // CHECK:           }
-// CHECK:           %[[VAL_23:.*]] = sparse_tensor.push_back %[[VAL_1]], %[[VAL_4]], %[[VAL_7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
-// CHECK:           return %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_24:.*]], %[[VAL_23]] : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+// CHECK:           %[[VAL_23:.*]] = sparse_tensor.push_back %[[VAL_0]], %[[VAL_3]], %[[VAL_7]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
+// CHECK:           return %[[VAL_0]], %[[VAL_1]], %[[VAL_24:.*]], %[[VAL_23]], %[[VAL_4]] : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 // CHECK:         }
 
 // CHECK-LABEL:   func.func @matmul(
-// CHECK-SAME:      %[[VAL_0:.*0]]: memref<2xindex>,
-// CHECK-SAME:      %[[VAL_1:.*1]]: memref<3xindex>,
+// CHECK-SAME:      %[[VAL_0:.*0]]: memref<3xindex>,
+// CHECK-SAME:      %[[VAL_1:.*1]]: memref<?xindex>,
 // CHECK-SAME:      %[[VAL_2:.*2]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_3:.*3]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_4:.*4]]: memref<?xf64>,
-// CHECK-SAME:      %[[VAL_5:.*5]]: memref<2xindex>,
-// CHECK-SAME:      %[[VAL_6:.*6]]: memref<3xindex>,
+// CHECK-SAME:      %[[VAL_3:.*3]]: memref<?xf64>,
+// CHECK-SAME:      %[[VAL_4:.*4]]: !llvm.struct<(array<2 x i64>)>,
+// CHECK-SAME:      %[[VAL_5:.*5]]: memref<3xindex>,
+// CHECK-SAME:      %[[VAL_6:.*6]]: memref<?xindex>,
 // CHECK-SAME:      %[[VAL_7:.*7]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_8:.*8]]: memref<?xindex>,
-// CHECK-SAME:      %[[VAL_9:.*9]]: memref<?xf64>) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-// CHECK-DAG:       %[[VAL_10:.*]] = arith.constant 4 : index
-// CHECK-DAG:       %[[VAL_11:.*]] = arith.constant 0.000000e+00 : f64
-// CHECK-DAG:       %[[VAL_12:.*]] = arith.constant 0 : index
-// CHECK-DAG:       %[[VAL_13:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[VAL_14:.*]] = arith.constant false
-// CHECK-DAG:       %[[VAL_15:.*]] = arith.constant true
-// CHECK:           %[[VAL_16:.*]] = memref.alloc() : memref<2xindex>
+// CHECK-SAME:      %[[VAL_8:.*8]]: memref<?xf64>,
+// CHECK-SAME:      %[[VAL_9:.*9]]: !llvm.struct<(array<2 x i64>)>) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+// CHECK:           %[[VAL_10:.*]] = arith.constant 4 : index
+// CHECK:           %[[VAL_11:.*]] = arith.constant 4 : i64
+// CHECK:           %[[VAL_12:.*]] = arith.constant 0.000000e+00 : f64
+// CHECK:           %[[VAL_13:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_14:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_15:.*]] = arith.constant false
+// CHECK:           %[[VAL_16:.*]] = arith.constant true
 // CHECK:           %[[VAL_17:.*]] = memref.alloc() : memref<3xindex>
 // CHECK:           %[[VAL_18:.*]] = memref.alloc() : memref<16xindex>
 // CHECK:           %[[VAL_19:.*]] = memref.cast %[[VAL_18]] : memref<16xindex> to memref<?xindex>
 // CHECK:           %[[VAL_20:.*]] = memref.alloc() : memref<16xindex>
 // CHECK:           %[[VAL_21:.*]] = memref.cast %[[VAL_20]] : memref<16xindex> to memref<?xindex>
 // CHECK:           %[[VAL_22:.*]] = memref.alloc() : memref<16xf64>
 // CHECK:           %[[VAL_23:.*]] = memref.cast %[[VAL_22]] : memref<16xf64> to memref<?xf64>
-// CHECK:           linalg.fill ins(%[[VAL_12]] : index) outs(%[[VAL_17]] : memref<3xindex>)
-// CHECK:           memref.store %[[VAL_10]], %[[VAL_16]]{{\[}}%[[VAL_12]]] : memref<2xindex>
-// CHECK:           memref.store %[[VAL_10]], %[[VAL_16]]{{\[}}%[[VAL_13]]] : memref<2xindex>
-// CHECK:           %[[VAL_24:.*]] = sparse_tensor.push_back %[[VAL_17]], %[[VAL_19]], %[[VAL_12]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index
-// CHECK:           %[[VAL_25:.*]] = sparse_tensor.push_back %[[VAL_17]], %[[VAL_24]], %[[VAL_12]], %[[VAL_10]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index, index
-// CHECK:           %[[VAL_26:.*]] = memref.alloc() : memref<4xf64>
-// CHECK:           %[[VAL_27:.*]] = memref.alloc() : memref<4xi1>
-// CHECK:           %[[VAL_28:.*]] = memref.alloc() : memref<4xindex>
-// CHECK:           %[[VAL_29:.*]] = memref.cast %[[VAL_28]] : memref<4xindex> to memref<?xindex>
-// CHECK:           linalg.fill ins(%[[VAL_11]] : f64) outs(%[[VAL_26]] : memref<4xf64>)
-// CHECK:           linalg.fill ins(%[[VAL_14]] : i1) outs(%[[VAL_27]] : memref<4xi1>)
-// CHECK:           %[[VAL_30:.*]]:5 = scf.for %[[VAL_31:.*]] = %[[VAL_12]] to %[[VAL_10]] step %[[VAL_13]] iter_args(%[[VAL_32:.*]] = %[[VAL_16]], %[[VAL_33:.*]] = %[[VAL_17]], %[[VAL_34:.*]] = %[[VAL_25]], %[[VAL_35:.*]] = %[[VAL_21]], %[[VAL_36:.*]] = %[[VAL_23]]) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-// CHECK:             %[[VAL_37:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_31]]] : memref<?xindex>
-// CHECK:             %[[VAL_38:.*]] = arith.addi %[[VAL_31]], %[[VAL_13]] : index
-// CHECK:             %[[VAL_39:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_38]]] : memref<?xindex>
-// CHECK:             %[[VAL_40:.*]] = scf.for %[[VAL_41:.*]] = %[[VAL_37]] to %[[VAL_39]] step %[[VAL_13]] iter_args(%[[VAL_42:.*]] = %[[VAL_12]]) -> (index) {
-// CHECK:               %[[VAL_43:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_41]]] : memref<?xindex>
-// CHECK:               %[[VAL_44:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_41]]] : memref<?xf64>
-// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_43]]] : memref<?xindex>
-// CHECK:               %[[VAL_46:.*]] = arith.addi %[[VAL_43]], %[[VAL_13]] : index
-// CHECK:               %[[VAL_47:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_46]]] : memref<?xindex>
-// CHECK:               %[[VAL_48:.*]] = scf.for %[[VAL_49:.*]] = %[[VAL_45]] to %[[VAL_47]] step %[[VAL_13]] iter_args(%[[VAL_50:.*]] = %[[VAL_42]]) -> (index) {
-// CHECK:                 %[[VAL_51:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_49]]] : memref<?xindex>
-// CHECK:                 %[[VAL_52:.*]] = memref.load %[[VAL_26]]{{\[}}%[[VAL_51]]] : memref<4xf64>
-// CHECK:                 %[[VAL_53:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_49]]] : memref<?xf64>
-// CHECK:                 %[[VAL_54:.*]] = arith.mulf %[[VAL_44]], %[[VAL_53]] : f64
-// CHECK:                 %[[VAL_55:.*]] = arith.addf %[[VAL_52]], %[[VAL_54]] : f64
-// CHECK:                 %[[VAL_56:.*]] = memref.load %[[VAL_27]]{{\[}}%[[VAL_51]]] : memref<4xi1>
-// CHECK:                 %[[VAL_57:.*]] = arith.cmpi eq, %[[VAL_56]], %[[VAL_14]] : i1
-// CHECK:                 %[[VAL_58:.*]] = scf.if %[[VAL_57]] -> (index) {
-// CHECK:                   memref.store %[[VAL_15]], %[[VAL_27]]{{\[}}%[[VAL_51]]] : memref<4xi1>
-// CHECK:                   memref.store %[[VAL_51]], %[[VAL_28]]{{\[}}%[[VAL_50]]] : memref<4xindex>
-// CHECK:                   %[[VAL_59:.*]] = arith.addi %[[VAL_50]], %[[VAL_13]] : index
-// CHECK:                   scf.yield %[[VAL_59]] : index
+// CHECK:           %[[VAL_24:.*]] = llvm.mlir.undef : !llvm.struct<(array<2 x i64>)>
+// CHECK:           linalg.fill ins(%[[VAL_13]] : index) outs(%[[VAL_17]] : memref<3xindex>)
+// CHECK:           %[[VAL_25:.*]] = llvm.insertvalue %[[VAL_11]], %[[VAL_24]][0, 0] : !llvm.struct<(array<2 x i64>)>
+// CHECK:           %[[VAL_26:.*]] = llvm.insertvalue %[[VAL_11]], %[[VAL_25]][0, 1] : !llvm.struct<(array<2 x i64>)>
+// CHECK:           %[[VAL_27:.*]] = sparse_tensor.push_back %[[VAL_17]], %[[VAL_19]], %[[VAL_13]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index
+// CHECK:           %[[VAL_28:.*]] = sparse_tensor.push_back %[[VAL_17]], %[[VAL_27]], %[[VAL_13]], %[[VAL_10]] {idx = 0 : index} : memref<3xindex>, memref<?xindex>, index, index
+// CHECK:           %[[VAL_29:.*]] = memref.alloc() : memref<4xf64>
+// CHECK:           %[[VAL_30:.*]] = memref.alloc() : memref<4xi1>
+// CHECK:           %[[VAL_31:.*]] = memref.alloc() : memref<4xindex>
+// CHECK:           %[[VAL_32:.*]] = memref.cast %[[VAL_31]] : memref<4xindex> to memref<?xindex>
+// CHECK:           linalg.fill ins(%[[VAL_12]] : f64) outs(%[[VAL_29]] : memref<4xf64>)
+// CHECK:           linalg.fill ins(%[[VAL_15]] : i1) outs(%[[VAL_30]] : memref<4xi1>)
+// CHECK:           %[[VAL_33:.*]]:5 = scf.for %[[VAL_34:.*]] = %[[VAL_13]] to %[[VAL_10]] step %[[VAL_14]] iter_args(%[[VAL_35:.*]] = %[[VAL_17]], %[[VAL_36:.*]] = %[[VAL_28]], %[[VAL_37:.*]] = %[[VAL_21]], %[[VAL_38:.*]] = %[[VAL_23]], %[[VAL_39:.*]] = %[[VAL_26]]) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+// CHECK:             %[[VAL_40:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_34]]] : memref<?xindex>
+// CHECK:             %[[VAL_41:.*]] = arith.addi %[[VAL_34]], %[[VAL_14]] : index
+// CHECK:             %[[VAL_42:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_41]]] : memref<?xindex>
+// CHECK:             %[[VAL_43:.*]] = scf.for %[[VAL_44:.*]] = %[[VAL_40]] to %[[VAL_42]] step %[[VAL_14]] iter_args(%[[VAL_45:.*]] = %[[VAL_13]]) -> (index) {
+// CHECK:               %[[VAL_46:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_44]]] : memref<?xindex>
+// CHECK:               %[[VAL_47:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_44]]] : memref<?xf64>
+// CHECK:               %[[VAL_48:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_46]]] : memref<?xindex>
+// CHECK:               %[[VAL_49:.*]] = arith.addi %[[VAL_46]], %[[VAL_14]] : index
+// CHECK:               %[[VAL_50:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_49]]] : memref<?xindex>
+// CHECK:               %[[VAL_51:.*]] = scf.for %[[VAL_52:.*]] = %[[VAL_48]] to %[[VAL_50]] step %[[VAL_14]] iter_args(%[[VAL_53:.*]] = %[[VAL_45]]) -> (index) {
+// CHECK:                 %[[VAL_54:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_52]]] : memref<?xindex>
+// CHECK:                 %[[VAL_55:.*]] = memref.load %[[VAL_29]]{{\[}}%[[VAL_54]]] : memref<4xf64>
+// CHECK:                 %[[VAL_56:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_52]]] : memref<?xf64>
+// CHECK:                 %[[VAL_57:.*]] = arith.mulf %[[VAL_47]], %[[VAL_56]] : f64
+// CHECK:                 %[[VAL_58:.*]] = arith.addf %[[VAL_55]], %[[VAL_57]] : f64
+// CHECK:                 %[[VAL_59:.*]] = memref.load %[[VAL_30]]{{\[}}%[[VAL_54]]] : memref<4xi1>
+// CHECK:                 %[[VAL_60:.*]] = arith.cmpi eq, %[[VAL_59]], %[[VAL_15]] : i1
+// CHECK:                 %[[VAL_61:.*]] = scf.if %[[VAL_60]] -> (index) {
+// CHECK:                   memref.store %[[VAL_16]], %[[VAL_30]]{{\[}}%[[VAL_54]]] : memref<4xi1>
+// CHECK:                   memref.store %[[VAL_54]], %[[VAL_31]]{{\[}}%[[VAL_53]]] : memref<4xindex>
+// CHECK:                   %[[VAL_62:.*]] = arith.addi %[[VAL_53]], %[[VAL_14]] : index
+// CHECK:                   scf.yield %[[VAL_62]] : index
 // CHECK:                 } else {
-// CHECK:                   scf.yield %[[VAL_50]] : index
-// CHECK:                 }
-// CHECK:                 memref.store %[[VAL_55]], %[[VAL_26]]{{\[}}%[[VAL_51]]] : memref<4xf64>
-// CHECK:                 scf.yield %[[VAL_60:.*]] : index
-// CHECK:               } {"Emitted from" = "linalg.generic"}
-// CHECK:             sparse_tensor.sort %[[VAL_62:.*]], %[[VAL_29]] : memref<?xindex>
-// CHECK:             %[[VAL_63:.*]]:5 = scf.for %[[VAL_64:.*]] = %[[VAL_12]] to %[[VAL_62]] step %[[VAL_13]] iter_args(%[[VAL_65:.*]] = %[[VAL_32]], %[[VAL_66:.*]] = %[[VAL_33]], %[[VAL_67:.*]] = %[[VAL_34]], %[[VAL_68:.*]] = %[[VAL_35]], %[[VAL_69:.*]] = %[[VAL_36]]) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
-// CHECK:               %[[VAL_70:.*]] = memref.load %[[VAL_28]]{{\[}}%[[VAL_64]]] : memref<4xindex>
-// CHECK:               %[[VAL_71:.*]] = memref.load %[[VAL_26]]{{\[}}%[[VAL_70]]] : memref<4xf64>
-// CHECK:               %[[VAL_72:.*]]:5 = func.call @_insert_D_C_4_4_f64_0_0(%[[VAL_65]], %[[VAL_66]], %[[VAL_67]], %[[VAL_68]], %[[VAL_69]], %[[VAL_31]], %[[VAL_70]], %[[VAL_71]]) : (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, index, index, f64) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>)
-// CHECK:               memref.store %[[VAL_11]], %[[VAL_26]]{{\[}}%[[VAL_70]]] : memref<4xf64>
-// CHECK:               memref.store %[[VAL_14]], %[[VAL_27]]{{\[}}%[[VAL_70]]] : memref<4xi1>
-// CHECK:               scf.yield %[[VAL_72]]#0, %[[VAL_72]]#1, %[[VAL_72]]#2, %[[VAL_72]]#3, %[[VAL_72]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
-// CHECK:             }
-// CHECK:             scf.yield %[[VAL_73:.*]]#0, %[[VAL_73]]#1, %[[VAL_73]]#2, %[[VAL_73]]#3, %[[VAL_73]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
-// CHECK:           } {"Emitted from" = "linalg.generic"}
-// CHECK:           memref.dealloc %[[VAL_26]] : memref<4xf64>
-// CHECK:           memref.dealloc %[[VAL_27]] : memref<4xi1>
-// CHECK:           memref.dealloc %[[VAL_28]] : memref<4xindex>
-// CHECK:           %[[VAL_74:.*]] = memref.load %[[VAL_75:.*]]#1{{\[}}%[[VAL_12]]] : memref<3xindex>
-// CHECK:           %[[VAL_76:.*]] = memref.load %[[VAL_75]]#2{{\[}}%[[VAL_12]]] : memref<?xindex>
-// CHECK:           %[[VAL_77:.*]] = scf.for %[[VAL_78:.*]] = %[[VAL_13]] to %[[VAL_74]] step %[[VAL_13]] iter_args(%[[VAL_79:.*]] = %[[VAL_76]]) -> (index) {
-// CHECK:             %[[VAL_80:.*]] = memref.load %[[VAL_75]]#2{{\[}}%[[VAL_78]]] : memref<?xindex>
-// CHECK:             %[[VAL_81:.*]] = arith.cmpi eq, %[[VAL_80]], %[[VAL_12]] : index
-// CHECK:             %[[VAL_82:.*]] = arith.select %[[VAL_81]], %[[VAL_79]], %[[VAL_80]] : index
-// CHECK:             scf.if %[[VAL_81]] {
-// CHECK:               memref.store %[[VAL_79]], %[[VAL_75]]#2{{\[}}%[[VAL_78]]] : memref<?xindex>
+// CHECK:                   scf.yield %[[VAL_53]] : index
+// CHECK:                 }
+// CHECK:                 memref.store %[[VAL_58]], %[[VAL_29]]{{\[}}%[[VAL_54]]] : memref<4xf64>
+// CHECK:                 scf.yield %[[VAL_63:.*]] : index
+// CHECK:               }
+// CHECK:               scf.yield %[[VAL_64:.*]] : index
+// CHECK:             }
+// CHECK:             sparse_tensor.sort %[[VAL_65:.*]], %[[VAL_32]] : memref<?xindex>
+// CHECK:             %[[VAL_66:.*]]:5 = scf.for %[[VAL_67:.*]] = %[[VAL_13]] to %[[VAL_65]] step %[[VAL_14]] iter_args(%[[VAL_68:.*]] = %[[VAL_35]], %[[VAL_69:.*]] = %[[VAL_36]], %[[VAL_70:.*]] = %[[VAL_37]], %[[VAL_71:.*]] = %[[VAL_38]], %[[VAL_72:.*]] = %[[VAL_39]]) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>) {
+// CHECK:               %[[VAL_73:.*]] = memref.load %[[VAL_31]]{{\[}}%[[VAL_67]]] : memref<4xindex>
+// CHECK:               %[[VAL_74:.*]] = memref.load %[[VAL_29]]{{\[}}%[[VAL_73]]] : memref<4xf64>
+// CHECK:               %[[VAL_75:.*]]:5 = func.call @_insert_D_C_4_4_f64_0_0(%[[VAL_68]], %[[VAL_69]], %[[VAL_70]], %[[VAL_71]], %[[VAL_72]], %[[VAL_34]], %[[VAL_73]], %[[VAL_74]]) : (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>, index, index, f64) -> (memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>)
+// CHECK:               memref.store %[[VAL_12]], %[[VAL_29]]{{\[}}%[[VAL_73]]] : memref<4xf64>
+// CHECK:               memref.store %[[VAL_15]], %[[VAL_30]]{{\[}}%[[VAL_73]]] : memref<4xi1>
+// CHECK:               scf.yield %[[VAL_75]]#0, %[[VAL_75]]#1, %[[VAL_75]]#2, %[[VAL_75]]#3, %[[VAL_75]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
+// CHECK:             }
+// CHECK:             scf.yield %[[VAL_76:.*]]#0, %[[VAL_76]]#1, %[[VAL_76]]#2, %[[VAL_76]]#3, %[[VAL_76]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
+// CHECK:           }
+// CHECK:           memref.dealloc %[[VAL_29]] : memref<4xf64>
+// CHECK:           memref.dealloc %[[VAL_30]] : memref<4xi1>
+// CHECK:           memref.dealloc %[[VAL_31]] : memref<4xindex>
+// CHECK:           %[[VAL_77:.*]] = memref.load %[[VAL_78:.*]]#0{{\[}}%[[VAL_13]]] : memref<3xindex>
+// CHECK:           %[[VAL_79:.*]] = memref.load %[[VAL_78]]#1{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:           %[[VAL_80:.*]] = scf.for %[[VAL_81:.*]] = %[[VAL_14]] to %[[VAL_77]] step %[[VAL_14]] iter_args(%[[VAL_82:.*]] = %[[VAL_79]]) -> (index) {
+// CHECK:             %[[VAL_83:.*]] = memref.load %[[VAL_78]]#1{{\[}}%[[VAL_81]]] : memref<?xindex>
+// CHECK:             %[[VAL_84:.*]] = arith.cmpi eq, %[[VAL_83]], %[[VAL_13]] : index
+// CHECK:             %[[VAL_85:.*]] = arith.select %[[VAL_84]], %[[VAL_82]], %[[VAL_83]] : index
+// CHECK:             scf.if %[[VAL_84]] {
+// CHECK:               memref.store %[[VAL_82]], %[[VAL_78]]#1{{\[}}%[[VAL_81]]] : memref<?xindex>
 // CHECK:             }
-// CHECK:             scf.yield %[[VAL_82]] : index
+// CHECK:             scf.yield %[[VAL_85]] : index
 // CHECK:           }
-// CHECK:           return %[[VAL_75]]#0, %[[VAL_75]]#1, %[[VAL_75]]#2, %[[VAL_75]]#3, %[[VAL_75]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+// CHECK:           return %[[VAL_78]]#0, %[[VAL_78]]#1, %[[VAL_78]]#2, %[[VAL_78]]#3, %[[VAL_78]]#4 : memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, !llvm.struct<(array<2 x i64>)>
 // CHECK:         }
 func.func @matmul(%A: tensor<4x8xf64, #CSR>,
                   %B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> {
   %C = bufferization.alloc_tensor() : tensor<4x4xf64, #CSR>
   %D = linalg.matmul
     ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>)
        outs(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
   return %D: tensor<4x4xf64, #CSR>
 }