[mlir][sparse] Converting SparseTensorCOO to use standard C++-style iterators.

[wrengr]

Summary

This differential comprises three related changes: (1) it gives SparseTensorCOO standard C++-style iterators; (2) it removes the old iterator stuff from SparseTensorCOO; and (3) it introduces SparseTensorIterator which behaves like the old SparseTensorCOO iterator stuff used to.

The SparseTensorIterator class is needed because the MLIR codegen cannot easily use the C++-style iterators (hence why SparseTensorCOO had the old iterator stuff). Distinguishing SparseTensorIterator from SparseTensorCOO also helps improve API hygiene since these two classes are used for distinct purposes. And having SparseTensorIterator as its own class enables changing the underlying implementation in the future, without needing to worry about updating all the codegen tests etc.

[wrengr]

removing obsolete todos

[wrengr]

rebase

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

99

yeah, this is a pretty neat "forwarding" to avoid the problem

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

887

note that we don't really have this method, so perhaps

new SparseTensor(src, Iterator)

or something like that will make that more clear/

[wrengr]

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

887

but we don't have that method either ;) Or did you mean spelling it more like the generated MLIR code rather than like the pseudo-C++ code used in the rest of the comment?

fwiw, I was thinking about changing the SparseTensorIterator ctor to take the SparseTensorStorage object directly (and call toCOO internally), but got a bit hung up on the details since toCOO also takes a permutation argument but the current version takes it as just a pointer rather than using a cleaner/safer API. So I'll probably end up making that change a bit down the road once I work out the details a bit more

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

 using ElementConsumer =
     const std::function<void(const std::vector<uint64_t> &, V)> &;
 
 /// A memory-resident sparse tensor in coordinate-scheme representation
 /// (a collection of `Element`s).  This data structure is used as
 /// an intermediate representation; e.g., for reading sparse tensors
 /// from external formats into memory, or for certain conversions between
 /// different `SparseTensorStorage` formats.
