[Mlir-commits] [mlir] [mlir][sparse] use shared value storage between wrapped iterator and the wrapper. (PR #80046)

[Aart Bik]

================
@@ -205,34 +205,48 @@ static Value offsetFromMinCrd(OpBuilder &b, Location l, Value minCrd,
 
 namespace {
 
-class TrivialIterator : public SparseIterator {
-  Value getLoopLo(OpBuilder &b, Location l) const {
-    // Dense loop are traversed by coordinate, delinearize the position to get
-    // the coordinate.
-    if (randomAccessible())
-      return SUBI(itPos, posLo);
-    return itPos;
+// The iterator that that traverse a concrete sparse tensor levels. High-level
+// abstract iterators wrap it to achieve more complex goals (such as collapsing
+// several levels). It also holds the common storage to hold the mlir::Values
+// for itself as well as for wrappers.
+class ConcreteIterator : public SparseIterator {
+protected:
+  ConcreteIterator(const SparseTensorLevel &stl, IterKind kind,
+                   unsigned itValCnt)
+      : SparseIterator(kind, stl.tid, stl.lvl, itValCnt, itValsStorage),
+        stl(stl) {
+    // Allocate enough storage for iterator values.
+    itValsStorage.resize(itValCnt);
   }
 
 public:
-  TrivialIterator(const SparseTensorLevel &stl,
-                  const IterKind kind = IterKind::kTrivial)
-      : SparseIterator(kind, stl.tid, stl.lvl, itPos), stl(stl) {}
-
   // For LLVM-style RTTI.
   static bool classof(const SparseIterator *from) {
     return from->kind == IterKind::kTrivial;
   }
 
   bool randomAccessible() const override { return isDenseLT(stl.getLT()); };
-  bool iteratableByFor() const override { return true; };
+  bool iteratableByFor() const override { return kind != IterKind::kDedup; };
   Value upperBound(OpBuilder &b, Location l) const override {
     return stl.size();
   };
 
+protected:
+  // Owner of the storage, all wrappers build on top of a concrete iterator
+  // shares the same storage such that the iterator values are always
----------------
aartbik wrote:

shares -> share (subject is plural)

https://github.com/llvm/llvm-project/pull/80046