revert unintended change

Author: Peiming Liu

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorLevel.cpp

@@ -45,6 +45,19 @@ using ValueTuple = std::tuple<Value, Value, Value>;
 //===----------------------------------------------------------------------===//
 
 namespace {
+class SparseLevel : public SparseTensorLevel {
+public:
+  SparseLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
+              Value crdBuffer)
+      : SparseTensorLevel(tid, lvl, lt, lvlSize), crdBuffer(crdBuffer) {}
+
+  Value peekCrdAt(OpBuilder &b, Location l, Value iv) const override {
+    return genIndexLoad(b, l, crdBuffer, iv);
+  }
+
+protected:
+  const Value crdBuffer;
+};
 
 class DenseLevel : public SparseTensorLevel {
 public:
@@ -60,27 +73,53 @@ class DenseLevel : public SparseTensorLevel {
                         Value max) const override {
     assert(max == nullptr && "Dense level can not be non-unique.");
     if (encoded) {
-      Value posLo = MULI(p, getSize());
-      return {posLo, getSize()};
+      Value posLo = MULI(p, lvlSize);
+      return {posLo, lvlSize};
     }
     // No need to linearize the position for non-annotated tensors.
-    return {C_IDX(0), getSize()};
+    return {C_IDX(0), lvlSize};
   }
 
   const bool encoded;
 };
 
-class SparseLevel : public SparseTensorLevel {
+class CompressedLevel : public SparseLevel {
 public:
-  SparseLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
-              ValueRange lvlBuf)
-      : SparseTensorLevel(tid, lvl, lt, lvlSize, lvlBuf) {
-    assert(!lvlBuf.empty());
+  CompressedLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
+                  Value posBuffer, Value crdBuffer)
+      : SparseLevel(tid, lvl, lt, lvlSize, crdBuffer), posBuffer(posBuffer) {}
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    if (max == nullptr) {
+      Value pLo = genIndexLoad(b, l, posBuffer, p);
+      Value pHi = genIndexLoad(b, l, posBuffer, ADDI(p, C_IDX(1)));
+      return {pLo, pHi};
+    }
+    llvm_unreachable("compressed-nu should be the first non-unique level.");
   }
 
-  Value peekCrdAt(OpBuilder &b, Location l, Value iv) const override {
-    return genIndexLoad(b, l, getLvlBufs().front(), iv);
+private:
+  const Value posBuffer;
+};
+
+class LooseCompressedLevel : public SparseLevel {
+public:
+  LooseCompressedLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
+                       Value posBuffer, Value crdBuffer)
+      : SparseLevel(tid, lvl, lt, lvlSize, crdBuffer), posBuffer(posBuffer) {}
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    assert(max == nullptr && "loss compressed level can not be non-unique.");
+    p = MULI(p, C_IDX(2));
+    Value pLo = genIndexLoad(b, l, posBuffer, p);
+    Value pHi = genIndexLoad(b, l, posBuffer, ADDI(p, C_IDX(1)));
+    return {pLo, pHi};
   }
+
+private:
+  const Value posBuffer;
 };
 
 class SingletonLevel : public SparseLevel {
@@ -102,8 +141,8 @@ class SingletonLevel : public SparseLevel {
 class TwoOutFourLevel : public SparseLevel {
 public:
   TwoOutFourLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
-                  Value crdBuf)
-      : SparseLevel(tid, lvl, lt, lvlSize, crdBuf) {}
+                  Value crdBuffer)
+      : SparseLevel(tid, lvl, lt, lvlSize, crdBuffer) {}
 
   ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
                         Value max) const override {
@@ -114,39 +153,6 @@ class TwoOutFourLevel : public SparseLevel {
   }
 };
 