+///
+/// This class provides all the typedefs required by the "Container"
+/// concept (<https://en.cppreference.com/w/cpp/named_req/Container>);
+/// however, beware that it cannot fully implement that concept since
+/// it cannot have a default ctor (because the `dimSizes` field is const).
+/// Thus these typedefs are provided for familiarity reasons, rather
+/// than as a proper implementation of the concept.
 template <typename V>
 class SparseTensorCOO final {
 public:
+  using value_type = const Element<V>;
+  using reference = value_type &;
+  using const_reference = reference;
+  // The types associated with `std::vector` differ significantly between
+  // C++11/17 vs C++20; so we explicitly defer to whatever `std::vector`
+  // says the types should be.
+  using vector_type = std::vector<Element<V>>;
+  using iterator = typename vector_type::const_iterator;
+  using const_iterator = iterator;
+  using difference_type = typename vector_type::difference_type;
+  using size_type = typename vector_type::size_type;
+
   SparseTensorCOO(const std::vector<uint64_t> &dimSizes, uint64_t capacity)
-      : dimSizes(dimSizes), isSorted(true), iteratorLocked(false),
-        iteratorPos(0) {
+      : dimSizes(dimSizes), isSorted(true) {
     if (capacity) {
       elements.reserve(capacity);
       indices.reserve(capacity * getRank());
     }
   }
 
   /// Factory method. Permutes the original dimensions according to
   /// the given ordering and expects subsequent add() calls to honor
public:
   /// Returns the `operator<` closure object for the COO's element type.
   ElementLT<V> getElementLT() const { return ElementLT<V>(getRank()); }
 
   /// Adds an element to the tensor.  This method does not check whether
   /// `ind` is already associated with a value, it adds it regardless.
   /// Resolving such conflicts is left up to clients of the iterator
   /// interface.
   ///
+  /// This method invalidates all iterators.
+  ///
   /// Asserts:
-  /// * is not in iterator mode
   /// * the `ind` is valid for `rank`
   /// * the elements of `ind` are valid for `dimSizes`.
   void add(const std::vector<uint64_t> &ind, V val) {
-    assert(!iteratorLocked && "Attempt to add() after startIterator()");
     const uint64_t *base = indices.data();
     uint64_t size = indices.size();
     uint64_t rank = getRank();
     assert(ind.size() == rank && "Element rank mismatch");
     for (uint64_t r = 0; r < rank; ++r) {
       assert(ind[r] < dimSizes[r] && "Index is too large for the dimension");
       indices.push_back(ind[r]);
     }
void add(const std::vector<uint64_t> &ind, V val) {
     }
     // Add the new element and update the sorted bit.
     Element<V> addedElem(base + size, val);
     if (!elements.empty() && isSorted)
       isSorted = getElementLT()(elements.back(), addedElem);
     elements.push_back(addedElem);
   }
 
+  const_iterator begin() const { return elements.cbegin(); }
+  const_iterator end() const { return elements.cend(); }
+
   /// Sorts elements lexicographically by index.  If an index is mapped to
   /// multiple values, then the relative order of those values is unspecified.
   ///
-  /// Asserts: is not in iterator mode.
+  /// This method invalidates all iterators.
   void sort() {
-    assert(!iteratorLocked && "Attempt to sort() after startIterator()");
     if (isSorted)
       return;
     std::sort(elements.begin(), elements.end(), getElementLT());
     isSorted = true;
   }
 
-  /// Switches into iterator mode.  If already in iterator mode, then
-  /// resets the position to the first element.
-  void startIterator() {
-    iteratorLocked = true;
-    iteratorPos = 0;
-  }
-
-  /// Gets the next element.  If there are no remaining elements, then
-  /// returns nullptr and switches out of iterator mode.
-  ///
-  /// Asserts: is in iterator mode.
-  const Element<V> *getNext() {
-    assert(iteratorLocked && "Attempt to getNext() before startIterator()");
-    if (iteratorPos < elements.size())
-      return &(elements[iteratorPos++]);
-    iteratorLocked = false;
-    return nullptr;
-  }
-
 private:
   const std::vector<uint64_t> dimSizes; // per-dimension sizes
   std::vector<Element<V>> elements;     // all COO elements
   std::vector<uint64_t> indices;        // shared index pool
   bool isSorted;
-  bool iteratorLocked;
-  unsigned iteratorPos;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_COO_H

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 // functions for each supported primary- and overhead-type.
 //
 //===----------------------------------------------------------------------===//
 
 /// The @newSparseTensor function for constructing a new sparse tensor.
 /// This is the "swiss army knife" method for materializing sparse
 /// tensors into the computation.  The types of the `ptr` argument and
 /// the result depend on the action, as explained in the following table
-/// (where "STS" means a sparse-tensor-storage object, and "COO" means
-/// a coordinate-scheme object).
+/// (where "STS" means a sparse-tensor-storage object, "COO" means
+/// a coordinate-scheme object, and "Iterator" means an iterator object).
 ///
 /// Action:         `ptr`:          Returns:
 /// kEmpty          unused          STS, empty
 /// kEmptyCOO       unused          COO, empty
 /// kFromFile       char* filename  STS, read from the file
 /// kFromCOO        COO             STS, copied from the COO source
 /// kToCOO          STS             COO, copied from the STS source
 /// kSparseToSparse STS             STS, copied from the STS source
-/// kToIterator     STS             COO-Iterator, call @getNext to use
+/// kToIterator     STS             Iterator, call @getNext to use and
+///                                 @delSparseTensorIterator to free.
 MLIR_CRUNNERUTILS_EXPORT void *
 _mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
                              StridedMemRefType<index_type, 1> *sref,
                              StridedMemRefType<index_type, 1> *pref,
                              OverheadType ptrTp, OverheadType indTp,
                              PrimaryType valTp, Action action, void *ptr);
 
 /// Tensor-storage method to obtain direct access to the values array.

 MLIR_CRUNNERUTILS_EXPORT void delSparseTensor(void *tensor);
 
 /// Releases the memory for the coordinate-scheme object.
 #define DECL_DELCOO(VNAME, V)                                                  \
   MLIR_CRUNNERUTILS_EXPORT void delSparseTensorCOO##VNAME(void *coo);
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_DELCOO)
 #undef DECL_DELCOO
 
