[mlir][sparse] support loop range query using SparseTensorLevel. (#75670)

Author: PeimingLiu

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

@@ -244,12 +244,12 @@ Value LoopEmitter::genAddress(OpBuilder &builder, Location loc, TensorId tid,
 Value LoopEmitter::genSegmentHigh(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl, Value pLo,
                                   Value pHi) {
-  SparseTensorLevel &level = *lvls[tid][lvl];
-  const Value sameCrd = level.peekCrdAt(builder, loc, pLo);
+  SparseTensorLevel &stl = *lvls[tid][lvl];
+  const Value sameCrd = stl.peekCrdAt(builder, loc, pLo);
   auto whileOp = builder.create<scf::WhileOp>(
       loc, builder.getIndexType(), pLo,
       /*beforeBuilder=*/
-      [pHi, &level, sameCrd](OpBuilder &builder, Location loc, ValueRange ivs) {
+      [pHi, &stl, sameCrd](OpBuilder &builder, Location loc, ValueRange ivs) {
         const auto pos = ivs[0];
         Value inBound = builder.create<arith::CmpIOp>(
             loc, arith::CmpIPredicate::ult, pos, pHi);
@@ -260,7 +260,7 @@ Value LoopEmitter::genSegmentHigh(OpBuilder &builder, Location loc,
           // Load the next coordinates only when inbound (to avoid OOB
           // accesses).
           builder.setInsertionPointToStart(ifInBound.thenBlock());
-          Value crd = level.peekCrdAt(builder, loc, pos);
+          Value crd = stl.peekCrdAt(builder, loc, pos);
           Value isSameCrd = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::eq, crd, sameCrd);
           YIELD(isSameCrd);
@@ -1226,27 +1226,19 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
 
   const Value c0 = C_IDX(0);
   const Value c1 = C_IDX(1);
-  const Value c2 = C_IDX(2);
   // Either the first level, or the previous level has been set.
   /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
   assert(lvl == 0 || posits[tid][lvl - 1]);
-  if (isCompressedLT(lvlTp) || isLooseCompressedLT(lvlTp)) {
-    // TODO: eliminate the cast upon feature complete.
-    const Value mem =
-        isCompressedLT(lvlTp)
-            ? static_cast<CompressedLevel &>(*lvls[tid][lvl]).posBuffer
-            : static_cast<LooseCompressedLevel &>(*lvls[tid][lvl]).posBuffer;
-
-    Value pLo = lvl == 0 ? c0 : posits[tid][lvl - 1];
-    if (isLooseCompressedLT(lvlTp))
-      pLo = builder.create<arith::MulIOp>(loc, pLo, c2);
-    posits[tid][lvl] = genIndexLoad(builder, loc, mem, pLo);
-
-    const Value pHi = ADDI(pLo, c1);
-    highs[tid][lvl] = genIndexLoad(builder, loc, mem, pHi);
+  if (isCompressedLT(lvlTp) || isLooseCompressedLT(lvlTp) ||
+      is2OutOf4LT(lvlTp)) {
+
+    Value pos = lvl == 0 ? c0 : posits[tid][lvl - 1];
+    std::tie(posits[tid][lvl], highs[tid][lvl]) =
+        lvls[tid][lvl]->peekRangeAt(builder, loc, pos);
     return;
   }
   if (isSingletonLT(lvlTp)) {
+    // TODO: merge this as well when SparseTensorLevel support dedup.
     const Value pLo = lvl == 0 ? c0 : posits[tid][lvl - 1];
     posits[tid][lvl] = pLo;
 
@@ -1262,13 +1254,6 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                           : ADDI(pLo, c1);
     return;
   }
-  if (is2OutOf4LT(lvlTp)) {
-    const Value pLo = lvl == 0 ? c0 : posits[tid][lvl - 1];
-    // Each 2:4 block has exactly two specified elements.
-    posits[tid][lvl] = MULI(pLo, c2);
-    highs[tid][lvl] = ADDI(posits[tid][lvl], c2);
-    return;
-  }
   llvm_unreachable("Unrecognized level-type!");
 }
 
@@ -1824,18 +1809,11 @@ void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
   auto [nxSz, stride] = sliceMeta[tid][lvl][1];
   assert(stride == 1 && "Not yet implemented");
   Value sPtrBuf = slicePosBuffer[tid][lvl][0];
-  Value pHi, pLo;
-  if (lvl == 0) {
-    pLo = c0;
-    // TODO: eliminate the cast upon feature complete.pLo = c0;
-    Value pBuf = static_cast<CompressedLevel &>(*lvls[tid][0]).posBuffer;
-    pHi = genIndexLoad(builder, loc, pBuf, c1);
-  } else {
-    // TODO: eliminate the cast upon feature complete.} else {
-    Value pBuf = static_cast<CompressedLevel &>(*lvls[tid][lvl]).posBuffer;
-    pLo = genIndexLoad(builder, loc, pBuf, posits[tid][lvl - 1]);
-    pHi = genIndexLoad(builder, loc, pBuf, ADDI(posits[tid][lvl - 1], c1));
-  }
+  const SparseTensorLevel &stl = *lvls[tid][lvl];
+
+  Value p = lvl == 0 ? c0 : posits[tid][lvl - 1];
+  auto [pLo, pHi] = stl.peekRangeAt(builder, loc, p);
+
   // Fills out pIdxBuffer[tid][lvl][0] with [pLo, pHi]
   updateSlicePos(builder, loc, sPtrBuf, pLo, c0, SlicePosKind::kLo);
   updateSlicePos(builder, loc, sPtrBuf, pHi, c0, SlicePosKind::kHi);
@@ -1849,7 +1827,7 @@ void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
   // nonempty. though we assume that even on empty sparse tensors, a non-empty
   // ptr/idx buffer is allocated for each level so it would not cause OOB to
   // avoid generating a ifOp here.
-  Value minCrd = lvls[tid][lvl]->peekCrdAt(builder, loc, pLo);
+  Value minCrd = stl.peekCrdAt(builder, loc, pLo);
 
   // FIXME: We need the relative offset related to the base slice.
   Value absOffset = offsetFromMinCoord(builder, loc, minCrd, nxSz, isNonEmpty);
@@ -1879,7 +1857,7 @@ void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
 // }
 void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
                                           TensorId tid, Level lvl) {
-  Value c0 = C_IDX(0), c1 = C_IDX(1);
+  Value c0 = C_IDX(0);
   unsigned depth = levelReducedDep[tid][lvl];
   // The remaining slice size after reduction.
   Value remSz = sliceMeta[tid][lvl][depth + 1].first;
@@ -1929,17 +1907,14 @@ void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
 
   ValueRange result = genUnResolvedSliceTreeTraverse(
       builder, loc, tid, unResSlices, firstResLvl, reduc,
-      [this, c1, tid, lvl, sPtrBuf](OpBuilder &builder, Location loc, Value iv,
-                                    MutableArrayRef<Value> reduc) {
+      [this, tid, lvl, sPtrBuf](OpBuilder &builder, Location loc, Value iv,
+                                MutableArrayRef<Value> reduc) {
         Value &nonEmpty = reduc[0];
         Value &minCrd = reduc[1];
         Value &curTupleCnt = reduc[2];
 
-        Value pHi = ADDI(iv, c1);
-        // TODO: eliminate the cast upon feature complete.
-        Value pBuf = static_cast<CompressedLevel &>(*lvls[tid][lvl]).posBuffer;
-        Value sPLo = genIndexLoad(builder, loc, pBuf, iv);
-        Value sPHi = genIndexLoad(builder, loc, pBuf, pHi);
+        const SparseTensorLevel &stl = *lvls[tid][lvl];
+        auto [sPLo, sPHi] = stl.peekRangeAt(builder, loc, iv);
 
         // isNonEmpty = isNonEmpty || lvlNonEmpty, i.e., as long as there is
         // one non-empty lvl, the slice is non-empty.
@@ -1957,7 +1932,7 @@ void LoopEmitter::genUnResolvedSliceBegin(OpBuilder &builder, Location loc,
           // }
           OpBuilder::InsertionGuard guard(builder);
           builder.setInsertionPointToStart(ifNonEmpty.thenBlock());
-          Value curC = lvls[tid][lvl]->peekCrdAt(builder, loc, sPLo);
+          Value curC = stl.peekCrdAt(builder, loc, sPLo);
           Value isSmaller = CMPI(ult, curC, minCrd);
           Value newMin = SELECT(isSmaller, curC, minCrd);
           YIELD(newMin);

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

@@ -13,6 +13,133 @@
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
+using ValuePair = std::pair<Value, Value>;
+
+//===----------------------------------------------------------------------===//
+// File local helper functions/macros.
+//===----------------------------------------------------------------------===//
+#define CMPI(p, lhs, rhs)                                                      \
+  (b.create<arith::CmpIOp>(l, arith::CmpIPredicate::p, (lhs), (rhs)))
+
+#define C_IDX(v) (constantIndex(b, l, (v)))
+#define YIELD(vs) (b.create<scf::YieldOp>(l, (vs)))
+#define ADDI(lhs, rhs) (b.create<arith::AddIOp>(l, (lhs), (rhs)))
+#define ANDI(lhs, rhs) (b.create<arith::AndIOp>(l, (lhs), (rhs)))
+#define SUBI(lhs, rhs) (b.create<arith::SubIOp>(l, (lhs), (rhs)))
+#define MULI(lhs, rhs) (b.create<arith::MulIOp>(l, (lhs), (rhs)))
+#define REMUI(lhs, rhs) (b.create<arith::RemUIOp>(l, (lhs), (rhs)))
+#define DIVUI(lhs, rhs) (b.create<arith::DivUIOp>(l, (lhs), (rhs)))
+#define SELECT(c, lhs, rhs) (b.create<arith::SelectOp>(l, (c), (lhs), (rhs)))
+
+static ValuePair constantRange(OpBuilder &b, Location l, Value lo, Value sz) {
+  return std::make_pair(lo, ADDI(lo, sz));
+}
+
+//===----------------------------------------------------------------------===//
+// SparseTensorLevel derived classes.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class SparseLevel : public SparseTensorLevel {
+public:
+  SparseLevel(LevelType lt, Value lvlSize, Value crdBuffer)
+      : SparseTensorLevel(lt, lvlSize), crdBuffer(crdBuffer) {}
+
+  Value peekCrdAt(OpBuilder &b, Location l, Value pos) const override {
+    return genIndexLoad(b, l, crdBuffer, pos);
+  }
+
+protected:
+  const Value crdBuffer;
+};
+
+class DenseLevel : public SparseTensorLevel {
+public:
+  DenseLevel(Value lvlSize) : SparseTensorLevel(LevelType::Dense, lvlSize) {
+    // Dense level, loop upper bound equals to the level size.
+    loopHi = lvlSize;
+  }
+
+  Value peekCrdAt(OpBuilder &, Location, Value pos) const override {
+    return pos;
+  }
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    assert(max == nullptr && "Dense level can not be non-unique.");
+    return constantRange(b, l, C_IDX(0), lvlSize);
+  }
+};
+
+class CompressedLevel : public SparseLevel {
+public:
+  CompressedLevel(LevelType lt, Value lvlSize, Value posBuffer, Value crdBuffer)
+      : SparseLevel(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("TODO: dedup not implemented");
+  }
+
+private:
+  const Value posBuffer;
+};
+
+class LooseCompressedLevel : public SparseLevel {
+public:
+  LooseCompressedLevel(LevelType lt, Value lvlSize, Value posBuffer,
+                       Value crdBuffer)
+      : SparseLevel(lt, lvlSize, crdBuffer), posBuffer(posBuffer) {}
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    // Allows this?
+    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 {
+public:
+  SingletonLevel(LevelType lt, Value lvlSize, Value crdBuffer)
+      : SparseLevel(lt, lvlSize, crdBuffer) {}
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    if (max == nullptr)
+      return constantRange(b, l, p, C_IDX(1));
+    llvm_unreachable("TODO: dedup not implemented");
+  }
+};
+
+class TwoOutFourLevel : public SparseLevel {
+public:
+  TwoOutFourLevel(LevelType lt, Value lvlSize, Value crdBuffer)
+      : SparseLevel(lt, lvlSize, crdBuffer) {}
+
+  ValuePair peekRangeAt(OpBuilder &b, Location l, Value p,
+                        Value max) const override {
+    assert(max == nullptr && "2:4 level can not be non-unique.");
+    // Each 2:4 block has exactly two specified elements.
+    Value c2 = C_IDX(2);
+    return constantRange(b, l, MULI(p, c2), c2);
+  }
+};
+
+} // namespace
 
 std::unique_ptr<SparseTensorLevel>
 sparse_tensor::makeSparseTensorLevel(OpBuilder &builder, Location loc, Value t,
@@ -49,6 +176,11 @@ sparse_tensor::makeSparseTensorLevel(OpBuilder &builder, Location loc, Value t,
   llvm_unreachable("unrecognizable level format");
 }
 
-Value SparseLevel::peekCrdAt(OpBuilder &b, Location l, Value pos) const {
-  return genIndexLoad(b, l, crdBuffer, pos);
-}
+#undef CMPI
+#undef C_IDX
+#undef YIELD
+#undef ADDI
+#undef ANDI
+#undef SUBI
+#undef MULI
+#undef SELECT

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

@@ -17,13 +17,22 @@ namespace sparse_tensor {
 class SparseTensorLevel {
   SparseTensorLevel(SparseTensorLevel &&) = delete;
   SparseTensorLevel(const SparseTensorLevel &) = delete;
+  SparseTensorLevel &operator=(SparseTensorLevel &&) = delete;
+  SparseTensorLevel &operator=(const SparseTensorLevel &) = delete;
 
 public:
   SparseTensorLevel() : SparseTensorLevel(LevelType::Undef, nullptr){};
   virtual ~SparseTensorLevel() = default;
 
   virtual Value peekCrdAt(OpBuilder &b, Location l, Value p) const = 0;
 
+  /// Peeks the lower and upper bound to *fully* traverse the level with
+  /// the given position `p` that the immediate parent level is current at.
+  /// `bound` is only used when the level is `non-unique` and deduplication is
+  /// required. It specifies the max upper bound of the non-unique segment.
+  virtual std::pair<Value, Value> peekRangeAt(OpBuilder &b, Location l, Value p,
+                                              Value bound = Value()) const = 0;
+
   LevelType getLT() const { return lt; }
   Value getPos() const { return pos; }
   Value getCrd() const { return crd; }
@@ -49,60 +58,6 @@ class SparseTensorLevel {
 std::unique_ptr<SparseTensorLevel>
 makeSparseTensorLevel(OpBuilder &builder, Location loc, Value t, Level l);
 
-class DenseLevel : public SparseTensorLevel {
-public:
-  DenseLevel(Value lvlSize) : SparseTensorLevel(LevelType::Dense, lvlSize) {
-    // Dense level, loop upper bound equals to the level size.
-    loopHi = lvlSize;
-  }
-
-  Value peekCrdAt(OpBuilder &, Location, Value pos) const override {
-    return pos;
-  }
-};
-
-class SparseLevel : public SparseTensorLevel {
-public:
-  SparseLevel(LevelType lt, Value lvlSize, Value crdBuffer)
-      : SparseTensorLevel(lt, lvlSize), crdBuffer(crdBuffer) {}
-
-  Value peekCrdAt(OpBuilder &b, Location l, Value pos) const override;
-
-public: // TODO: make these values private upon feature complete.
-  const Value crdBuffer;
-};
-
-class CompressedLevel : public SparseLevel {
-public:
-  CompressedLevel(LevelType lt, Value lvlSize, Value posBuffer, Value crdBuffer)
-      : SparseLevel(lt, lvlSize, crdBuffer), posBuffer(posBuffer) {}
-
-public: // TODO: make these values private upon feature complete.
-  const Value posBuffer;
-};
-
-class LooseCompressedLevel : public SparseLevel {
-public:
-  LooseCompressedLevel(LevelType lt, Value lvlSize, Value posBuffer,
-                       Value crdBuffer)
-      : SparseLevel(lt, lvlSize, crdBuffer), posBuffer(posBuffer) {}
-
-public: // TODO: make these values private upon feature complete.
-  const Value posBuffer;
-};
-
-class SingletonLevel : public SparseLevel {
-public:
-  SingletonLevel(LevelType lt, Value lvlSize, Value crdBuffer)
-      : SparseLevel(lt, lvlSize, crdBuffer) {}
-};
-
-class TwoOutFourLevel : public SparseLevel {
-public:
-  TwoOutFourLevel(LevelType lt, Value lvlSize, Value crdBuffer)
-      : SparseLevel(lt, lvlSize, crdBuffer) {}
-};
-
 } // namespace sparse_tensor
 } // namespace mlir