-class CompressedLevel : public SparseLevel {
-public:
-  CompressedLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
-                  Value posBuffer, Value crdBuffer)
-      : SparseLevel(tid, lvl, lt, lvlSize, {crdBuffer, posBuffer}) {}
-
-  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
-                        Value max) const override {
-    if (max == nullptr) {
-      Value pLo = genIndexLoad(b, l, getPosBuf(), p);
-      Value pHi = genIndexLoad(b, l, getPosBuf(), ADDI(p, C_IDX(1)));
-      return {pLo, pHi};
-    }
-    llvm_unreachable("compressed-nu should be the first non-unique level.");
-  }
-};
-
-class LooseCompressedLevel : public SparseLevel {
-public:
-  LooseCompressedLevel(unsigned tid, Level lvl, LevelType lt, Value lvlSize,
-                       Value posBuffer, Value crdBuffer)
-      : SparseLevel(tid, lvl, lt, lvlSize, {crdBuffer, posBuffer}) {}
-
-  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
-                        Value max) const override {
-    assert(max == nullptr && "loss compressed level can not be non-unique.");
-    p = MULI(p, C_IDX(2));
-    Value pLo = genIndexLoad(b, l, getPosBuf(), p);
-    Value pHi = genIndexLoad(b, l, getPosBuf(), ADDI(p, C_IDX(1)));
-    return {pLo, pHi};
-  }
-};
-
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -195,9 +201,7 @@ static Value offsetFromMinCrd(OpBuilder &b, Location l, Value minCrd,
 //===----------------------------------------------------------------------===//
 // SparseIterator derived classes.
 //===----------------------------------------------------------------------===//
-
-namespace mlir {
-namespace sparse_tensor {
+namespace {
 
 // The iterator that traverses a concrete sparse tensor levels. High-level
 // abstract iterators wrap it to achieve more complex goals (such as collapsing
@@ -232,11 +236,6 @@ class ConcreteIterator : public SparseIterator {
   SmallVector<Value> cursorValsStorage;
 };
 
-} // namespace sparse_tensor
-} // namespace mlir
-
-namespace {
-
 class TrivialIterator : public ConcreteIterator {
 public:
   TrivialIterator(const SparseTensorLevel &stl)

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorLevel.h

@@ -15,11 +15,9 @@
 namespace mlir {
 namespace sparse_tensor {
 
-class ConcreteIterator;
-
-/// The base class for all types of sparse tensor levels. It provides
-/// interfaces to query the loop range (see `peekRangeAt`) and look up the
-/// coordinates (see `peekCrdAt`).
+/// The base class for all types of sparse tensor levels. It provides interfaces
+/// to query the loop range (see `peekRangeAt`) and look up the coordinates (see
+/// `peekCrdAt`).
 class SparseTensorLevel {
   SparseTensorLevel(SparseTensorLevel &&) = delete;
   SparseTensorLevel(const SparseTensorLevel &) = delete;
@@ -33,6 +31,7 @@ class SparseTensorLevel {
     return std::string(toMLIRString(lt)) + "[" + std::to_string(tid) + "," +
            std::to_string(lvl) + "]";
   }
+
   virtual Value peekCrdAt(OpBuilder &b, Location l, Value iv) const = 0;
 
   /// Peeks the lower and upper bound to *fully* traverse the level with
@@ -53,17 +52,7 @@ class SparseTensorLevel {
 
   Level getLevel() const { return lvl; }
   LevelType getLT() const { return lt; }
-  Value getSize() const { return lvlVals.front(); }
-  Value getCrdBuf() const {
-    assert(lvlVals.size() > 1);
-    return lvlVals[1];
-  }
-  Value getPosBuf() const {
-    assert(lvlVals.size() > 2);
-    return lvlVals[2];
-  }
-  ValueRange getLvlVals() const { return lvlVals; }
-  ValueRange getLvlBufs() const { return ValueRange(lvlVals).drop_front(); }
+  Value getSize() const { return lvlSize; }
 
   //
   // Level properties
@@ -72,24 +61,12 @@ class SparseTensorLevel {
 
 protected:
   SparseTensorLevel(unsigned tid, unsigned lvl, LevelType lt, Value lvlSize)
-      : tid(tid), lvl(lvl), lt(lt), lvlVals() {
-    lvlVals.push_back(lvlSize);
-  };
-
-  SparseTensorLevel(unsigned tid, unsigned lvl, LevelType lt, Value lvlSize,
-                    ValueRange lvlBufs)
-      : tid(tid), lvl(lvl), lt(lt), lvlVals() {
-    lvlVals.push_back(lvlSize);
-    lvlVals.append(lvlBufs.begin(), lvlBufs.end());
-  };
+      : tid(tid), lvl(lvl), lt(lt), lvlSize(lvlSize){};
 
 public:
   const unsigned tid, lvl;
   const LevelType lt;
-  // The first value in the vector is always lvlsize; for sparse levels, the
-  // second value is always the coordinate buffer; for sparse level with
-  // position buffers, the third value is always the position buffer.
-  SmallVector<Value, 3> lvlVals;
+  const Value lvlSize;
 };
 
 enum class IterKind : uint8_t {