+/// Releases the memory for an iterator object.
+#define DECL_DELITER(VNAME, V)                                                 \
+  MLIR_CRUNNERUTILS_EXPORT void delSparseTensorIterator##VNAME(void *iter);
+MLIR_SPARSETENSOR_FOREVERY_V(DECL_DELITER)
+#undef DECL_DELITER
+
 /// Helper function to read a sparse tensor filename from the environment,
 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
 MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id);
 
 /// Helper function to read the header of a file and return the
 /// shape/sizes, without parsing the elements of the file.
 MLIR_CRUNNERUTILS_EXPORT void readSparseTensorShape(char *filename,
                                                     std::vector<uint64_t> *out);

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

 
 /// Generates a call to release/delete a `SparseTensorCOO`.
 static void genDelCOOCall(OpBuilder &builder, Location loc, Type elemTp,
                           Value coo) {
   SmallString<21> name{"delSparseTensorCOO", primaryTypeFunctionSuffix(elemTp)};
   createFuncCall(builder, loc, name, {}, coo, EmitCInterface::Off);
 }
 
+/// Generates a call to release/delete a `SparseTensorIterator`.
+static void genDelIteratorCall(OpBuilder &builder, Location loc, Type elemTp,
+                               Value iter) {
+  SmallString<26> name{"delSparseTensorIterator",
+                       primaryTypeFunctionSuffix(elemTp)};
+  createFuncCall(builder, loc, name, {}, iter, EmitCInterface::Off);
+}
+
 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];
 ///   if val != 0
 ///     t->add(&val, [i1,..,ik], [p1,..,pk]);
 static void genAddEltCall(OpBuilder &builder, Location loc, Type eltType,
                           Value ptr, Value valPtr, Value ind, Value perm) {
   SmallString<9> name{"addElt", primaryTypeFunctionSuffix(eltType)};
   SmallVector<Value, 4> params{ptr, valPtr, ind, perm};
   Type pTp = getOpaquePointerType(builder);
   createFuncCall(builder, loc, name, pTp, params, EmitCInterface::On);
 }
 
 /// Generates a call to `iter->getNext()`.  If there is a next element,
 /// then it is copied into the out-parameters `ind` and `elemPtr`,
 /// and the return value is true.  If there isn't a next element, then
-/// the memory for `iter` is freed and the return value is false.
+/// the return value is false.
 static Value genGetNextCall(OpBuilder &builder, Location loc, Value iter,
                             Value ind, Value elemPtr) {
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
   SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params{iter, ind, elemPtr};
   Type i1 = builder.getI1Type();
   return createFuncCall(builder, loc, name, i1, params, EmitCInterface::On)
       .getResult(0);
genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
   genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
   rewriter.create<scf::YieldOp>(loc);
   // Final call to construct sparse tensor storage and free temporary resources.
   rewriter.setInsertionPointAfter(whileOp);
   params[6] = constantAction(rewriter, loc, Action::kFromCOO);
   params[7] = coo;
   Value dst = genNewCall(rewriter, loc, params);
   genDelCOOCall(rewriter, loc, elemTp, coo);
-  genDelCOOCall(rewriter, loc, elemTp, iter);
+  genDelIteratorCall(rewriter, loc, elemTp, iter);
   rewriter.replaceOp(op, dst);
   return success();
 }
 
 // Generates a while loop that iterates over the COO list extracted
 // from `t`, using `bodyBuilder` to build the loop body.
 //   while (elem = coo->getNext()) {
 //     bodyBuilder
 //   }
 // TODO: It can be used by other operators (ReshapeOp, ConvertOP) conversion to
 // reduce code repetition!
