[mlir][sparse] Introducing iterators for TensorId, LoopId, TensorLoopId

[wrengr]

Summary

Client code frequently wants to iterate over the values of these three types. This commit introduces the iterators themselves, whereas the next commit will perform the actual conversion from typedefs to newtypes.

Depends On D146690

[wrengr]

git-clang-format

[wrengr]

rebase

[wrengr]

rebase

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/Utils/MergerNewtypes.h

269–277

Isn't it too heavy for a such a simple concept? I would prefer a foreach function for the same purpose.

[wrengr]

rebasing for the postponement splitting of D146684

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/BitVector.h"
 #include <optional>
 
 namespace mlir {
 namespace sparse_tensor {
 
-// TODO: These type aliases currently only serve to make the code more
-// self-documenting, however because they are not type-checked they can
-// do nothing to prevent mixups.  They should be made into newtypes and
-// moved into MergerNewtypes.h with the other newtypes.
-
-/// Tensor identifiers.  The valid set of identifiers is defined by the
-/// first argument passed to the `Merger` ctor.
-using TensorId = unsigned;
-
-/// Loop identifiers.  The valid set of identifiers is defined by the
-/// second two arguments to the `Merger` ctor.
-///
-/// These identifiers serve as proxies for the `$dim` argument to
-/// `linalg::IndexOp`, however the numerical value of a `LoopId` should
-/// not necessarily be equated with the numerical value of the corresponding
-/// `$dim` argument.  The `$dim` arguments are De Bruijn indices: that
-/// is, they identify the loop which binds the loop-variable by counting
-/// the enclosing loops from innermost to outermost, starting from zero.
-/// Whereas `LoopId` are considered to be arbitrary names for identifying
-/// loops; since the `Merger` does not care about the actual ordering of
-/// loops, and leaves it up to the `LoopEmitter` to specify the actual
-/// loop ordering (`LoopOrd`).
-///
-/// TODO: Despite the above claim that `$dim` and `LoopId` need not be
-/// numerically equal, some code in the `Merger` class does equate them
-/// (e.g., `buildTensorExp`).  So we need to explicate the exact relationship
-/// between `$dim`, `LoopId`, and `LoopOrd`; especially with regards to their
-/// providence.  If `LoopId` really is supposed to be equated with `$dim`,
-/// then we should change the name to `LoopIdx` or similar, to capture the
-/// fact that its numerical value is not invariant when entering/exiting
-/// loops (unlike `TensorId`, `ExprId`, `LatPointId`, and `LatSetId` which
-/// are invariant identifiers).
-using LoopId = unsigned;
-
-/// A compressed representation of `std::pair<TensorId, LoopId>`.
-/// The compression scheme is such that this also serves as an index
-/// into the bitvector stored in `LatPoint` (since that bitvector is
-/// just the implementation for a set of `TensorLoopId` values).
-using TensorLoopId = unsigned;
-
 /// Tensor expression. Represents an MLIR expression in tensor index notation.
 struct TensorExp final {
   enum class Kind;
 
   /// Child subexpressions for non-leaf expressions.
   struct Children final {
     ExprId e0;
     ExprId e1;
public:
   /// Gets the tensor-identifier of the `TensorLoopId`.
   constexpr TensorId tensor(TensorLoopId b) const { return b % numTensors; }
   /// Gets the loop-identifier of the `TensorLoopId`.
   constexpr LoopId loop(TensorLoopId b) const { return b / numTensors; }
 
   /// Get the total number of tensors (including the output-tensor and
   /// synthetic-tensor).
   constexpr unsigned getNumTensors() const { return numTensors; }
+  /// Get the range of all tensor identifiers.
+  constexpr tensor_id::Range getTensorIds() const {
+    return tensor_id::Range(0, numTensors);
+  }
 
   /// Get the total number of loops (native loops + filter loops).
   constexpr unsigned getNumLoops() const { return numLoops; }
   /// Get the number of native loops.
   constexpr unsigned getNumNativeLoops() const { return numNativeLoops; }
   /// Get the number of filter loops.
   constexpr unsigned getNumFilterLoops() const {
     return numLoops - numNativeLoops;
   }
-  /// Get the identifier of the first filter-loop.
-  constexpr LoopId getStartingFilterLoopId() const {
-    return getNumNativeLoops();
+  /// Get the range of all loop identifiers.
+  constexpr loop_id::Range getLoopIds() const {
+    return loop_id::Range(0, numLoops);
+  }
+  /// Get the range of native-loop identifiers.
+  constexpr loop_id::Range getNativeLoopIds() const {
+    return loop_id::Range(0, numNativeLoops);
+  }
+  /// Get the range of filter-loop identifiers.
+  constexpr loop_id::Range getFilterLoopIds() const {
+    return loop_id::Range(numNativeLoops, numLoops);
   }
 
   /// Returns true if `b` is the `i`th loop of the output tensor.
   constexpr bool isOutTensor(TensorLoopId b, LoopId i) const {
     return b == makeTensorLoopId(outTensor, i);
   }
 
   /// Get the output tensor's identifier.

mlir/include/mlir/Dialect/SparseTensor/Utils/MergerNewtypes.h

 // values in registers rather than on the stack (cf.,
 // <https://www.agner.org/optimize/calling_conventions.pdf>).
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_UTILS_MERGERNEWTYPES_H_
 #define MLIR_DIALECT_SPARSETENSOR_UTILS_MERGERNEWTYPES_H_
 
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/iterator.h"
+
 #include <cassert>
 #include <type_traits>
 
 namespace mlir {
 namespace sparse_tensor {
 
 namespace detail {
 /// A constant serving as the canonically invalid identifier,
 /// regardless of the identifier type.
 static constexpr unsigned kInvalidId = -1u;
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
+/// Tensor identifiers.
+///
+/// Semantically, tensor identifiers could be chosen to be anything;
+/// but operationally, they must be chosen such that the `Merger`
+/// and `GenericOpSparsifier` agree.  Therefore, the numeric values of
+/// tensor identifiers are chosen to be the `BlockArgument::getArgNumber`
+/// of the value passed to `Merger::buildTensorExp`, which ranges from
+/// zero to `linalg::GenericOp::getNumOperands` for the op passed to
+/// `GenericOpSparsifier::matchAndRewrite`.
+using TensorId = unsigned;
+
+// NOTE: We use this namespace to simulate having turned `TensorId`
+// into a newtype, so that we can split the patch for adding the iterators
+// from the patch for actually making it a newtype.
+namespace tensor_id {
+class Iterator;
+class Range;
+} // namespace tensor_id
+
+/// An iterator for `TensorId`.  We define this as a separate class because
+/// it wouldn't be generally safe/meaningful to define `TensorId::operator++`.
+/// The ctor is private for similar reasons, so client code should create
+/// iterators via `tensor_id::Range` instead.
+class tensor_id::Iterator final
+    : public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                        TensorId> {
+  friend class tensor_id::Range;
+  explicit constexpr Iterator(TensorId tid) : tid(tid) {}
+
+public:
+  using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                   TensorId>::operator++;
+  Iterator &operator++() {
+    ++tid;
+    return *this;
+  }
+  const TensorId *operator->() const { return &tid; }
+  constexpr TensorId operator*() const { return tid; }
+  constexpr bool operator==(Iterator rhs) const { return tid == rhs.tid; }
+  constexpr bool operator!=(Iterator rhs) const { return !(*this == rhs); }
+
+private:
+  TensorId tid;
+};
+static_assert(std::is_trivially_copyable_v<tensor_id::Iterator> &&
+              std::is_trivially_destructible_v<tensor_id::Iterator>);
+
+/// An iterator range for `TensorId`.
+class tensor_id::Range final {
+public:
+  explicit constexpr Range(TensorId lo, TensorId hi)
+      : begin_(lo <= hi ? lo : hi), end_(hi) {}
+  constexpr Iterator begin() const { return begin_; }
+  constexpr Iterator end() const { return end_; }
+  constexpr bool empty() const { return begin_ == end_; }
+
+private:
+  Iterator begin_;
+  Iterator end_;
+};
+static_assert(std::is_trivially_copyable_v<tensor_id::Range> &&
+              std::is_trivially_destructible_v<tensor_id::Range>);
+
+//===----------------------------------------------------------------------===//
+/// Loop identifiers.
+///
+/// These identifiers serve as proxies for the `$dim` argument to
+/// `linalg::IndexOp`, however the numerical value of a `LoopId` should
+/// not necessarily be equated with the numerical value of the corresponding
+/// `$dim` argument.  The `$dim` arguments are De Bruijn indices: that
+/// is, they identify the loop which binds the loop-variable by counting
+/// the enclosing loops from innermost to outermost, starting from zero.
+/// Whereas `LoopId` are considered to be arbitrary names for identifying
+/// loops; since the `Merger` does not care about the actual ordering of
+/// loops, and leaves it up to the `LoopEmitter` to specify the actual
+/// loop ordering (`LoopOrd`).
+///
+/// TODO: Despite the above claim that `$dim` and `LoopId` need not be
+/// numerically equal, some code in the `Merger` class does equate them
+/// (e.g., `buildTensorExp`).  So we need to explicate the exact relationship
+/// between `$dim`, `LoopId`, and `LoopOrd`; especially with regards to their
+/// providence.  If `LoopId` really is supposed to be equated with `$dim`,
+/// then we should change the name to `LoopIdx` or similar, to capture the
+/// fact that its numerical value is not invariant when entering/exiting
+/// loops (unlike `TensorId`, `ExprId`, `LatPointId`, and `LatSetId` which
+/// are invariant identifiers).
+using LoopId = unsigned;
+
+// NOTE: We use this namespace to simulate having turned `LoopId` into
+// a newtype, so that we can split the patch for adding the iterators from
+// the patch for actually making it a newtype.
+namespace loop_id {
+class Iterator;
+class Range;
+} // namespace loop_id
+
+/// An iterator for `LoopId`.  We define this as a separate class because
+/// it wouldn't be generally safe/meaningful to define `LoopId::operator++`.
+/// The ctor is private for similar reasons, so client code should create
+/// iterators via `loop_id::Range` instead.
+class loop_id::Iterator final
+    : public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                        LoopId> {
+  friend class loop_id::Range;
+  explicit constexpr Iterator(LoopId i) : loop(i) {}
+
+public:
+  using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                   LoopId>::operator++;
+  Iterator &operator++() {
+    ++loop;
+    return *this;
+  }
+  const LoopId *operator->() const { return &loop; }
+  constexpr LoopId operator*() const { return loop; }
+  constexpr bool operator==(Iterator rhs) const { return loop == rhs.loop; }
+  constexpr bool operator!=(Iterator rhs) const { return !(*this == rhs); }
+
+private:
+  LoopId loop;
+};
+static_assert(std::is_trivially_copyable_v<loop_id::Iterator> &&
+              std::is_trivially_destructible_v<loop_id::Iterator>);
+
+/// An iterator range for `LoopId`.
+class loop_id::Range final {
+public:
+  explicit constexpr Range(LoopId lo, LoopId hi)
+      : begin_((lo != detail::kInvalidId && hi != detail::kInvalidId)
+                   ? (lo <= hi ? lo : hi)
+                   : detail::kInvalidId),
+        end_(lo != detail::kInvalidId ? hi : detail::kInvalidId) {}
+  constexpr Iterator begin() const { return begin_; }
+  constexpr Iterator end() const { return end_; }
+  constexpr bool empty() const { return begin_ == end_; }
+
+private:
+  Iterator begin_;
+  Iterator end_;
+};
+static_assert(std::is_trivially_copyable_v<loop_id::Range> &&
+              std::is_trivially_destructible_v<loop_id::Range>);
+
+//===----------------------------------------------------------------------===//
+/// A compressed representation of `std::pair<TensorId, LoopId>`.
+/// The compression scheme is such that this also serves as an index
+/// into the bitvector stored in `LatPoint` (since that bitvector is
+/// just the implementation for a set of `TensorLoopId` values).
+using TensorLoopId = unsigned;
+
+// NOTE: We use this namespace to simulate having turned `TensorLoopId` into
+// a newtype, so that we can split the patch for adding the iterators from
+// the patch for actually making it a newtype.
+namespace tensor_loop_id {
+class Iterator;
+class Range;
+} // namespace tensor_loop_id
+
+/// An iterator of the `TensorLoopId`s which are included/set in a `BitVector`.
+class tensor_loop_id::Iterator final
+    : public llvm::iterator_adaptor_base<
+          Iterator, llvm::BitVector::const_set_bits_iterator,
+          // Since `const_set_bits_iterator` doesn't define its own
+          // `std::iterator_traits`, we must manually do so here.
+          /*iterator_category=*/std::forward_iterator_tag,
+          /*value_type=*/unsigned,
+          /*difference_type=*/std::ptrdiff_t, // irrelevant, since not
+                                              // random-access
+          /*pointer=*/const unsigned *,
+          /*reference=*/const unsigned &> {
+  friend class tensor_loop_id::Range;
+  explicit Iterator(llvm::BitVector::const_set_bits_iterator I)
+      : Iterator::iterator_adaptor_base(I) {}
+
+public:
+  TensorLoopId operator*() const { return TensorLoopId{*this->I}; }
+};
+
+/// An iterator range for the `TensorLoopId`s which are included/set
+/// in a `BitVector`.
+class tensor_loop_id::Range final {
+public:
+  explicit Range(const llvm::BitVector &bits) : set_bits(bits.set_bits()) {}
+  Iterator begin() const { return Iterator{set_bits.begin()}; }
+  Iterator end() const { return Iterator{set_bits.end()}; }
+
+private:
+  llvm::iterator_range<llvm::BitVector::const_set_bits_iterator> set_bits;
+};
+
+//===----------------------------------------------------------------------===//
 /// `TensorExp` identifiers.  These are allocated by `Merger::addExp`,
 /// and serve as unique identifiers for the corresponding `TensorExp` object.
 class ExprId final {
   friend class Merger;
   friend struct TensorExp; // Must be "struct" for [-Wmismatched-tags].
   explicit constexpr ExprId(unsigned value) : value(value) {}
 
 public:

 
 //===----------------------------------------------------------------------===//
 /// `LatPoint` identifiers.  These are allocated by `Merger::addLat`,
 /// and serve as unique identifiers for the corresponding `LatPoint` object.
 class LatPointId final {
   friend class Merger;
   explicit constexpr LatPointId(unsigned value) : value(value) {}
 
 public:
   /// Constructs a new lattice-point identifier with a dedicated
   /// known-invalid value.
   explicit constexpr LatPointId() : value(detail::kInvalidId) {}
   /// Checks whether the lattice-point identifier has the dedicated
   /// known-invalid value.
   constexpr bool isValid() const { return value != detail::kInvalidId; }
   constexpr bool operator==(LatPointId rhs) const { return value == rhs.value; }
   constexpr bool operator!=(LatPointId rhs) const { return value != rhs.value; }
 
 private:
   unsigned value;
 };
 static_assert(std::is_trivially_copyable_v<LatPointId> &&
               std::is_trivially_destructible_v<LatPointId>);
 
 //===----------------------------------------------------------------------===//

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

private:
   /// in the given slice.
   /// Returns std::pair<Transformed coordinates, Predicate>
   std::pair<Value, Value> genSliceLegitPredicate(OpBuilder &builder,
                                                  Location loc, Value crd,
                                                  TensorId tid, Level lvl);
 
   unsigned getNumTensors() const { return tensors.size(); }
 
+  tensor_id::Range getTensorIds() const {
+    return tensor_id::Range(0, getNumTensors());
+  }
+
   bool isOutputTensor(TensorId tid) const {
     return hasOutput && tid == getNumTensors() - 1;
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 

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

void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   this->loopIdToOrd.assign(numLoops, 0);
   this->loopStack.reserve(numLoops);
   this->loopSeqStack.reserve(numLoops);
 
   this->dependentLvlMap.assign(
       numTensors, std::vector<std::vector<std::pair<TensorId, Level>>>());
 
   // Initialize nested types of `TensorId`-indexed fields.
-  for (TensorId tid = 0; tid < numTensors; tid++) {
+  for (const TensorId tid : getTensorIds()) {
     const Value t = tensors[tid];
     // a scalar or 0-dimension tensors
     if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
 
     auto rtp = getRankedTensorType(t);
     if (auto reshape = t.getDefiningOp<tensor::CollapseShapeOp>();
         isUniqueCOOType(rtp) && reshape) {

 void LoopEmitter::initializeLoopEmit(OpBuilder &builder, Location loc,
                                      LoopEmitter::OutputUpdater updater) {
   // For every tensor:
   // * get the values buffer.
   // * For every level:
   //   * get the positions and coordinates buffers
   //   * get/compute the level-size, which is also used as the upper-bound
   //     on positions.
-  for (TensorId t = 0, numTensors = getNumTensors(); t < numTensors; t++) {
+  for (const TensorId t : getTensorIds()) {
     const Value tensor = tensors[t];
     const auto rtp = tensor.getType().dyn_cast<RankedTensorType>();
     if (!rtp)
       // Skips only scalar, zero ranked tensor still need to be bufferized and
       // (probably) filled with zeros by users.
       continue;
     // FIXME: the definition of `lvlRank` looks more like a dim-rank;
     // but the variable is used as a level everywhere below, which

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

 
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// same index is used more than once. Also rejects compound affine
 /// expressions in sparse dimensions.
 /// filterIdx stores the current filter loop idx should be used for the next
 /// compound affine sparse level, and it will be incremented by one when
 /// used.
 static bool findAffine(Merger &merger, TensorId tid, Level lvl, AffineExpr a,
-                       DimLevelType dlt, LoopId &filterLdx,
+                       DimLevelType dlt, loop_id::Iterator &filterLdx,
                        bool setLvlFormat = true) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     const LoopId idx = merger.makeLoopId(a.cast<AffineDimExpr>().getPosition());
     if (!isUndefDLT(merger.getDimLevelType(tid, idx)))
       return false; // used more than once
 
     if (setLvlFormat)
       merger.setLevelAndType(tid, idx, lvl, dlt);
     return true;
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul:
   case AffineExprKind::Constant: {
     if (!isDenseDLT(dlt) && setLvlFormat) {
-      assert(isUndefDLT(merger.getDimLevelType(tid, filterLdx)));
+      assert(isUndefDLT(merger.getDimLevelType(tid, *filterLdx)));
       // Use a filter loop for sparse affine expression.
-      merger.setLevelAndType(tid, filterLdx, lvl, dlt);
+      merger.setLevelAndType(tid, *filterLdx, lvl, dlt);
       ++filterLdx;
     }
 
     if (auto binOp = a.dyn_cast<AffineBinaryOpExpr>()) {
       // We do not set dim level format for affine expresssion like d0 + d1 on
       // either loop index at d0 or d1.
       // We continue the recursion merely to check whether current affine is
       // admissible or not.

 /// expression on sparse tensors, and they accept different affine expressions.
 /// When using filter-loop-based approach, it accept (almost) arbitrary affine
 /// index expression on sparse tensor but it is much less efficient, and will be
 /// gradually removed from the codebase.
 /// When using dependent index reducton-based approach, it currently only
 /// supports affine addition index expression.
 static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased) {
   bool annotated = false;
-  // `filterLdx` may be mutated by `findAffine`.
-  LoopId filterLdx = env.merger().getStartingFilterLoopId();
+  const loop_id::Range filterLdxRange = env.merger().getFilterLoopIds();
+  loop_id::Iterator filterLdx = filterLdxRange.begin();
   for (OpOperand &t : env.op()->getOpOperands()) {
     const TensorId tid = env.makeTensorId(t.getOperandNumber());
     const auto map = env.op().getMatchingIndexingMap(&t);
     const auto enc = getSparseTensorEncoding(t.get().getType());
     if (enc)
       annotated = true;
 
     const Level lvlRank = map.getNumResults();
for (Level l = 0; l < lvlRank; l++) {
         if (!findDepIdxSet(env.merger(), tid, l, a, dlt))
           return false; // inadmissible affine expression
       } else {
         if (!findAffine(env.merger(), tid, l, a, dlt, filterLdx))
           return false; // inadmissible affine expression
       }
     }
   }
-  assert(filterLdx == env.merger().getNumLoops());
+  assert(filterLdx == filterLdxRange.end());
   return annotated;
 }
 
 /// A helper to compute a topological sort. O(n^2) time complexity
 /// as we use adj matrix for the graph.
 /// The sorted result will put the first Reduction iterator to the
 /// latest possible `LoopOrd`.
 ///
 /// The `inDegree` is indexed by `LoopId`, and the `adjM` is indexed by
 /// `(LoopId,LoopId)`.
 static bool topSortOptimal(CodegenEnv &env,
                            ArrayRef<utils::IteratorType> iteratorTypes,
                            std::vector<unsigned> &inDegree,
                            std::vector<std::vector<bool>> &adjM) {
   std::vector<LoopId> redIt;    // reduce iterator with 0 degree
   std::vector<LoopId> parIt;    // parallel iterator with 0 degree
   std::vector<LoopId> filterIt; // filter loop with 0 degree
-  const LoopId numLoops = env.merger().getNumLoops();
-  for (LoopId i = 0; i < numLoops; i++) {
+  for (const LoopId i : env.merger().getLoopIds()) {
     if (inDegree[i] == 0) {
       if (env.merger().isFilterLoop(i))
         filterIt.push_back(i);
       else if (linalg::isReductionIterator(iteratorTypes[i]))
         redIt.push_back(i);
       else
         parIt.push_back(i);
     }
while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
     //    if (xxx)
     //      for (1 to M)
     //        O(X) computation  => O(NK+NMX) time complexity
     auto &it = !filterIt.empty() ? filterIt : (!parIt.empty() ? parIt : redIt);
     auto src = it.back();
     env.topSortPushBack(src);
     it.pop_back();
     // Update in-degree, and push 0-degree node into worklist.
-    for (LoopId dst = 0; dst < numLoops; dst++) {
+    for (const LoopId dst : env.merger().getLoopIds()) {
       if (adjM[src][dst] && --inDegree[dst] == 0) {
         if (env.merger().isFilterLoop(dst))
           filterIt.push_back(dst);
         else if (linalg::isReductionIterator(iteratorTypes[dst]))
           redIt.push_back(dst);
         else
           parIt.push_back(dst);
       }
     }
   }
-  return env.topSortSize() == numLoops;
+  return env.topSortSize() == env.merger().getNumLoops();
 }
 
 /// Helper method to add all constraints from the indices in one affine
 /// expression before all indices in the other affine expression. For
 /// example i0+i1 < i2+i3+1 yields i0<i2, i0<i3, i1<i2, and i1<i3.
 /// The affine expression `a` is empty iff `fidx` have a value, leading to
 /// b = (i0 + i1) < fidx => i0 < fidx, i1 < fidx.
 /// The affine expression `b` is empty iff `tidx` have a value, leading to
if ((isDenseInput && !includesDenseInput(mask)) ||
         (isDenseOutput && !includesDenseOutput(mask)))
       continue;
 
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     // TODO: Do we really need this?
     if (includesUndef(mask)) {
       const TensorId tid = env.makeTensorId(t.getOperandNumber());
-      for (LoopId i = 0; i < numLoops; i++) {
+      for (const LoopId i : env.merger().getLoopIds()) {
         const auto dltI = env.dlt(tid, i);
         if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
-          for (LoopId j = 0; j < numLoops; j++)
+          for (const LoopId j : env.merger().getLoopIds())
             if (isUndefDLT(env.dlt(tid, j))) {
               adjM[i][j] = true;
               inDegree[j]++;
             }
         } else {
           assert(isDenseDLT(dltI) || isUndefDLT(dltI));
         }
       }

 }
 
 static void genInitConstantDenseAddress(CodegenEnv &env,
                                         RewriterBase &rewriter) {
   // We can generate address for constant affine expression before any loops
   // starting from the first level as they do not depend on any thing.
   // E.g., [Dense, Dense, Sparse] -> (1, 2, d0), the addresses for the first two
   // levels can be determined before loops.
-  for (TensorId tid = 0, e = env.op().getNumDpsInputs(); tid < e; tid++)
+  const TensorId hi = env.makeTensorId(env.op().getNumDpsInputs());
+  for (const TensorId tid : tensor_id::Range(0, hi))
     genConstantDenseAddressFromLevel(env, rewriter, tid, 0);
 }
 
 /// Return true if the lattices bit can be iterated by a for loop.
 static bool translateBitsToTidLvlPairs(
     CodegenEnv &env, LatPointId li, LoopId ldx, SmallVectorImpl<TensorId> &tids,
     SmallVectorImpl<Level> &lvls, SmallVectorImpl<TensorId> &affineTids,
     SmallVectorImpl<Level> &affineLvls, SmallVectorImpl<AffineExpr> &exps) {

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
   return simple;
 }
 
 bool Merger::latGT(LatPointId i, LatPointId j) const {
   const BitVector &bitsi = lat(i).bits;
   const BitVector &bitsj = lat(j).bits;
   assert(bitsi.size() == bitsj.size());
   if (bitsi.count() > bitsj.count()) {
-    for (TensorLoopId b = 0, be = bitsj.size(); b < be; b++)
-      if (bitsj[b] && !bitsi[b])
+    for (const TensorLoopId b : tensor_loop_id::Range(bitsj))
+      if (!bitsi[b])
         return false;
     return true;
   }
   return false;
 }
 
 bool Merger::onlyDenseDiff(LatPointId i, LatPointId j) const {
   BitVector tmp(lat(j).bits);
bool Merger::isSingleCondition(TensorId t, ExprId e) const {
   case TensorExp::Kind::kBinary:
   case TensorExp::Kind::kReduce:
     return false;
   }
   llvm_unreachable("unexpected kind");
 }
 
 bool Merger::hasAnySparse(const BitVector &bits) const {
-  for (TensorLoopId b = 0, be = bits.size(); b < be; b++)
-    if (bits[b]) {
-      const auto dlt = getDimLevelType(b);
-      if (isCompressedDLT(dlt) || isSingletonDLT(dlt))
-        return true;
-    }
+  for (const TensorLoopId b : tensor_loop_id::Range(bits)) {
   return false;
 }
 
 bool Merger::hasSparseIdxReduction(const BitVector &bits) const {
   // TODO: return false on dense levels.
   for (unsigned b = 0, be = bits.size(); b < be; b++)
     if (bits[b] && isLvlWithNonTrivialIdxExp(b))
       return true;
void Merger::dumpSet(LatSetId s) const {
   for (const LatPointId p : ss) {
     llvm::dbgs() << "  ";
     dumpLat(p);
   }
   llvm::dbgs() << "}\n";
 }
 
 void Merger::dumpBits(const BitVector &bits) const {
-  for (TensorLoopId b = 0, be = bits.size(); b < be; b++) {
-    if (bits[b]) {
-      const TensorId t = tensor(b);
-      const LoopId i = loop(b);
-      const auto dlt = lvlTypes[t][i];
-      if (isLvlWithNonTrivialIdxExp(b))
-        llvm::dbgs() << " DEP_" << t << "_" << i;
-      else
-        llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
-    }
-  }
-}
+  for (const TensorLoopId b : tensor_loop_id::Range(bits)) {
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.
 //===----------------------------------------------------------------------===//
 
 LatSetId Merger::buildLattices(ExprId e, LoopId i) {