+// TODO: rename to `genSparseIterationLoop`?
 static void genSparseCOOIterationLoop(
     ConversionPatternRewriter &rewriter, Location loc, Value t,
     RankedTensorType tensorTp,
     function_ref<void(OpBuilder &, Location, Value, Value)> bodyBuilder) {
   auto enc = getSparseTensorEncoding(tensorTp);
   assert(enc && "Generating Sparse Tensor COO Loop on a Dense Tensor!");
 
   unsigned rank = tensorTp.getRank();
static void genSparseCOOIterationLoop(
   rewriter.setInsertionPointToStart(after);
   // Callback here to build loop body.
   bodyBuilder(rewriter, loc, srcIdx, elemPtr);
   rewriter.create<scf::YieldOp>(loc);
   // Finish generating loop.
   rewriter.setInsertionPointAfter(whileOp);
 
   // Free memory for iterator.
-  genDelCOOCall(rewriter, loc, elemTp, iter);
+  genDelIteratorCall(rewriter, loc, elemTp, iter);
 }
 
 // Generate loop that iterates over a dense tensor.
 //   for i1 in dim1
 //    ..
 //     for ik in dimk
 //       val = a[i1,..,ik]
 //       if val != 0
if (encDst && encSrc) {
         genDelCOOCall(rewriter, loc, stp.getElementType(), coo);
         rewriter.replaceOp(op, dst);
       }
       return success();
     }
     if (!encDst && encSrc) {
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
-      //   iter = src->toCOO();
-      //   iter->startIterator();
+      //   iter = new SparseTensorIterator(src);
       //   while (elem = iter->getNext()) {
       //     dst[elem.indices] = elem.value;
       //   }
+      //   delete iter;
       RankedTensorType dstTensorTp = resType.cast<RankedTensorType>();
       RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
       unsigned rank = dstTensorTp.getRank();
       Type elemTp = dstTensorTp.getElementType();
       // Fabricate a no-permutation encoding for newParams().
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
if (!encDst && encSrc) {
       Value cond = genGetNextCall(rewriter, loc, iter, ind, elemPtr);
       rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
       Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
       rewriter.setInsertionPointToStart(after);
       SmallVector<Value, 4> ivs = loadIndices(rewriter, loc, rank, ind);
       insertScalarIntoDenseTensor(rewriter, loc, elemPtr, dst, ivs);
       rewriter.create<scf::YieldOp>(loc);
       rewriter.setInsertionPointAfter(whileOp);
-      genDelCOOCall(rewriter, loc, elemTp, iter);
+      genDelIteratorCall(rewriter, loc, elemTp, iter);
       rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, dst);
       // Deallocate the buffer.
       if (bufferization::allocationDoesNotEscape(op->getOpResult(0))) {
         rewriter.setInsertionPoint(insertionBlock->getTerminator());
         deallocDenseTensor(rewriter, loc, dst);
       }
       return success();
     }

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 //
 // Implementation details for public functions, which don't have a good
 // place to live in the C++ library this file is wrapping.
 //
 //===----------------------------------------------------------------------===//
 
 namespace {
 
+/// Wrapper class to avoid memory leakage issues.  The `SparseTensorCOO<V>`
+/// class provides a standard C++ iterator interface, where the iterator
+/// is implemented as per `std::vector`'s iterator.  However, for MLIR's
+/// usage we need to have an iterator which also holds onto the underlying
+/// `SparseTensorCOO<V>` so that it can be freed whenever the iterator
+/// is freed.
+//
+// We name this `SparseTensorIterator` rather than `SparseTensorCOOIterator`
+// for future-proofing, since the use of `SparseTensorCOO` is an
+// implementation detail that we eventually want to change (e.g., to
+// use `SparseTensorEnumerator` directly, rather than constructing the
+// intermediate `SparseTensorCOO` at all).
+template <typename V>
+class SparseTensorIterator final {
+public:
+  /// This ctor requires `coo` to be a non-null pointer to a dynamically
+  /// allocated object, and takes ownership of that object.  Therefore,
+  /// callers must not free the underlying COO object, since the iterator's
+  /// dtor will do so.
+  explicit SparseTensorIterator(const SparseTensorCOO<V> *coo)
+      : coo(coo), it(coo->begin()), end(coo->end()) {}
+
+  ~SparseTensorIterator() { delete coo; }
+
+  // Disable copy-ctor and copy-assignment, to prevent double-free.
+  SparseTensorIterator(const SparseTensorIterator<V> &) = delete;
+  SparseTensorIterator<V> &operator=(const SparseTensorIterator<V> &) = delete;
+
+  /// Gets the next element.  If there are no remaining elements, then
+  /// returns nullptr.
+  const Element<V> *getNext() { return it < end ? &*it++ : nullptr; }
+
+private:
+  const SparseTensorCOO<V> *const coo; // Owning pointer.
+  typename SparseTensorCOO<V>::const_iterator it;
+  const typename SparseTensorCOO<V>::const_iterator end;
+};
+
 /// Initializes sparse tensor from an external COO-flavored format.
 /// Used by `IMPL_CONVERTTOMLIRSPARSETENSOR`.
 // TODO: generalize beyond 64-bit indices.
 template <typename V>
 static SparseTensorStorage<uint64_t, uint64_t, V> *
 toMLIRSparseTensor(uint64_t rank, uint64_t nse, const uint64_t *shape,
                    const V *values, const uint64_t *indices,
                    const uint64_t *perm, const DimLevelType *sparsity) {
if (action == Action::kSparseToSparse) {                                   \
       auto *tensor = static_cast<SparseTensorStorageBase *>(ptr);              \
       return SparseTensorStorage<P, I, V>::newSparseTensor(rank, shape, perm,  \
                                                            sparsity, tensor);  \
     }                                                                          \
     if (action == Action::kEmptyCOO)                                           \
       return SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm);        \
     coo = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);       \
     if (action == Action::kToIterator) {                                       \
-      coo->startIterator();                                                    \
+      return new SparseTensorIterator<V>(coo);                                 \
     } else {                                                                   \
       assert(action == Action::kToCOO);                                        \
     }                                                                          \
     return coo;                                                                \
   }
 
 #define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
 
void *_mlir_ciface_addElt##VNAME(void *coo, StridedMemRefType<V, 0> *vref,   \
     V *value = vref->data + vref->offset;                                      \
     static_cast<SparseTensorCOO<V> *>(coo)->add(indices, *value);              \
     return coo;                                                                \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_ADDELT)
 #undef IMPL_ADDELT
 
 #define IMPL_GETNEXT(VNAME, V)                                                 \
-  bool _mlir_ciface_getNext##VNAME(void *coo,                                  \
+  bool _mlir_ciface_getNext##VNAME(void *iter,                                 \
                                    StridedMemRefType<index_type, 1> *iref,     \
                                    StridedMemRefType<V, 0> *vref) {            \
-    assert(coo &&iref &&vref);                                                 \
+    assert(iter &&iref &&vref);                                                \
     assert(iref->strides[0] == 1);                                             \
     index_type *indx = iref->data + iref->offset;                              \
     V *value = vref->data + vref->offset;                                      \
     const uint64_t isize = iref->sizes[0];                                     \
     const Element<V> *elem =                                                   \
-        static_cast<SparseTensorCOO<V> *>(coo)->getNext();                     \
+        static_cast<SparseTensorIterator<V> *>(iter)->getNext();               \
     if (elem == nullptr)                                                       \
       return false;                                                            \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indx[r] = elem->indices[r];                                              \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)

 
 #define IMPL_DELCOO(VNAME, V)                                                  \
   void delSparseTensorCOO##VNAME(void *coo) {                                  \
     delete static_cast<SparseTensorCOO<V> *>(coo);                             \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_DELCOO)
 #undef IMPL_DELCOO
 
+#define IMPL_DELITER(VNAME, V)                                                 \
+  void delSparseTensorIterator##VNAME(void *iter) {                            \
+    delete static_cast<SparseTensorIterator<V> *>(iter);                       \
+  }
+MLIR_SPARSETENSOR_FOREVERY_V(IMPL_DELITER)
+#undef IMPL_DELITER
+
 char *getTensorFilename(index_type id) {
   char var[80];
   sprintf(var, "TENSOR%" PRIu64, id);
   char *env = getenv(var);
   if (!env)
     MLIR_SPARSETENSOR_FATAL("Environment variable %s is not set\n", var);
   return env;
 }

mlir/test/Dialect/SparseTensor/sparse_concat.mlir

 // CHECK:         } do {
 // CHECK:           %[[TMP_12:.*]] = memref.load %[[TMP_8]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           %[[TMP_13:.*]] = arith.addi %[[TMP_12]], %[[TMP_c2]] : index
 // CHECK:           %[[TMP_14:.*]] = memref.load %[[TMP_8]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           %[[TMP_15:.*]] = memref.load %[[TMP_10]][] : memref<f64>
 // CHECK:           memref.store %[[TMP_15]], %[[TMP_0]][%[[TMP_13]], %[[TMP_14]]] : memref<5x4xf64>
 // CHECK:           scf.yield
 // CHECK:         }
-// CHECK:         call @delSparseTensorCOOF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
+// CHECK:         call @delSparseTensorIteratorF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         %[[TMP_11:.*]] = bufferization.to_tensor %[[TMP_0]] : memref<5x4xf64>
 // CHECK:         return %[[TMP_11]] : tensor<5x4xf64>
 // CHECK:       }
 func.func @concat_mix_dense(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #SparseMatrix>) -> tensor<5x4xf64> {
   %0 = sparse_tensor.concatenate %arg0, %arg1 {dimension = 0 : index}
        : tensor<2x4xf64>, tensor<3x4xf64, #SparseMatrix> to tensor<5x4xf64>
   return %0 : tensor<5x4xf64>
 }

 // CHECK:           %[[TMP_22:.*]] = memref.load %[[TMP_18]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           %[[TMP_23:.*]] = arith.addi %[[TMP_22]], %[[TMP_c2]] : index
 // CHECK:           %[[TMP_24:.*]] = memref.load %[[TMP_18]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           memref.store %[[TMP_23]], %[[TMP_9]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           memref.store %[[TMP_24]], %[[TMP_9]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           %[[TMP_25:.*]] = func.call @addEltF64(%[[TMP_7]], %[[TMP_20]], %[[TMP_10]], %[[TMP_5]]) : (!llvm.ptr<i8>, memref<f64>, memref<?xindex>, memref<?xindex>) -> !llvm.ptr<i8>
 // CHECK:           scf.yield
 // CHECK:         }
-// CHECK:         call @delSparseTensorCOOF64(%[[TMP_17]]) : (!llvm.ptr<i8>) -> ()
+// CHECK:         call @delSparseTensorIteratorF64(%[[TMP_17]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         %[[TMP_21:.*]] = call @newSparseTensor(%[[TMP_1]], %[[TMP_3]], %[[TMP_5]], %[[TMP_c0_i32]], %[[TMP_c0_i32]], %[[TMP_c1_i32]], %[[TMP_c2_i32]], %[[TMP_7]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 // CHECK:         call @delSparseTensorCOOF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         return %[[TMP_21]] : !llvm.ptr<i8>
 // CHECK:       }
 func.func @concat_mix_sparse(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #SparseMatrix>) -> tensor<5x4xf64, #SparseMatrix> {
   %0 = sparse_tensor.concatenate %arg0, %arg1 {dimension = 0 : index}
        : tensor<2x4xf64>, tensor<3x4xf64, #SparseMatrix> to tensor<5x4xf64, #SparseMatrix>
   return %0 : tensor<5x4xf64, #SparseMatrix>

 // CHECK:           %[[TMP_22:.*]] = memref.load %[[TMP_18]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           %[[TMP_23:.*]] = memref.load %[[TMP_18]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           %[[TMP_24:.*]] = arith.addi %[[TMP_23]], %[[TMP_c2]] : index
 // CHECK:           memref.store %[[TMP_22]], %[[TMP_9]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           memref.store %[[TMP_24]], %[[TMP_9]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           %[[TMP_25:.*]] = func.call @addEltF64(%[[TMP_7]], %[[TMP_20]], %[[TMP_10]], %[[TMP_5]]) : (!llvm.ptr<i8>, memref<f64>, memref<?xindex>, memref<?xindex>) -> !llvm.ptr<i8>
 // CHECK:           scf.yield
 // CHECK:         }
-// CHECK:         call @delSparseTensorCOOF64(%[[TMP_17]]) : (!llvm.ptr<i8>) -> ()
+// CHECK:         call @delSparseTensorIteratorF64(%[[TMP_17]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         %[[TMP_21:.*]] = call @newSparseTensor(%[[TMP_1]], %[[TMP_3]], %[[TMP_5]], %[[TMP_c0_i32]], %[[TMP_c0_i32]], %[[TMP_c1_i32]], %[[TMP_c2_i32]], %[[TMP_7]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 // CHECK:         call @delSparseTensorCOOF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         return %[[TMP_21]] : !llvm.ptr<i8>
 // CHECK:       }
 func.func @concat_mix_sparse_perm_dim1(%arg0: tensor<4x2xf64>, %arg1: tensor<4x3xf64, #SparseMatrix_P>) -> tensor<4x5xf64, #SparseMatrix_P> {
   %0 = sparse_tensor.concatenate %arg0, %arg1 {dimension = 1 : index}
        : tensor<4x2xf64>, tensor<4x3xf64, #SparseMatrix_P> to tensor<4x5xf64, #SparseMatrix_P>
   return %0 : tensor<4x5xf64, #SparseMatrix_P>

 // CHECK:         } do {
 // CHECK:           %[[TMP_12:.*]] = memref.load %[[TMP_8]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:           %[[TMP_13:.*]] = memref.load %[[TMP_8]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:           %[[TMP_14:.*]] = arith.addi %[[TMP_13]], %[[TMP_c2]] : index
 // CHECK:           %[[TMP_15:.*]] = memref.load %[[TMP_10]][] : memref<f64>
 // CHECK:           memref.store %[[TMP_15]], %[[TMP_0]][%[[TMP_12]], %[[TMP_14]]] : memref<4x5xf64>
 // CHECK:           scf.yield
 // CHECK:         }
-// CHECK:         call @delSparseTensorCOOF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
+// CHECK:         call @delSparseTensorIteratorF64(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:         %[[TMP_11:.*]] = bufferization.to_tensor %[[TMP_0]] : memref<4x5xf64>
 // CHECK:         return %[[TMP_11]] : tensor<4x5xf64>
 // CHECK:       }
 func.func @concat_mix_dense_perm_dim1(%arg0: tensor<4x2xf64>, %arg1: tensor<4x3xf64, #SparseMatrix_P>) -> tensor<4x5xf64> {
   %0 = sparse_tensor.concatenate %arg0, %arg1 {dimension = 1 : index}
        : tensor<4x2xf64>, tensor<4x3xf64, #SparseMatrix_P> to tensor<4x5xf64>
   return %0 : tensor<4x5xf64>
 }

 // CHECK:           } do {
 // CHECK:             %[[TMP_13:.*]] = memref.load %[[TMP_9]][%[[TMP_c0]]] : memref<2xindex>
 // CHECK:             %[[TMP_14:.*]] = memref.load %[[TMP_9]][%[[TMP_c1]]] : memref<2xindex>
 // CHECK:             %[[TMP_15:.*]] = arith.addi %[[TMP_14]], %[[TMP_c2]] : index
 // CHECK:             %[[TMP_16:.*]] = memref.load %[[TMP_11]][] : memref<f64>
 // CHECK:             memref.store %[[TMP_16]], %[[TMP_0]][%[[TMP_13]], %[[TMP_15]]] : memref<3x5xf64>
 // CHECK:             scf.yield
 // CHECK:           }
-// CHECK:           call @delSparseTensorCOOF64(%[[TMP_8]]) : (!llvm.ptr<i8>) -> ()
+// CHECK:           call @delSparseTensorIteratorF64(%[[TMP_8]]) : (!llvm.ptr<i8>) -> ()
 // CHECK:           %[[TMP_12:.*]] = bufferization.to_tensor %[[TMP_1]] : memref<?x?xf64>
 // CHECK:           return %[[TMP_12]] : tensor<?x?xf64>
 // CHECK:         }
 // CHECK:       }
 func.func @concat_mix_dense_perm_dim1_dyn(%arg0: tensor<3x2xf64>, %arg1: tensor<3x3xf64, #SparseMatrix>) -> tensor<?x?xf64> {
   %0 = sparse_tensor.concatenate %arg0, %arg1 {dimension = 1 : index}
        : tensor<3x2xf64>, tensor<3x3xf64, #SparseMatrix> to tensor<?x?xf64>
   return %0 : tensor<?x?xf64>
 }

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // CHECK-CONV:        %[[R:.*]] = arith.remui %[[X]], %[[C10]] : index
 // CHECK-CONV:        memref.store %[[D]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        memref.store %[[R]], %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
-// CHECK-CONV:      call @delSparseTensorCOOF64
+// CHECK-CONV:      call @delSparseTensorIteratorF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
 // rewrite for codegen:
 //
 // CHECK-RWT-LABEL:   func.func @sparse_expand(
 // CHECK-RWT-SAME:    %[[S:.*]]:
 // CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
 // CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index

 // CHECK-CONV:        %[[M:.*]] = arith.muli %[[X]], %[[C10]] : index
 // CHECK-CONV:        %[[A:.*]] = arith.addi %[[M]], %[[Y]] : index
 // CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
-// CHECK-CONV:      call @delSparseTensorCOOF64
+// CHECK-CONV:      call @delSparseTensorIteratorF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
 // rewrite for codegen:
 //
 // CHECK-RWT-LABEL:   func.func @sparse_collapse(
 // CHECK-RWT-SAME:    %[[S:.*]]:
 // CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
 // CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index

 // CHECK-CONV:        %[[D5:.*]] = arith.divui %[[R]], %[[D4]] : index
 // CHECK-CONV:        memref.store %[[D3]], %{{.*}}[%[[C0]]] : memref<2xindex>
 // CHECK-CONV:        memref.store %[[D5]], %{{.*}}[%[[C1]]] : memref<2xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
-// CHECK-CONV:      call @delSparseTensorCOOF64
+// CHECK-CONV:      call @delSparseTensorIteratorF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
 // rewrite for codegen:
 //
 // CHECK-RWT-LABEL:   func.func @dynamic_sparse_expand(
 // CHECK-RWT-SAME:    %[[S:.*]]:
 // CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
 // CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index

 // CHECK-CONV:        %[[M3:.*]] = arith.muli %[[Y]], %[[D2]] : index
 // CHECK-CONV:        %[[A:.*]] = arith.addi %[[M2]], %[[M3]] : index
 // CHECK-CONV:        memref.store %[[A]], %{{.*}}[%[[C0]]] : memref<1xindex>
 // CHECK-CONV:        call @addEltF64
 // CHECK-CONV:        scf.yield
 // CHECK-CONV:      }
 // CHECK-CONV:      %[[N:.*]] = call @newSparseTensor
 // CHECK-CONV:      call @delSparseTensorCOOF64
-// CHECK-CONV:      call @delSparseTensorCOOF64
+// CHECK-CONV:      call @delSparseTensorIteratorF64
 // CHECK-CONV:      return %[[N]] : !llvm.ptr<i8>
 //
 // rewrite for codegen:
 //
 // CHECK-RWT-LABEL:   func.func @dynamic_sparse_collapse(
 // CHECK-RWT-SAME:    %[[S:.*]]:
 // CHECK-RWT-DAG:     %[[C10:.*]] = arith.constant 10 : index
 // CHECK-RWT-DAG:     %[[C0:.*]] = arith.constant 0 : index