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

[wrengr]

Summary

This commit converts these types from typedefs to newtypes. In addition this allows cleaning up some other code (because we couldn't overload things by type before), such as the TensorExp constructor.

Depends On D146691

[wrengr]

git-clang-format

[Peiming]

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

308

Why using operation * instead of operator unsigned() to convert it implicitly to unsigned for array indexing?

[wrengr]

rebasing for the postponement splitting of D146684

[wrengr]

rebasing for the postponement splitting of D146684

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

 namespace sparse_tensor {
 
 /// 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 {
+    explicit constexpr Children() = default;
+    explicit constexpr Children(ExprId e0, ExprId e1) : e0(e0), e1(e1) {}
     ExprId e0;
     ExprId e1;
   };
 
-  // The `x` parameter has different types depending on the value of the
-  // `k` parameter.  The correspondences are:
-  // * `kTensor`    -> `TensorId`
-  // * `kInvariant` -> `kInvalidId`
-  // * `kLoopVar`   -> `LoopId`
-  // * else         -> `ExprId`
-  //
-  // The `y`, `v`, and `op` parameters either must or must not be
-  // `kInvalidId`/`nullptr`, depending on the value of the `k` parameter;
-  // however, they have uniform C++ types regardless of the value of `k`.
-  TensorExp(Kind k, unsigned x, ExprId y, Value v, Operation *op);
+  explicit constexpr TensorExp(TensorId t);
+  explicit constexpr TensorExp(LoopId i);
+  explicit constexpr TensorExp(Value v);
+
+  /// This constructor handles all unary, sesquinary, and binary
+  /// expressions.  The `e0` argument must always be a valid identifier.
+  /// Whereas, the other three arguments either must or must not be valid,
+  /// depending on the value of `k`.
+  constexpr TensorExp(Kind k, ExprId e0, ExprId e1, Value v, Operation *op);
 
   /// Tensor expression kind.
   Kind kind;
 
   union {
     /// `kTensor` expressions simply have a tensor identifier.
     TensorId tensor;
 
Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
          unsigned numFilterLoops, unsigned maxLvlRank);
 
   //
   // Constructing valid tensor and loop identifiers.
   //
 
   /// Safely converts the argument to a tensor identifier.
   constexpr TensorId makeTensorId(unsigned t) const {
-    assert(isValidTensorId(t));
-    return t;
+    return factory.makeTensorId(t);
   }
 
   /// Safely converts the argument to a loop identifier.
   constexpr LoopId makeLoopId(unsigned i) const {
-    assert(isValidLoopId(i));
-    return i;
+    return factory.makeLoopId(i);
   }
 
   /// Safely converts the arguments to a pair of (tensor,loop) identifiers.
   constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const {
-    assert(isValidTensorId(t) && isValidLoopId(i));
-    return numTensors * i + t;
+    return factory.makeTensorLoopId(t, i);
+  }
+  constexpr TensorLoopId makeTensorLoopId(TensorId t, LoopId i) const {
+    return factory.makeTensorLoopId(t, i);
   }
 
   //
   // Allocating new expressions, points, and sets.
   //
 
   /// Constructs a new tensor expression, and returns its identifier.
   ExprId addTensorExp(TensorId t);
public:
 
   /// Returns true if p0 > p1.
   bool latGT(LatPointId p0, LatPointId p1) const;
 
   /// Returns true if p0 and p1 only differ in dense.
   bool onlyDenseDiff(LatPointId p0, LatPointId p1) const;
 
   /// Gets the tensor-identifier of the `TensorLoopId`.
-  constexpr TensorId tensor(TensorLoopId b) const { return b % numTensors; }
+  constexpr TensorId tensor(TensorLoopId b) const {
+    return factory.getTensorId(b);
+  }
   /// Gets the loop-identifier of the `TensorLoopId`.
-  constexpr LoopId loop(TensorLoopId b) const { return b / numTensors; }
+  constexpr LoopId loop(TensorLoopId b) const { return factory.getLoopId(b); }
 
   /// Get the total number of tensors (including the output-tensor and
   /// synthetic-tensor).
-  constexpr unsigned getNumTensors() const { return numTensors; }
+  constexpr unsigned getNumTensors() const { return factory.getNumTensors(); }
   /// Get the range of all tensor identifiers.
-  constexpr tensor_id::Range getTensorIds() const {
-    return tensor_id::Range(0, numTensors);
-  }
+  constexpr TensorId::Range getTensorIds() const { return factory.tensorIds(); }
 
   /// Get the total number of loops (native loops + filter loops).
-  constexpr unsigned getNumLoops() const { return numLoops; }
+  constexpr unsigned getNumLoops() const { return factory.getNumLoops(); }
   /// 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;
+    return getNumLoops() - numNativeLoops;
   }
   /// Get the range of all loop identifiers.
-  constexpr loop_id::Range getLoopIds() const {
-    return loop_id::Range(0, numLoops);
-  }
+  constexpr LoopId::Range getLoopIds() const { return factory.loopIds(); }
   /// Get the range of native-loop identifiers.
-  constexpr loop_id::Range getNativeLoopIds() const {
-    return loop_id::Range(0, numNativeLoops);
+  constexpr LoopId::Range getNativeLoopIds() const {
+    return factory.loopIdsUpTo(numNativeLoops);
   }
   /// Get the range of filter-loop identifiers.
-  constexpr loop_id::Range getFilterLoopIds() const {
-    return loop_id::Range(numNativeLoops, numLoops);
+  constexpr LoopId::Range getFilterLoopIds() const {
+    return factory.loopIdsFrom(numNativeLoops);
   }
 
   /// 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.
   constexpr TensorId getOutTensorID() const { return outTensor; }
   /// Get the synthetic tensor's identifier (used for all invariant
   /// tensor expressions).
   constexpr TensorId getSynTensorID() const { return syntheticTensor; }
 
   constexpr bool isFilterLoop(LoopId i) const {
     assert(isValidLoopId(i));
-    return i >= numNativeLoops;
+    return i.value >= numNativeLoops;
   }
 
   /// Returns true if the expression is `(kTensor t)`.
   bool expIsTensor(ExprId e, TensorId t) const {
     const auto &expr = exp(e);
     return expr.kind == TensorExp::Kind::kTensor && expr.tensor == t;
   }
 
public:
 
   /// Returns true if bits contains a dependent index reduction condition on
   /// sparse levels.
   bool hasSparseIdxReduction(const BitVector &bits) const;
 
   /// Gets the level-type of the `t`th tensor on `i`th loop.
   DimLevelType getDimLevelType(TensorId t, LoopId i) const {
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return lvlTypes[t][i];
+    return lvlTypes[t.value][i.value];
   }
 
   /// Gets the level-type of the TensorLoopId.
   DimLevelType getDimLevelType(TensorLoopId b) const {
     return getDimLevelType(tensor(b), loop(b));
   }
 
   /// Gets the loop identifier for the `lvl`th level of the `t`th tensor.
   std::optional<LoopId> getLoopId(TensorId t, Level lvl) const {
     assert(isValidLevel(t, lvl));
-    return lvlToLoop[t][lvl];
+    return lvlToLoop[t.value][lvl];
   }
 
   /// Gets the level number of the the `t`th tensor on `i`th loop.
   std::optional<Level> getLvl(TensorId t, LoopId i) const {
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return loopToLvl[t][i];
+    return loopToLvl[t.value][i.value];
   }
   std::optional<Level> getLvl(TensorLoopId b) const {
     return getLvl(tensor(b), loop(b));
   }
 
   /// Sets the level number and level-type of the `t`th tensor on
   /// `i`th loop.
   void setLevelAndType(TensorId t, LoopId i, Level lvl, DimLevelType dlt) {
     assert(isValidLevel(t, lvl) && isValidLoopId(i) && isValidDLT(dlt));
-    lvlTypes[t][i] = dlt;
-    loopToLvl[t][i] = lvl;
-    lvlToLoop[t][lvl] = i;
+    lvlTypes[t.value][i.value] = dlt;
+    loopToLvl[t.value][i.value] = lvl;
+    lvlToLoop[t.value][lvl] = i;
   }
 
   using ForeachTensorLoopIdCallback = function_ref<void(
       TensorLoopId, TensorId, std::optional<Level>, DimLevelType, bool)>;
 
   /// Iterates over a set of `TensorLoopId`s, invoking the callback
   /// for each `TensorLoopId` and passing it the corresponding tensor
   /// identifier, level, and level-type, following with a boolean value
   /// indicating whether it is a dependent index reduction loop condition.
   void foreachTensorLoopId(LatPointId p,
                            ForeachTensorLoopIdCallback callback) const {
     // TODO: the default ought to be simple=true; but we'll need to make
     // sure to update all the tests to make sure they do the right thing.
     foreachTensorLoopId(p, /*simple=*/false, callback);
   }
   void foreachTensorLoopId(LatPointId p, bool simple,
                            ForeachTensorLoopIdCallback callback) const {
     const auto &point = lat(p);
     const auto &bits = simple ? point.simple : point.bits;
-    for (const TensorLoopId b : bits.set_bits()) {
+    for (const TensorLoopId b : TensorLoopId::Range(bits)) {
       const TensorId t = tensor(b);
       const auto optLvl = getLvl(b);
       const auto lvlTp = getDimLevelType(b);
       if (isLvlWithNonTrivialIdxExp(b)) {
         // This must be an undefined level.
         assert(!optLvl.has_value());
         // Slice the tid along the dependent level to iterate current loop.
-        callback(b, t, loopToDependencies[loop(b)][t], lvlTp,
+        callback(b, t, loopToDependencies[loop(b).value][t.value], lvlTp,
                  /*isIdxReduc=*/true);
       } else {
         callback(b, t, optLvl, lvlTp, /*isIdxReduc=*/false);
       }
     }
   }
 
   /// Sets whether the output tensor is sparse or not.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
   /// Establishes the two-way map that i <-> <t, lvl>.
   void setLoopDependentTensorLevel(LoopId i, TensorId t, Level lvl) {
     assert(isValidLoopId(i) && isValidLevel(t, lvl));
-    loopToDependencies[i][t] = lvl;
-    levelToDependentIdx[t][lvl].push_back(i);
+    loopToDependencies[i.value][t.value] = lvl;
+    levelToDependentIdx[t.value][lvl].push_back(i);
   }
 
   /// Whether the loop has dependent slice.
   bool hasDependentLvl(LoopId i, TensorId t) {
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return loopToDependencies[i][t].has_value();
+    return loopToDependencies[i.value][t.value].has_value();
   }
 
   /// Returns the list of loop indices which appear in the non-trivial index
   /// expression on t_l, e.g., A[i+j] => {i, j}
   std::vector<LoopId> &getDependentLoops(TensorId t, Level lvl) {
     assert(isValidLevel(t, lvl));
-    return levelToDependentIdx[t][lvl];
+    return levelToDependentIdx[t.value][lvl];
   }
 
   /// Returns the defining [tid, lvl] for the loop.
   std::pair<TensorId, Level> getLoopDefiningLvl(LoopId i) const {
     assert(isValidLoopId(i));
-    return loopBounds[i];
+    return loopBounds[i.value];
   }
 
   /// Checks whether the TensorLoopId represents a tensor level with
   /// non-trivial index expression on it.
   bool isLvlWithNonTrivialIdxExp(TensorLoopId b) const {
     const TensorId t = tensor(b);
     const LoopId i = loop(b);
     assert(isValidTensorId(t) && isValidLoopId(i));
-    return loopToDependencies[i][t].has_value();
+    return loopToDependencies[i.value][t.value].has_value();
   }
 
   /// Convenience getters to immediately access the stored nodes.
   /// These methods return `const&` because the underlying objects must
   /// not be mutated by client code.  The only exception is for mutating
   /// the value associated with an expression, for which there are
   /// dedicated methods below.
   ///
#endif
   std::optional<ExprId> buildTensorExpFromLinalg(linalg::GenericOp op);
 
   /// Rebuilds SSA format from a tensor expression.
   Value buildExp(RewriterBase &rewriter, Location loc, ExprId e, Value v0,
                  Value v1) const;
 
 private:
   /// Private helpers.
-  constexpr bool isValidTensorId(TensorId t) const { return t < numTensors; }
+  constexpr bool isValidTensorId(TensorId t) const {
+    return factory.isValidTensorId(t);
+  }
   constexpr bool isValidLoopId(LoopId i) const {
-    return i != detail::kInvalidId && i < numLoops;
+    return factory.isValidLoopId(i);
   }
   bool isValidLevel(TensorId t, Level lvl) const {
-    return isValidTensorId(t) && lvl < lvlToLoop[t].size();
+    return isValidTensorId(t) && lvl < lvlToLoop[t.value].size();
   }
   bool isValidExprId(ExprId e) const {
     return e.isValid() && e.value < tensorExps.size();
   }
   bool isValidLatPointId(LatPointId p) const {
     return p.isValid() && p.value < latPoints.size();
   }
   bool isValidLatSetId(LatSetId s) const {
     return s.isValid() && s.value < latSets.size();
   }
   bool maybeZero(ExprId e) const;
   bool isInvariant(ExprId e) const {
     return exp(e).kind == TensorExp::Kind::kInvariant;
   }
   Type inferType(ExprId e, Value src) const;
 
   /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
   std::optional<ExprId> buildTensorExp(linalg::GenericOp op, Value v);
 
   /// Merger data structures.
   const TensorId outTensor;
   const TensorId syntheticTensor;
-  const unsigned numTensors;
+  const TensorLoopFactory factory;
   const unsigned numNativeLoops;
-  const unsigned numLoops;
   bool hasSparseOut;
 
   // Below we use `std::vector` for things which have a priori fixed
   // sizes, whereas we use `llvm::SmallVector` for things with variable
   // size.  Do beware that these two classes differ in the semantics of
   // `operator[]`: `SmallVector` performs OOB checks, whereas `std::vector`
   // does not.
 

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

 
 namespace detail {
 /// A constant serving as the canonically invalid identifier,
 /// regardless of the identifier type.
 static constexpr unsigned kInvalidId = -1u;
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
-/// Tensor identifiers.
+/// Tensor identifiers.  The `TensorLoopFactory` class can be used to
+/// construct identifiers from `unsigned` values while checking validity.
+/// (For more details, see that class.)  Since client code often needs
+/// to construct tensor identifiers from arbitrary values, we provide
+/// the `TensorId::Unchecked` factory-method for constructing arbitrary
+/// identifiers.  Unlike the other newtypes, we don't have a dedicated
+/// known-invalid tensor identifier (because it isn't needed so far).
 ///
 /// 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 TensorId final {
+public:
+  // Must forward-declare, because they can only be defined once the
+  // `TensorId` definition is complete.
   class Iterator;
   class Range;
-} // namespace tensor_id
+
+private:
+  friend class TensorId::Iterator;
+  friend class TensorLoopFactory;
+  friend class Merger;
+  explicit constexpr TensorId(unsigned value) : value(value) {}
+
+public:
+  /// Constructs a new tensor identifier without performing any validity
+  /// checks.  We provide this as a factory rather than as a ctor so that
+  /// the potential invalidity is explicit at all callsites.
+  static constexpr TensorId Unchecked(unsigned value) {
+    return TensorId{value};
+  }
+
+  // Since the set of valid identifiers and their numeric values are both
+  // fixed at the time the `Merger` (or `TensorLoopFactory`) is constructed,
+  // it's perfectly meaningful to allow client code to extract those values.
+  constexpr unsigned operator*() const { return value; }
+
+  constexpr bool operator==(TensorId rhs) const { return value == rhs.value; }
+  constexpr bool operator!=(TensorId rhs) const { return value != rhs.value; }
+
+  /// Returns the tensor identifiers `[0, hi)` as an iterable range.
+  static constexpr Range fromZeroUpTo(TensorId hi);
+  /// Returns the tensor identifiers `[this, hi)` as an iterable range.
+  constexpr Range upTo(TensorId hi) const;
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<TensorId> &&
+              std::is_trivially_destructible_v<TensorId>);
 
 /// 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
+/// iterators via `TensorId::Range` instead.
+class TensorId::Iterator final
     : public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
                                         TensorId> {
-  friend class tensor_id::Range;
+  friend class TensorId::Range;
   explicit constexpr Iterator(TensorId tid) : tid(tid) {}
 
 public:
   using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
                                    TensorId>::operator++;
   Iterator &operator++() {
-    ++tid;
+    ++tid.value;
     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>);
+static_assert(std::is_trivially_copyable_v<TensorId::Iterator> &&
+              std::is_trivially_destructible_v<TensorId::Iterator>);
 
 /// An iterator range for `TensorId`.
-class tensor_id::Range final {
+class TensorId::Range final {
 public:
   explicit constexpr Range(TensorId lo, TensorId hi)
-      : begin_(lo <= hi ? lo : hi), end_(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>);
+static_assert(std::is_trivially_copyable_v<TensorId::Range> &&
+              std::is_trivially_destructible_v<TensorId::Range>);
+
+// These two can only be defined once the definitions of both `TensorId`
+// and `TensorId::Range` are complete.
+inline constexpr TensorId::Range TensorId::fromZeroUpTo(TensorId hi) {
+  return TensorId{0}.upTo(hi);
+}
+inline constexpr TensorId::Range TensorId::upTo(TensorId hi) const {
+  return TensorId::Range(*this, hi);
+}
 
 //===----------------------------------------------------------------------===//
-/// Loop identifiers.
+/// Loop identifiers.  The `TensorLoopFactory` class can be used to
+/// construct identifiers from `unsigned` values while checking validity.
+/// (For more details, see that class.)  Since client code often needs
+/// to construct loop identifiers from arbitrary values, we provide
+/// the `LoopId::Unchecked` factory-method for constructing arbitrary
+/// identifiers.  In addition, the `LoopId::Invalid` factory-method can
+/// be used to construct a dedicated known-invalid loop identifier.
 ///
 /// 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 LoopId final {
+public:
+  // Must forward-declare, because they can only be defined once the
+  // `LoopId` definition is complete.
   class Iterator;
   class Range;
-} // namespace loop_id
+
+private:
+  friend class LoopId::Iterator;
+  friend class TensorLoopFactory;
+  friend class Merger;
+  explicit constexpr LoopId(unsigned value) : value(value) {}
+
+public:
+  /// Constructs a new loop identifier without performing any validity
+  /// checks.  We provide this as a factory rather than as a ctor so that
+  /// the potential invalidity is explicit at all callsites.
+  static constexpr LoopId Unchecked(unsigned value) { return LoopId{value}; }
+
+  /// Constructs a new loop identifier with a dedicated known-invalid value.
+  static constexpr LoopId Invalid() { return LoopId{detail::kInvalidId}; }
+
+  /// Checks whether the loop identifier has the dedicated known-invalid value.
+  constexpr bool isValid() const { return value != detail::kInvalidId; }
+
+  // Since the set of valid identifiers and their numeric values are both
+  // fixed at the time the `Merger` (or `TensorLoopFactory`) is constructed,
+  // it's perfectly meaningful to allow client code to extract those values.
+  constexpr unsigned operator*() const { return value; }
+
+  constexpr bool operator==(LoopId rhs) const { return value == rhs.value; }
+  constexpr bool operator!=(LoopId rhs) const { return value != rhs.value; }
+
+  /// Returns the loop identifiers `[0, hi)` as an iterable range.
+  static constexpr Range fromZeroUpTo(LoopId hi);
+  /// Returns the loop identifiers `[this, hi)` as an iterable range.
+  constexpr Range upTo(LoopId hi) const;
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<LoopId> &&
+              std::is_trivially_destructible_v<LoopId>);
 
 /// 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
+/// iterators via `LoopId::Range` instead.
+class LoopId::Iterator final
     : public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
                                         LoopId> {
-  friend class loop_id::Range;
+  friend class LoopId::Range;
   explicit constexpr Iterator(LoopId i) : loop(i) {}
 
 public:
   using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
                                    LoopId>::operator++;
   Iterator &operator++() {
-    ++loop;
+    ++loop.value;
     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>);
+static_assert(std::is_trivially_copyable_v<LoopId::Iterator> &&
+              std::is_trivially_destructible_v<LoopId::Iterator>);
 
 /// An iterator range for `LoopId`.
-class loop_id::Range final {
+class LoopId::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) {}
+      : begin_((lo.isValid() && hi.isValid()) ? (*lo <= *hi ? lo : hi)
+                                              : LoopId::Invalid()),
+        end_(lo.isValid() ? hi : LoopId::Invalid()) {}
   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>);
+static_assert(std::is_trivially_copyable_v<LoopId::Range> &&
+              std::is_trivially_destructible_v<LoopId::Range>);
+
+// These two can only be defined once the definitions of both `LoopId`
+// and `LoopId::Range` are complete.
+inline constexpr LoopId::Range LoopId::fromZeroUpTo(LoopId hi) {
+  return LoopId{0}.upTo(hi);
+}
+inline constexpr LoopId::Range LoopId::upTo(LoopId hi) const {
+  return LoopId::Range(*this, hi);
+}
 
 //===----------------------------------------------------------------------===//
 /// 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 TensorLoopId final {
+public:
+  // Must forward-declare, because they can only be defined once the
+  // `TensorId` definition is complete.
   class Iterator;
   class Range;
-} // namespace tensor_loop_id
+
+private:
+  friend class TensorLoopId::Iterator;
+  friend class TensorLoopFactory;
+  friend class Merger;
+  explicit constexpr TensorLoopId(unsigned value) : value(value) {}
+
+public:
+  // We provide this getter because some code in Sparsification.cpp
+  // needs it for indexing into `BitVector`s.  Unfortunately this leaks
+  // implementation details about the compression scheme; but so long as
+  // client code only interprets it as an arbitrary index into `BitVector`
+  // rather than any meaningful number, then that's still safe.
+  constexpr unsigned operator*() const { return value; }
+
+  constexpr bool operator==(TensorLoopId b) const { return value == b.value; }
+  constexpr bool operator!=(TensorLoopId b) const { return value != b.value; }
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<TensorLoopId> &&
+              std::is_trivially_destructible_v<TensorLoopId>);
 
 /// An iterator of the `TensorLoopId`s which are included/set in a `BitVector`.
-class tensor_loop_id::Iterator final
+class TensorLoopId::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;
+  friend class TensorLoopId::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 {
+class TensorLoopId::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;
 };
 
 //===----------------------------------------------------------------------===//
+/// A factory class for constructing `TensorId`, `LoopId`, and `TensorLoopId`
+/// while ensuring their validity.  Unfortunately, since we can't "brand" the
+/// resulting identifiers as having come from a particular factory instance,
+/// this construction-time validity checking does not ensure use-time validity.
+/// (The only way to fix that would be to template the identifier classes
+/// on `numTensors`/`numLoops`, which would introduce too much complexity.)
+class TensorLoopFactory {
+public:
+  constexpr TensorLoopFactory(unsigned numTensors, unsigned numLoops)
+      : numTensors(numTensors), numLoops(numLoops) {}
+
+  /// Get the number of tensors for this factory.
+  constexpr unsigned getNumTensors() const { return numTensors; }
+
+  /// Get the number of loops for this factory.
+  constexpr unsigned getNumLoops() const { return numLoops; }
+
+  /// Checks whether the tensor identifier is valid for this factory.
+  constexpr bool isValidTensorId(TensorId t) const {
+    return t.value < numTensors;
+  }
+
+  /// Checks whether the loop identifier is valid for this factory.
+  constexpr bool isValidLoopId(LoopId i) const {
+    return i.isValid() && i.value < numLoops;
+  }
+
+  /// Projects the tensor-identifier out of the `TensorLoopId`.
+  /// NOTE: This method doesn't/can't check that the `TensorLoopId`
+  /// was constructed for the same `numTensors` as the factory.
+  constexpr TensorId getTensorId(TensorLoopId b) const {
+    return TensorId{b.value % numTensors};
+  }
+
+  /// Projects the loop-identifier out of the `TensorLoopId`.
+  /// NOTE: This method doesn't/can't check that the `TensorLoopId`
+  /// was constructed for the same `numTensors` as the factory.
+  constexpr LoopId getLoopId(TensorLoopId b) const {
+    return LoopId{b.value / numTensors};
+  }
+
+  /// Safely converts the argument to a tensor identifier.
+  constexpr TensorId makeTensorId(unsigned t) const {
+    const TensorId tid{t};
+    assert(isValidTensorId(tid));
+    return tid;
+  }
+
+  /// Returns the iterable range `[0, numTensors)`.
+  constexpr TensorId::Range tensorIds() const {
+    return TensorId::fromZeroUpTo(TensorId{numTensors});
+  }
+  /// Returns the iterable range `[lo, numTensors)`.
+  constexpr TensorId::Range tensorIdsFrom(TensorId lo) const {
+    // The `TensorId::Range` ctor will handle the `lo > numTensors` case.
+    return TensorId::Range(lo, TensorId{numTensors});
+  }
+  constexpr TensorId::Range tensorIdsFrom(unsigned lo) const {
+    return tensorIdsFrom(TensorId{lo});
+  }
+  /// Returns the iterable range `[0, min(hi,numTensors))`.
+  constexpr TensorId::Range tensorIdsUpTo(TensorId hi) const {
+    return TensorId::fromZeroUpTo(
+        hi.value <= numTensors ? hi : TensorId{numTensors});
+  }
+  constexpr TensorId::Range tensorIdsUpTo(unsigned hi) const {
+    return tensorIdsUpTo(TensorId{hi});
+  }
+
+  /// Safely converts the argument to a loop identifier.
+  constexpr LoopId makeLoopId(unsigned i) const {
+    const LoopId lid{i};
+    assert(isValidLoopId(lid));
+    return lid;
+  }
+
+  /// Returns the iterable range `[0, numLoops)`.
+  constexpr LoopId::Range loopIds() const {
+    return LoopId::fromZeroUpTo(LoopId{numLoops});
+  }
+  /// Returns the iterable range `[lo, numLoops)`.
+  constexpr LoopId::Range loopIdsFrom(LoopId lo) const {
+    // The `LoopId::Range` ctor will handle the `lo > numLoops` case.
+    return LoopId::Range(lo, LoopId{numLoops});
+  }
+  constexpr LoopId::Range loopIdsFrom(unsigned lo) const {
+    return loopIdsFrom(LoopId{lo});
+  }
+  /// Returns the iterable range `[0, min(hi,numLoops))`.
+  constexpr LoopId::Range loopIdsUpTo(LoopId hi) const {
+    return LoopId::fromZeroUpTo(hi.value <= numLoops ? hi : LoopId{numLoops});
+  }
+  constexpr LoopId::Range loopIdsUpTo(unsigned hi) const {
+    return loopIdsUpTo(LoopId{hi});
+  }
+
+  /// Safely converts the arguments to a pair of (tensor,loop) identifiers.
+  constexpr TensorLoopId makeTensorLoopId(TensorId t, LoopId i) const {
+    assert(isValidTensorId(t) && isValidLoopId(i));
+    return TensorLoopId{numTensors * i.value + t.value};
+  }
+  constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const {
+    return makeTensorLoopId(TensorId{t}, LoopId{i});
+  }
+
+private:
+  unsigned numTensors;
+  unsigned numLoops;
+};
+
+//===----------------------------------------------------------------------===//
 /// `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 class TensorExp;
   explicit constexpr ExprId(unsigned value) : value(value) {}
 
 public:

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

bool CodegenEnv::isAdmissibleTopoOrder() {
   // into the computation and insertions occur in lexicographic index order.
   LoopOrd nest = 0;
   const auto iteratorTypes = linalgOp.getIteratorTypesArray();
   assert(topSortSize() == latticeMerger.getNumLoops());
   for (const LoopId i : topSort) {
     if (!latticeMerger.isFilterLoop(i)) {
       // We only count non-filter loops as filter loops should be considered
       // a special type of parallel loops.
-      if (linalg::isReductionIterator(iteratorTypes[i]))
+      if (linalg::isReductionIterator(iteratorTypes[*i]))
         break; // terminate at first reduction
       nest++;
     }
   }
   // Determine admissible dynamic insertion situations:
   // (1) fully injective, since there are no reductions,
   // (2) admissible 1-d expansion in innermost dimension.
   if (static_cast<int64_t>(nest) >= linalgOp.getRank(lhs) - 1) {

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());
+  TensorId::Range getTensorIds() const {
+    return TensorId::fromZeroUpTo(TensorId::Unchecked(getNumTensors()));
   }
 
   bool isOutputTensor(TensorId tid) const {
-    return hasOutput && tid == getNumTensors() - 1;
+    return hasOutput && *tid == getNumTensors() - 1;
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
   bool isValidLevel(TensorId tid, Level lvl) const {
-    return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
+    return *tid < lvlTypes.size() && lvl < lvlTypes[*tid].size();
   }
 
   /// Prepares loop for iterating over `tensor[lvl]`, under the assumption
   /// that `tensor[0...lvl-1]` loops have already been set up.
   void prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl);
 
   /// Emits extra locals, since the locals might not be in simplified lattices
private:
   ///
   /// NOTE: the result uses `Level` rather than the `int64_t` of
   /// `ReassociationIndices`, since the former gives clarity to what
   /// the values actually mean.
   ///
   /// TODO: why not do this computation when we first store the reassoc,
   /// instead of doing it every time we look it up?
   SmallVector<Level, 2> getCollapseReassociation(TensorId tid, Level dstLvl) {
-    assert(tid < getNumTensors() && "Invalid TensorId");
+    assert(*tid < getNumTensors() && "Invalid TensorId");
     assert(collapseReassoc.size() == getNumTensors());
-    if (const auto reassoc = collapseReassoc[tid]) {
+    if (const auto reassoc = collapseReassoc[*tid]) {
       // TODO: store the dstLvlRank in the LoopEmitter so that we can
       // check `dstLvl < dstLvlRank` at the top; and only here need to
       // assert that `reassoc.size() == dstLvlRank`.
       assert(dstLvl < reassoc.size() && "Level is out-of-bounds");
       const auto srcLvls = reassoc[dstLvl].cast<ArrayAttr>();
       return llvm::to_vector<2>(
           llvm::map_range(srcLvls, [&](Attribute srcLvl) -> Level {
             // TODO: replace this with the converter for `LevelAttr`.

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

static std::pair<Value, Value> fromSliceCrd(OpBuilder &builder, Location loc,
   Value rem = builder.create<arith::RemUIOp>(loc, crd, stride);
   crd = builder.create<arith::DivUIOp>(loc, crd, stride);
   return std::make_pair(crd, rem);
 }
 
 std::pair<Value, Value>
 LoopEmitter::genSliceLegitPredicate(OpBuilder &builder, Location loc, Value crd,
                                     TensorId tid, Level lvl) {
-  assert(isSparseSlices[tid]);
-  Value slice = tensors[tid];
-  Value offset = sliceOffsets[tid][lvl];
-  Value stride = sliceStrides[tid][lvl];
+  assert(isSparseSlices[*tid]);
+  Value slice = tensors[*tid];
+  Value offset = sliceOffsets[*tid][lvl];
+  Value stride = sliceStrides[*tid][lvl];
   auto enc = getSparseTensorEncoding(slice.getType());
 
   const auto [newCrd, crdRem] =
       fromSliceCrd(builder, loc, crd, offset, stride, slice, lvl);
 
   SmallVector<Value, 3> conds; // at most 3 conditions
 
   // First, coord >= offset (skip the check if offset is known to be 0).
   if (auto staticOffset = enc.getStaticLvlSliceOffset(lvl);
       !(staticOffset.has_value() && *staticOffset == 0)) {
     auto geOffset = builder.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::uge, crd, offset);
     conds.push_back(geOffset);
   }
 
   // Second, coord_in_slice < length
   auto ltLength = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
-                                                newCrd, lvlSizes[tid][lvl]);
+                                                newCrd, lvlSizes[*tid][lvl]);
   conds.push_back(ltLength);
 
   // Third, rem == 0 (skip the check if stride is known to be 1).
   if (auto staticStride = enc.getStaticLvlSliceStride(lvl);
       !(staticStride.has_value() && *staticStride == 1)) {
     auto fitStride = builder.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::eq, crdRem, constantIndex(builder, loc, 0));
     conds.push_back(fitStride);
   }
 
   // Must meet all condition to be a valid coordinate in slice.
   auto pred = conds.front();
   for (auto cond : ValueRange(conds).drop_front())
     pred = builder.create<arith::AndIOp>(loc, pred, cond);
 
   return {newCrd, pred};
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 Value LoopEmitter::genAddress(OpBuilder &builder, Location loc, TensorId tid,
                               Level lvl, Value crd) {
-  Value pos = lvl == 0 ? constantIndex(builder, loc, 0) : posits[tid][lvl - 1];
-  Value mul = builder.create<arith::MulIOp>(loc, highs[tid][lvl], pos);
-  if (isSparseSlices[tid])
-    crd = toSliceCrd(builder, loc, crd, sliceOffsets[tid][lvl],
-                     sliceStrides[tid][lvl], tensors[tid], lvl);
+  Value pos = lvl == 0 ? constantIndex(builder, loc, 0) : posits[*tid][lvl - 1];
+  Value mul = builder.create<arith::MulIOp>(loc, highs[*tid][lvl], pos);
+  if (isSparseSlices[*tid])
+    crd = toSliceCrd(builder, loc, crd, sliceOffsets[*tid][lvl],
+                     sliceStrides[*tid][lvl], tensors[*tid], lvl);
   Value add = builder.create<arith::AddIOp>(loc, mul, crd);
   return add;
 }
 
 Value LoopEmitter::genSegmentHigh(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl, Value pLo,
                                   Value pHi) {
-  const auto coordinates = coordinatesBuffers[tid][lvl];
+  const auto coordinates = coordinatesBuffers[*tid][lvl];
   const auto sameCrd = genIndexLoad(builder, loc, coordinates, pLo);
   auto whileOp = builder.create<scf::WhileOp>(
       loc, builder.getIndexType(), pLo,
       /*beforeBuilder=*/
       [pHi, coordinates, sameCrd](OpBuilder &builder, Location loc,
                                   ValueRange ivs) {
         const auto pos = ivs[0];
         Value inBound = builder.create<arith::CmpIOp>(

 Value LoopEmitter::genSparseCrd(OpBuilder &builder, Location loc, TensorId tid,
                                 Level dstLvl) {
   Value crd = constantIndex(builder, loc, 0);
   const auto reassoc = getCollapseReassociation(tid, dstLvl);
   const unsigned reassocSize = reassoc.size();
   for (unsigned i = 0; i < reassocSize; i++) {
     const Level srcLvl = reassoc[i];
     // A load on the coordinates array yields the coordinate.
-    const Value mem = coordinatesBuffers[tid][srcLvl];
+    const Value mem = coordinatesBuffers[*tid][srcLvl];
     /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
-    const Value pos = posits[tid][dstLvl];
+    const Value pos = posits[*tid][dstLvl];
     const Value off = genIndexLoad(builder, loc, mem, pos);
     // Linearized the coordinates within the same collapse reassociation.
     crd = builder.create<arith::AddIOp>(loc, crd, off);
     if (i != reassocSize - 1) {
       crd = builder.create<arith::MulIOp>(loc, crd,
-                                          this->lvlSizes[tid][reassoc[i + 1]]);
+                                          this->lvlSizes[*tid][reassoc[i + 1]]);
     }
   }
   return crd;
 }
 
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, ArrayRef<LoopId> topSort,
                          DependentLvlGetter dimGetter) {
void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   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 (const TensorId tid : getTensorIds()) {
-    const Value t = tensors[tid];
+    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) {
       // TODO: Supports more kinds of sparse tensors.
       // FIXME: We should instead lower reshape operations on sparse tensors to
       // view change.
-      collapseReassoc[tid] = reshape.getReassociation();
+      collapseReassoc[*tid] = reshape.getReassociation();
       rtp = reshape.getSrcType();
       // Overwrites the tensor to the source tensor of reshape operations.
-      tensors[tid] = reshape.getSrc();
+      tensors[*tid] = reshape.getSrc();
     }
     const SparseTensorType stt(rtp);
     const Level lvlRank = stt.getLvlRank();
     // We always treat sparse output tensor as dense so that we always iterate
     // it based on lvl size.
     if (stt.hasEncoding() && !(isOutputTensor(tid) && isSparseOut)) {
       const auto enc = stt.getEncoding();
-      isSparseSlices[tid] = enc.isSlice();
+      isSparseSlices[*tid] = enc.isSlice();
       for (auto lvlTp : enc.getDimLevelType())
-        lvlTypes[tid].push_back(lvlTp);
+        lvlTypes[*tid].push_back(lvlTp);
     } else {
-      lvlTypes[tid].assign(lvlRank, DimLevelType::Dense);
+      lvlTypes[*tid].assign(lvlRank, DimLevelType::Dense);
     }
 
     // Initialize using empty value.
-    lvlSizes[tid].assign(lvlRank, Value());
-    highs[tid].assign(lvlRank, Value());
-    segHi[tid].assign(lvlRank, Value());
-    posits[tid].assign(lvlRank, Value());
-    coords[tid].assign(lvlRank, Value());
-    positionsBuffers[tid].assign(lvlRank, Value());
-    coordinatesBuffers[tid].assign(lvlRank, Value());
-    sliceOffsets[tid].assign(lvlRank, Value());
-    sliceStrides[tid].assign(lvlRank, Value());
-    dependentLvlMap[tid].assign(lvlRank,
+    lvlSizes[*tid].assign(lvlRank, Value());
+    highs[*tid].assign(lvlRank, Value());
+    segHi[*tid].assign(lvlRank, Value());
+    posits[*tid].assign(lvlRank, Value());
+    coords[*tid].assign(lvlRank, Value());
+    positionsBuffers[*tid].assign(lvlRank, Value());
+    coordinatesBuffers[*tid].assign(lvlRank, Value());
+    sliceOffsets[*tid].assign(lvlRank, Value());
+    sliceStrides[*tid].assign(lvlRank, Value());
+    dependentLvlMap[*tid].assign(lvlRank,
                                 std::vector<std::pair<TensorId, Level>>());
     if (dimGetter) {
-      auto reassoc = collapseReassoc[tid];
+      auto reassoc = collapseReassoc[*tid];
       Level dstRank = reassoc ? reassoc.size() : lvlRank;
       for (Level l = 0; l < dstRank; l++) {
-        dependentLvlMap[tid][l] = dimGetter(tid, l);
+        dependentLvlMap[*tid][l] = dimGetter(tid, l);
         // TODO: View-base collapse and dependent index reduction are not
         // compatible right now.
-        assert(!reassoc || dependentLvlMap[tid][l].empty());
+        assert(!reassoc || dependentLvlMap[*tid][l].empty());
       }
     }
   }
 
   // Construct the inverse of the `topSort` from the sparsifier.
   // This is needed to map `AffineDimExpr`s back to the `LoopOrd`
   // used in loop emitter.
   // FIXME: This map should be maintained outside loop emitter.
   for (LoopOrd n = 0; n < numLoops; n++)
-    loopIdToOrd[topSort[n]] = n;
+    loopIdToOrd[*(topSort[n])] = n;
 }
 
 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 (const TensorId t : getTensorIds()) {
-    const Value tensor = tensors[t];
+    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
     // suggests there may be some dim/lvl confusion going on here.
     const Level lvlRank = rtp.getRank();
     const auto shape = rtp.getShape();
     const auto enc = getSparseTensorEncoding(rtp);
     const Level cooStart = enc ? getCOOStart(enc) : lvlRank;
     // Scan all levels of current tensor.
     for (Level l = 0; l < lvlRank; l++) {
       // This should be called only once at beginning.
-      assert(!positionsBuffers[t][l] && !coordinatesBuffers[t][l] &&
-             !highs[t][l]);
-      const auto lvlTp = lvlTypes[t][l];
+      assert(!positionsBuffers[*t][l] && !coordinatesBuffers[*t][l] &&
+             !highs[*t][l]);
+      const auto lvlTp = lvlTypes[*t][l];
       // Handle sparse storage schemes.
       if (isCompressedDLT(lvlTp)) {
         // Generate sparse primitives to obtain positions and coordinates.
-        positionsBuffers[t][l] = genToPositions(builder, loc, tensor, l);
-        coordinatesBuffers[t][l] =
+        positionsBuffers[*t][l] = genToPositions(builder, loc, tensor, l);
+        coordinatesBuffers[*t][l] =
             genToCoordinates(builder, loc, tensor, l, cooStart);
       } else if (isSingletonDLT(lvlTp)) {
         // Singleton level, fetch coordinates.
-        coordinatesBuffers[t][l] =
+        coordinatesBuffers[*t][l] =
             genToCoordinates(builder, loc, tensor, l, cooStart);
       } else {
         // Dense level, nothing to fetch.
         assert(isDenseDLT(lvlTp));
       }
 
       // FIXME: `toOrigDim` is deprecated.  For now this relies on the
       // 1:1 mapping between levels and dimensions, since nowhere else
       // in the code supports HigherOrdering yet either.
       Value lvlSz = mlir::linalg::createOrFoldDimOp(builder, loc, tensor,
                                                     toOrigDim(enc, l));
       // Find upper bound in current dimension.
-      highs[t][l] = lvlSizes[t][l] = lvlSz;
-      if (isSparseSlices[t]) {
-        sliceOffsets[t][l] = genSliceOffset(builder, loc, tensors[t], l);
-        sliceStrides[t][l] = genSliceStride(builder, loc, tensors[t], l);
+      highs[*t][l] = lvlSizes[*t][l] = lvlSz;
+      if (isSparseSlices[*t]) {
+        sliceOffsets[*t][l] = genSliceOffset(builder, loc, tensors[*t], l);
+        sliceStrides[*t][l] = genSliceStride(builder, loc, tensors[*t], l);
       }
     }
 
     // Perform the required bufferization. Dense inputs materialize
     // from the input tensors. Sparse inputs use sparse primitives to obtain the
     // values.
     // Delegates extra output initialization to clients.
     bool isOutput = isOutputTensor(t);
if (!enc) {
         denseTp = bufferization::getMemRefTypeWithFullyDynamicLayout(rtp);
 
       Value denseVal =
           builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
       // Dense outputs need special handling.
       if (isOutput && updater)
         denseVal = updater(builder, loc, denseVal, tensor);
 
-      valBuffer[t] = denseVal;
+      valBuffer[*t] = denseVal;
     } else {
       // Annotated sparse tensors.
       // We also need the value buffer for all-dense annotated "sparse" tensors.
-      valBuffer[t] = genToValues(builder, loc, tensor);
+      valBuffer[*t] = genToValues(builder, loc, tensor);
     }
     // NOTE: we can also prepare for 0 lvl here in advance, this will hoist
     // some loop preparation from tensor iteration, but will also (undesirably)
     // hoist the code ouside if-conditions.
   }
 }
 
 void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
Operation *LoopEmitter::enterLoopOverTensorAtLvl(
   bool isSparseInput = false;
   TensorId tid = tids.front();
   Level dstLvl = lvls.front();
   assert(tids.size() == lvls.size());
   for (auto [t, l] : llvm::zip(tids, lvls)) {
     // TODO: this check for validity of the (t,l) pairs should be
     // checked/enforced at the callsites, if possible.
     assert(isValidLevel(t, l));
-    assert(!coords[t][l]); // We cannot re-enter the same level
-    const auto lvlTp = lvlTypes[t][l];
+    assert(!coords[*t][l]); // We cannot re-enter the same level
+    const auto lvlTp = lvlTypes[*t][l];
     const bool isSparse = isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp);
     // Must be a recognizable level-type.
     assert(isSparse || isDenseDLT(lvlTp));
     // We can at most have one sparse input, otherwise, a while loop is required
     // to co-iterate multiple sparse tensors.
     assert(!isSparseInput || !isSparse);
     if (isSparse) {
       tid = t;
       dstLvl = l;
     }
     isSparseInput = isSparseInput || isSparse;
   }
 
   const auto reassoc = getCollapseReassociation(tid, dstLvl);
   // TODO: support dynamic slices.
   // Use the first source-level here to build the loop bound (which is
   // also the biggest range).
   const Level srcLvl = reassoc.front();
   const Value step = constantIndex(builder, loc, 1);
   /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
-  const Value lo = isSparseInput ? posits[tid][srcLvl]  // current position
+  const Value lo = isSparseInput ? posits[*tid][srcLvl] // current position
                                  : loopSeqStack.back(); // universal index
-  const Value hi = highs[tid][srcLvl];
+  const Value hi = highs[*tid][srcLvl];
 
   Operation *loop = nullptr;
   Value iv;
   if (isParallel) {
-    assert(collapseReassoc[tid] == nullptr);
+    assert(collapseReassoc[*tid] == nullptr);
     scf::ParallelOp parOp =
         builder.create<scf::ParallelOp>(loc, lo, hi, step, reduc);
     builder.setInsertionPointToStart(parOp.getBody());
     assert(parOp.getNumReductions() == reduc.size());
     iv = parOp.getInductionVars()[0];
 
     // In-place update on the reduction variable vector.
     // Note that the init vals is not the actual reduction variables but instead
if (isParallel) {
     for (int i = 0, e = reduc.size(); i < e; i++)
       reduc[i] = forOp.getRegionIterArg(i);
     loop = forOp;
   }
   assert(loop && iv);
 
   Value crd;
   if (isSparseInput) {
-    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[*tid].getType()));
     // For COO, the position is the same across consecutive levels.
     /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
-    llvm::for_each(reassoc,
-                   [this, tid, iv](Level srcLvl) { posits[tid][srcLvl] = iv; });
+    llvm::for_each(
+        reassoc, [this, tid, iv](Level srcLvl) { posits[*tid][srcLvl] = iv; });
     crd = genSparseCrd(builder, loc, tid, dstLvl);
   } else {
     // Dense tensor, the coordinate is the inducation variable.
     crd = iv;
   }
 
-  if (isSparseSlices[tid] && isSparseInput) {
+  if (isSparseSlices[*tid] && isSparseInput) {
     // For sparse level slices, we need to filter out invalid coordinates that
     // are not included in the slice.
     SmallVector<Type> types;
     for (Value red : reduc)
       types.push_back(red.getType());
 
     auto [trans, pred] = genSliceLegitPredicate(builder, loc, crd, tid, srcLvl);
     bool hasReduc = !types.empty();
if (hasReduc) {
       builder.create<scf::YieldOp>(loc, reduc);
     }
     // Set the insertion point to matched branch.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
     crd = trans;
   }
 
   assert(crd);
-  coords[tid][srcLvl] = crd;
+  coords[*tid][srcLvl] = crd;
   // NOTE: we can also prepare for next level here in advance
   // Push the loop into stack
   loopStack.emplace_back(ArrayRef<TensorId>(tid), ArrayRef<Level>(srcLvl), loop,
                          builder.getInsertionBlock(), crd, loopTag);
   // Emit extra locals.
   emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tids, lvls);
 
   return loop;
 }
 
 Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level lvl,
     AffineExpr affine, MutableArrayRef<Value> reduc) {
   assert(isValidLevel(tid, lvl));
-  assert(!affine.isa<AffineDimExpr>() && !isDenseDLT(lvlTypes[tid][lvl]));
+  assert(!affine.isa<AffineDimExpr>() && !isDenseDLT(lvlTypes[*tid][lvl]));
   // We can not re-enter the same level.
-  assert(!coords[tid][lvl]);
+  assert(!coords[*tid][lvl]);
 
   // TODO: We should instead use a whileOp for filter loop to allow early
   // break when exceeding (for ordered levels).
   // TODO: There are many other potiential opportunities that we might apply in
   // the future. E.g., we could use binary search to locate positions.
   const Value step = constantIndex(builder, loc, 1);
-  const Value pLo = posits[tid][lvl];
-  const Value pHi = highs[tid][lvl];
+  const Value pLo = posits[*tid][lvl];
+  const Value pHi = highs[*tid][lvl];
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, pLo, pHi, step, reduc);
 
   // In-place update on the reduction variable vector.
   assert(forOp.getNumRegionIterArgs() == reduc.size());
   for (int i = 0, e = reduc.size(); i < e; i++)
     reduc[i] = forOp.getRegionIterArg(i);
 
   builder.setInsertionPointToStart(forOp.getBody());
   // The induction variable gives the position.
   const Value pos = forOp.getInductionVar();
-  posits[tid][lvl] = pos;
+  posits[*tid][lvl] = pos;
   // Generating a load on the coordinates array yields the crd.
-  const Value mem = coordinatesBuffers[tid][lvl];
+  const Value mem = coordinatesBuffers[*tid][lvl];
   const Value crd = genIndexLoad(builder, loc, mem, pos);
-  coords[tid][lvl] = crd;
+  coords[*tid][lvl] = crd;
 
   // Generate an if-condition to filter out coordinates that are not
   // equal to the result of the affine expression.
   Value expected = genAffine(builder, loc, affine);
   auto pred = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, crd,
                                             expected);
   SmallVector<Type> types;
   for (Value red : reduc) {
Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
   loopStack.emplace_back(ArrayRef<TensorId>(tid), ArrayRef<Level>(lvl), forOp,
                          builder.getInsertionBlock(), crd, nullptr);
   return forOp;
 }
 
 void LoopEmitter::genDenseAffineAddress(OpBuilder &builder, Location loc,
                                         TensorId tid, Level lvl,
                                         AffineExpr lvlExpr) {
-  assert(isDenseDLT(lvlTypes[tid][lvl]));
+  assert(isDenseDLT(lvlTypes[*tid][lvl]));
   // For dense levels, the level-coordinate also serves as the position.
   Value lvlCrd = genAffine(builder, loc, lvlExpr);
-  posits[tid][lvl] = genAddress(builder, loc, tid, lvl, lvlCrd);
+  posits[*tid][lvl] = genAddress(builder, loc, tid, lvl, lvlCrd);
 }
 
 Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
     OpBuilder &builder, Location loc, ArrayRef<TensorId> tids,
     ArrayRef<Level> lvls, bool needsUniv, MutableArrayRef<Value> reduc) {
   assert(tids.size() == lvls.size());
   SmallVector<Type> types;
   SmallVector<Value> operands;
   // Construct the while-loop with a parameter for each coordinate.
   const Type indexType = builder.getIndexType();
   for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
-    const auto lvlTp = lvlTypes[tid][lvl];
+    const auto lvlTp = lvlTypes[*tid][lvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
-        if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
+        if (!isUniqueDLT(lvlTypes[*tid][reassoc[i]])) {
           // This is the segment high for each non-unique levels.
           types.push_back(indexType);
           operands.push_back(constantIndex(builder, loc, 0));
         }
       }
-      const auto pos = posits[tid][reassoc.front()];
+      const auto pos = posits[*tid][reassoc.front()];
       assert(pos);
       types.push_back(indexType);
       operands.push_back(pos);
     }
   }
   // The position where user-supplied reduction variable starts.
   for (Value rec : reduc) {
     types.push_back(rec.getType());
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
   for (auto [t, lvl] : llvm::zip(tids, lvls)) {
     const TensorId tid = t; // Why `t` can not be captured by lambda?
-    const auto lvlTp = lvlTypes[tid][lvl];
+    const auto lvlTp = lvlTypes[*tid][lvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, lvl);
-      assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+      assert(reassoc.size() == 1 || isUniqueCOOType(tensors[*tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
-        if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
+        if (!isUniqueDLT(lvlTypes[*tid][reassoc[i]])) {
           // Links the SSA chain for segHi.
-          segHi[tid][reassoc[i]] = after->getArgument(o++);
+          segHi[*tid][reassoc[i]] = after->getArgument(o++);
         }
       }
       Value op1 = before->getArgument(o);
       // We used the first level bound as the bound the collapsed set of levels.
-      Value op2 = highs[tid][reassoc.front()];
+      Value op2 = highs[*tid][reassoc.front()];
       Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
                                                 op1, op2);
       cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
       // Update positions
       Value pos = after->getArgument(o++);
       // For COO, the position is the same across consecutive levels.
       /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
       llvm::for_each(reassoc, [this, tid, pos](Level srcLvl) {
-        posits[tid][srcLvl] = pos;
+        posits[*tid][srcLvl] = pos;
       });
     }
   }
   builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
 
   // Generates while body.
   builder.setInsertionPointToStart(&whileOp.getAfter().front());
 
   SmallVector<std::pair<Value, unsigned>> slicesPreds;
   unsigned i = 0;
   for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
     // Prepares for next level.
-    const auto lvlTp = lvlTypes[tid][lvl];
+    const auto lvlTp = lvlTypes[*tid][lvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
-      coords[tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
-      if (isSparseSlices[tid]) {
+      coords[*tid][lvl] = genSparseCrd(builder, loc, tid, lvl);
+      if (isSparseSlices[*tid]) {
         auto [trans, pred] =
-            genSliceLegitPredicate(builder, loc, coords[tid][lvl], tid, lvl);
+            genSliceLegitPredicate(builder, loc, coords[*tid][lvl], tid, lvl);
         slicesPreds.emplace_back(pred, i);
         // Updates to the relative coordinate to the slice.
-        coords[tid][lvl] = trans;
+        coords[*tid][lvl] = trans;
       }
       i++;
     }
   }
 
   if (!slicesPreds.empty()) {
     // Skips invalid loop iteration when slice coordinate is inapplicable.
     SmallVector<Value> yields(after->getArguments());
if (!slicesPreds.empty()) {
     // If all slices are legit, start the user generated code.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   }
 
   Value min;
   // Finds the minimum coordinate
   if (!needsUniv) {
     for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
-      const auto lvlTp = lvlTypes[tid][lvl];
+      const auto lvlTp = lvlTypes[*tid][lvl];
       if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
-        const auto crd = coords[tid][lvl];
+        const auto crd = coords[*tid][lvl];
         if (min) {
           Value cmp = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ult, crd, min);
           min = builder.create<arith::SelectOp>(loc, cmp, crd, min);
         } else {
           min = crd;
         }
       }
     }
   } else {
     assert(!min);
     // Otherwise, universal index is the minimal pos.
     min = after->getArguments().back();
   }
 
   // Sets up the loop stack.
   loopStack.emplace_back(tids, lvls, whileOp, builder.getInsertionBlock(), min,
                          loopTag);
   assert(loopStack.size() == loopSeqStack.size());
 
   for (auto [tid, dstLvl] : llvm::zip(tids, lvls)) {
     const auto reassoc = getCollapseReassociation(tid, dstLvl);
-    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+    assert(reassoc.size() == 1 || isUniqueCOOType(tensors[*tid].getType()));
     // TODO: Refactors this into smaller functions.
     // NOTE: For all the collapsed level (except for the last one, that is why
     // the loop ends with `reassoc.size() - 1`), as each iteration is advanced
     // by the segment size of the last level, which does not always invalidate
     // the segment size for the previous levels, thus we need to propagate the
     // segment sizes across loop iterations and only forward if needed.
     //
     // E.g., for a COO tensor with the following coordinates array.
     // (0, 0, 1),
     // (0, 0, 2),
     // (1, 1, 1),
     // segHi[lvl=0] = segHi[lvl=1] = 2
     // segHi[lvl=2] = 1,
     // the first iteration does not invalidate segHi[0] and segHi[1]
     for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
       const Level srcLvl = reassoc[i];
-      if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
-        const Value pos = posits[tid][srcLvl];
-        const auto oldSegHi = segHi[tid][srcLvl];
+      if (!isUniqueDLT(lvlTypes[*tid][srcLvl])) {
+        const Value pos = posits[*tid][srcLvl];
+        const auto oldSegHi = segHi[*tid][srcLvl];
         assert(oldSegHi);
         Value newSegHi = builder.create<arith::CmpIOp>(
             loc, arith::CmpIPredicate::uge, pos, oldSegHi);
         auto ifNewSegHi = builder.create<scf::IfOp>(loc, builder.getIndexType(),
                                                     newSegHi, true);
         {
           OpBuilder::InsertionGuard guard(builder);
           builder.setInsertionPointToStart(ifNewSegHi.thenBlock());
-          builder.create<scf::YieldOp>(loc,
-                                       genSegmentHigh(builder, loc, tid, srcLvl,
-                                                      pos, highs[tid][srcLvl]));
+          builder.create<scf::YieldOp>(
+              loc, genSegmentHigh(builder, loc, tid, srcLvl, pos,
+                                  highs[*tid][srcLvl]));
           // Else, resues the same segment high.
           builder.setInsertionPointToStart(ifNewSegHi.elseBlock());
           builder.create<scf::YieldOp>(loc, oldSegHi);
         }
-        highs[tid][srcLvl + 1] = segHi[tid][srcLvl] = ifNewSegHi.getResult(0);
+        highs[*tid][srcLvl + 1] = segHi[*tid][srcLvl] = ifNewSegHi.getResult(0);
       }
     };
     const auto srcLvl = reassoc.back();
-    if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
-      segHi[tid][srcLvl] = genSegmentHigh(
-          builder, loc, tid, srcLvl, posits[tid][srcLvl], highs[tid][srcLvl]);
+    if (!isUniqueDLT(lvlTypes[*tid][srcLvl])) {
+      segHi[*tid][srcLvl] = genSegmentHigh(
+          builder, loc, tid, srcLvl, posits[*tid][srcLvl], highs[*tid][srcLvl]);
     }
   }
 
   // Emits extra locals
   emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tids, lvls);
 
   // Updates reduction variables
   assert(after->getNumArguments() == o + reduc.size() + (needsUniv ? 1 : 0));
   // In-place update on reduction variable.
   for (unsigned i = 0, e = reduc.size(); i < e; i++)
     reduc[i] = after->getArgument(o + i);
 
   return whileOp;
 }
 
 void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                              TensorId tid, Level dstLvl) {
   assert(isValidLevel(tid, dstLvl));
-  const auto lvlTp = lvlTypes[tid][dstLvl];
+  const auto lvlTp = lvlTypes[*tid][dstLvl];
 
   if (isDenseDLT(lvlTp))
     return;
 
   const Value c0 = constantIndex(builder, loc, 0);
   const Value c1 = constantIndex(builder, loc, 1);
   for (const Level srcLvl : getCollapseReassociation(tid, dstLvl)) {
     // Either the first level, or the previous level has been set.
     /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
-    assert(srcLvl == 0 || posits[tid][srcLvl - 1]);
+    assert(srcLvl == 0 || posits[*tid][srcLvl - 1]);
     if (!isCompressedDLT(lvlTp) && !isSingletonDLT(lvlTp))
       continue;
     if (isCompressedDLT(lvlTp)) {
-      const Value mem = positionsBuffers[tid][srcLvl];
+      const Value mem = positionsBuffers[*tid][srcLvl];
 
-      const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
-      posits[tid][srcLvl] = genIndexLoad(builder, loc, mem, pLo);
+      const Value pLo = srcLvl == 0 ? c0 : posits[*tid][srcLvl - 1];
+      posits[*tid][srcLvl] = genIndexLoad(builder, loc, mem, pLo);
 
       const Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
-      highs[tid][srcLvl] = genIndexLoad(builder, loc, mem, pHi);
+      highs[*tid][srcLvl] = genIndexLoad(builder, loc, mem, pHi);
       return;
     }
     if (isSingletonDLT(lvlTp)) {
-      const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
-      posits[tid][srcLvl] = pLo;
+      const Value pLo = srcLvl == 0 ? c0 : posits[*tid][srcLvl - 1];
+      posits[*tid][srcLvl] = pLo;
 
       // If we are coiterating non-unique levels, then use pHi=segHi;
       // otherwise use pHi=pLo+1.
       // NOTE: Just because the level is non-unique, that does not
       // guarantee that segHi is defined: because we only generate segHi
       // whenever coiterating, in order to improve code quality for the
       // non-coiterating cases.
-      const auto parentSegHi = segHi[tid][srcLvl - 1];
-      highs[tid][srcLvl] =
-          (!isUniqueDLT(lvlTypes[tid][srcLvl - 1]) && parentSegHi)
+      const auto parentSegHi = segHi[*tid][srcLvl - 1];
+      highs[*tid][srcLvl] =
+          (!isUniqueDLT(lvlTypes[*tid][srcLvl - 1]) && parentSegHi)
               ? parentSegHi
               : builder.create<arith::AddIOp>(loc, pLo, c1);
       return;
     }
   }
 
   llvm_unreachable("Unrecognized level-type!");
 }
 
 void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder,
                                                        Location loc,
                                                        ArrayRef<TensorId> tids,
                                                        ArrayRef<Level> lvls) {
   // Initialize dense positions. Note that we generate dense coordinates of the
   // output tensor unconditionally, since they may not appear in the lattice,
   // but may be needed for linearized codegen.
   assert(tids.size() == lvls.size());
   for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
-    if (isDenseDLT(lvlTypes[tid][lvl])) {
-      auto enc = getSparseTensorEncoding(tensors[tid].getType());
+    if (isDenseDLT(lvlTypes[*tid][lvl])) {
+      auto enc = getSparseTensorEncoding(tensors[*tid].getType());
       if (enc && !isSparseOutput(tid)) {
-        bool validPos = lvl == 0 || posits[tid][lvl - 1];
+        bool validPos = lvl == 0 || posits[*tid][lvl - 1];
         if (!validPos) {
           // We might not find the pos for the sparse output tensor as it is
           // unconditionally required by the sparsification.
           assert(isOutputTensor(tid));
           continue;
         }
-        posits[tid][lvl] =
+        posits[*tid][lvl] =
             genAddress(builder, loc, tid, lvl, loopStack.back().iv);
         // NOTE: we can also prepare for next lvl here in advance
       }
     }
   }
 }
 
 void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
for (unsigned i = 0, e = parOp.getResults().size(); i < e; i++)
       reduc[i] = parOp.getResult(i);
   }
 
   // Finished iterating a tensor, clean up
   // We only do the clean up on for loop as while loops do not necessarily
   // finish the iteration on a sparse tensor
   for (auto [tid, lvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
     // Reset to null.
-    coords[tid][lvl] = Value();
-    posits[tid][lvl] = Value();
+    coords[*tid][lvl] = Value();
+    posits[*tid][lvl] = Value();
     // Dense level, high is fixed.
-    if (!isDenseDLT(lvlTypes[tid][lvl]))
-      highs[tid][lvl] = Value();
+    if (!isDenseDLT(lvlTypes[*tid][lvl]))
+      highs[*tid][lvl] = Value();
   }
 }
 
 void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
                                 MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
   builder.setInsertionPointToEnd(loopInfo.userCodeBlock);
   Value iv = loopInfo.iv;
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value> operands;
   Value one = constantIndex(builder, loc, 1);
   for (auto [tid, dstLvl] : llvm::zip(loopInfo.tids, loopInfo.lvls)) {
-    const auto lvlTp = lvlTypes[tid][dstLvl];
+    const auto lvlTp = lvlTypes[*tid][dstLvl];
     if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
       const auto reassoc = getCollapseReassociation(tid, dstLvl);
-      assert(reassoc.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+      assert(reassoc.size() == 1 || isUniqueCOOType(tensors[*tid].getType()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         const Level srcLvl = reassoc[i];
-        if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
-          operands.push_back(segHi[tid][srcLvl]);
+        if (!isUniqueDLT(lvlTypes[*tid][srcLvl])) {
+          operands.push_back(segHi[*tid][srcLvl]);
           o++;
         }
       }
-      const Value crd = coords[tid][dstLvl];
-      const Value pos = posits[tid][dstLvl];
+      const Value crd = coords[*tid][dstLvl];
+      const Value pos = posits[*tid][dstLvl];
       Value cmp =
           builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, crd, iv);
       // If the loop contains a coiteration with non-unique level, we fast
       // forward all the duplicated coords by setting the position to the
       // segment high.
-      Value add = !isUniqueDLT(lvlTypes[tid][reassoc.back()])
-                      ? segHi[tid][reassoc.back()]
+      Value add = !isUniqueDLT(lvlTypes[*tid][reassoc.back()])
+                      ? segHi[*tid][reassoc.back()]
                       : builder.create<arith::AddIOp>(loc, pos, one);
 
       operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, pos));
       // Following loops continue iteration from the break point of the
       // current while loop.
       const Value newPos = whileOp->getResult(o++);
       // We need to define a new local variable for `tid` to avoid
       // warnings about "captured structured bindings are a C++20 extension".
       // FIXME(wrengr): define a helper function to capture this idiom!
       const TensorId newTid = tid;
       llvm::for_each(reassoc, [this, newTid, newPos](Level srcLvl) {
-        posits[newTid][srcLvl] = newPos;
+        posits[*newTid][srcLvl] = newPos;
       });
       // The coordinate is invalid now.
-      coords[tid][dstLvl] = nullptr;
+      coords[*tid][dstLvl] = nullptr;
       // The segment high is invalid now.
-      segHi[tid][dstLvl] = nullptr;
+      segHi[*tid][dstLvl] = nullptr;
       // highs remains unchanged.
     }
   }
 
   // Reduction value from users.
   for (auto &i : reduc) {
     operands.push_back(i);
     // In place update reduction variable.

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

LoopEmitter loopEmitter(
         ValueRange{input},
         StringAttr::get(getContext(), ForeachOp::getOperationName()));
     loopEmitter.initializeLoopEmit(rewriter, loc);
     for (Level l = 0; l < lvlRank; l++) {
       // TODO: provide utility function for loop sequences that only contains
       // one for loop?
       // FIXME(wrengr): what is this "ld" supposed to be really?
       const Level ld = op.getOrder() ? op.getOrder()->getDimPosition(l) : l;
-      const SmallVector<TensorId, 1> tids{0};
+      const SmallVector<TensorId, 1> tids{TensorId::Unchecked(0)};
       loopEmitter.enterNewLoopSeq(rewriter, loc, tids, ld);
       // Note that reduc will be taken care of by loop emitter and get updated
       // in place.
 
       loopEmitter.enterLoopOverTensorAtLvl(rewriter, loc, tids, l, reduc);
     }
 
     SmallVector<Value> lcvs;

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

 // to a more self-explanatory naming convention like "curLoop"-vs-"prevLoop"
 // (assuming that's the actual meaning behind the "idx"-vs-"ldx" convention).
 
 /// Determines if affine expression is invariant.
 static bool isInvariantAffine(AffineExpr a, ArrayRef<LoopId> loopStack,
                               LoopId ldx, bool &isAtLoop) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
-    const LoopId i = a.cast<AffineDimExpr>().getPosition();
+    const LoopId i = LoopId::Unchecked(a.cast<AffineDimExpr>().getPosition());
     if (i == ldx) {
       isAtLoop = true;
       // Must be invariant if we are at the given loop.
       return true;
     }
     bool isInvariant = false;
     for (LoopId l : loopStack) {
       isInvariant = (l == i);

 
 /// 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, loop_id::Iterator &filterLdx,
+                       DimLevelType dlt, LoopId::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)

 /// 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;
-  const loop_id::Range filterLdxRange = env.merger().getFilterLoopIds();
-  loop_id::Iterator filterLdx = filterLdxRange.begin();
+  const LoopId::Range filterLdxRange = env.merger().getFilterLoopIds();
+  LoopId::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();

 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
   for (const LoopId i : env.merger().getLoopIds()) {
-    if (inDegree[i] == 0) {
+    if (inDegree[*i] == 0) {
       if (env.merger().isFilterLoop(i))
         filterIt.push_back(i);
-      else if (linalg::isReductionIterator(iteratorTypes[i]))
+      else if (linalg::isReductionIterator(iteratorTypes[*i]))
         redIt.push_back(i);
       else
         parIt.push_back(i);
     }
   }
 
   while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
     // We always choose in order of filter loop -> parallel loop -> reduction
while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
     //      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 (const LoopId dst : env.merger().getLoopIds()) {
-      if (adjM[src][dst] && --inDegree[dst] == 0) {
+      if (adjM[*src][*dst] && --inDegree[*dst] == 0) {
         if (env.merger().isFilterLoop(dst))
           filterIt.push_back(dst);
-        else if (linalg::isReductionIterator(iteratorTypes[dst]))
+        else if (linalg::isReductionIterator(iteratorTypes[*dst]))
           redIt.push_back(dst);
         else
           parIt.push_back(dst);
       }
     }
   }
   return env.topSortSize() == env.merger().getNumLoops();
 }

 /// `(LoopId,LoopId)`.
 static void addAffineOrderings(std::vector<std::vector<bool>> &adjM,
                                std::vector<unsigned> &inDegree, AffineExpr a,
                                AffineExpr b, std::optional<LoopId> fidx,
                                std::optional<LoopId> tidx) {
   if (!a && !b) {
     // Recursion leaf.
     assert(fidx && tidx);
-    const LoopId f = *fidx, t = *tidx;
+    const unsigned f = **fidx, t = **tidx;
     if (!adjM[f][t]) {
       adjM[f][t] = true;
       inDegree[t]++;
     }
     return;
   }
   // Picks an affine expression and expand (recurse into) it.
   const auto toExpand = a ? a : b;
   switch (toExpand.getKind()) {
   case AffineExprKind::DimId: {
     const std::optional<LoopId> idx{
-        toExpand.cast<AffineDimExpr>().getPosition()};
+        LoopId::Unchecked(toExpand.cast<AffineDimExpr>().getPosition())};
     if (toExpand == a)
       addAffineOrderings(adjM, inDegree, AffineExpr(), b, idx, tidx);
     else // toExpand == b
       addAffineOrderings(adjM, inDegree, a, AffineExpr(), fidx, idx);
     break;
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
if (!tldx) {
     // no ordering like d0->d2->d1).
     // TODO: this is obviously a sub optimal solution.
     if (!fldx && !fa.isa<AffineConstantExpr>()) {
       // Heuristic: we prefer parallel loop for lhs to reduce the chance
       // we add reduce < parallel ordering.
       finder.setPickedIterType(utils::IteratorType::parallel);
       finder.walkPostOrder(fa);
       fa = finder.getDimExpr();
-      fldx = finder.getDimExpr().getPosition();
+      fldx = LoopId::Unchecked(finder.getDimExpr().getPosition());
     }
     if (!ta.isa<AffineConstantExpr>()) {
       // Heuristic: we prefer reduction loop for rhs to reduce the chance
       // adding reduce < parallel ordering.
       finder.setPickedIterType(utils::IteratorType::reduction);
       finder.walkPostOrder(ta);
       ta = finder.getDimExpr();
-      tldx = finder.getDimExpr().getPosition();
+      tldx = LoopId::Unchecked(finder.getDimExpr().getPosition());
     }
   }
 }
 
 static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
                                           OpOperand *skip, SortMask mask,
                                           std::vector<std::vector<bool>> &adjM,
                                           std::vector<unsigned> &inDegree) {
for (Level lvl = 1; lvl < lvlRank; lvl++) {
     const AffineDimExpr fexp = finder.getDimExpr();
     const LoopId fldx = env.makeLoopId(fexp.getPosition());
 
     finder.walkPostOrder(ta);
     const AffineDimExpr texp = finder.getDimExpr();
     const LoopId tldx = env.makeLoopId(texp.getPosition());
 
     // d_x > d_y
-    if (!adjM[fldx][tldx]) {
-      adjM[fldx][tldx] = true;
-      inDegree[tldx]++;
+    if (!adjM[*fldx][*tldx]) {
+      adjM[*fldx][*tldx] = true;
+      inDegree[*tldx]++;
     }
 
     AffineDimCollector fCollector;
     fCollector.walkPostOrder(fa);
     AffineDimCollector tCollector;
     tCollector.walkPostOrder(ta);
 
     // make sure dx and dy is the last;
     for (auto fd : fCollector.dims) {
       const LoopId f = env.makeLoopId(fd.getPosition());
       if (f == fldx)
         continue;
-      if (!adjM[f][fldx]) {
-        adjM[f][fldx] = true;
-        inDegree[fldx]++;
+      if (!adjM[*f][*fldx]) {
+        adjM[*f][*fldx] = true;
+        inDegree[*fldx]++;
       }
     }
     for (auto td : tCollector.dims) {
       const LoopId t = env.makeLoopId(td.getPosition());
       if (t == tldx)
         continue;
-      if (!adjM[t][tldx]) {
-        adjM[t][tldx] = true;
-        inDegree[tldx]++;
+      if (!adjM[*t][*tldx]) {
+        adjM[*t][*tldx] = true;
+        inDegree[*tldx]++;
       }
     }
     // Since we only support affine addition, the order between two dim
     // expression does not really matters.
     // {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
     // This is to ensure that the affine expressions are reduced in sparse
     // tensor level ordering.
     // TODO: this ordering could probably be loosen if we support out-of-order
     // reduction.
     // TODO: the evaluation order need to be ensure to
     // support affine multiplication.
     for (auto fd : fCollector.dims) {
       const LoopId f = env.makeLoopId(fd.getPosition());
       if (f == fldx) // skip d_x
         continue;
 
       for (auto td : tCollector.dims) {
         const LoopId t = env.makeLoopId(td.getPosition());
         if (t == tldx) // skip d_y
           continue;
-        if (!adjM[f][t]) {
-          adjM[f][t] = true;
-          inDegree[t]++;
+        if (!adjM[*f][*t]) {
+          adjM[*f][*t] = true;
+          inDegree[*t]++;
         }
       }
     }
   }
 }
 
 /// Computes a topologically sorted iteration graph for the linalg operation.
 /// Ensures all tensors are visited in natural index order. This is
for (OpOperand &t : env.op()->getOpOperands()) {
     // TODO: Do we really need this?
     if (includesUndef(mask)) {
       const TensorId tid = env.makeTensorId(t.getOperandNumber());
       for (const LoopId i : env.merger().getLoopIds()) {
         const auto dltI = env.dlt(tid, i);
         if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
           for (const LoopId j : env.merger().getLoopIds())
             if (isUndefDLT(env.dlt(tid, j))) {
-              adjM[i][j] = true;
-              inDegree[j]++;
+              adjM[*i][*j] = true;
+              inDegree[*j]++;
             }
         } else {
           assert(isDenseDLT(dltI) || isUndefDLT(dltI));
         }
       }
     }
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.

 static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
                           SmallVectorImpl<Value> &args) {
   const Location loc = env.op().getLoc();
   const TensorId tid = env.makeTensorId(t->getOperandNumber());
   const auto map = env.op().getMatchingIndexingMap(t);
   const auto stt = getSparseTensorType(t->get());
   if (stt.hasEncoding()) {
     // For sparse tensors we only push the last-level's position onto `args`.
-    const auto pos = env.emitter().getPosits()[tid].back();
+    const auto pos = env.emitter().getPosits()[*tid].back();
     assert(pos);
     args.push_back(pos);
   } else {
     // For dense tensors we push all level's coordinates onto `args`.
     const Level lvlRank = stt.getLvlRank();
     assert(static_cast<Level>(map.getNumResults()) == lvlRank);
     for (Level l = 0; l < lvlRank; l++) {
       const auto lvlExpr = map.getResult(l);
       const auto lvlCrd = env.emitter().genAffine(builder, loc, lvlExpr);
       args.push_back(lvlCrd);
     }
   }
-  return env.emitter().getValBuffer()[tid];
+  return env.emitter().getValBuffer()[*tid];
 }
 
 /// Generates insertion code to implement dynamic tensor load.
 static Value genInsertionLoad(CodegenEnv &env, OpBuilder &builder,
                               OpOperand *t) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   // Direct lexicographic coordinate order, tensor loads as zero.

 static Value genTensorLoad(CodegenEnv &env, OpBuilder &builder, ExprId exp) {
   // Test if the load was hoisted to a higher loop nest.
   Value val = env.exp(exp).val;
   if (val)
     return val;
 
   // Load during insertion.
   linalg::GenericOp op = env.op();
-  OpOperand *t = &op->getOpOperand(env.exp(exp).tensor);
+  OpOperand *t = &op->getOpOperand(*(env.exp(exp).tensor));
   if (env.isSparseOutput(t)) {
     if (env.isCustomReduc())
       return genInsertionLoadReduce(env, builder, t);
     return genInsertionLoad(env, builder, t);
   }
   // Actual load.
   SmallVector<Value> args;
   Value ptr = genSubscript(env, builder, t, args);

 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
                           LoopId ldx, bool atStart) {
   if (!exp.isValid())
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
     // Inspect tensor indices.
-    bool isAtLoop = ldx == ::mlir::sparse_tensor::detail::kInvalidId;
+    bool isAtLoop = !ldx.isValid();
     linalg::GenericOp op = env.op();
-    OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
+    OpOperand &t = op->getOpOperand(*(env.exp(exp).tensor));
     const auto map = op.getMatchingIndexingMap(&t);
     const auto stt = getSparseTensorType(t.get());
     const Level lvlRank = stt.getLvlRank();
     assert(static_cast<Level>(map.getNumResults()) == lvlRank);
     for (Level l = 0; l < lvlRank; l++) {
       // FIXME: `toOrigDim` is deprecated.
       // FIXME: above we asserted that there are `lvlRank` many results,
       // but this is assuming there are in fact `dimRank` many results instead.
static Operation *genFor(CodegenEnv &env, OpBuilder &builder, bool isOuter,
 
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     if (env.merger().isFilterLoop(ldx)) {
       const TensorId tid = tids.front();
       const Level lvl = lvls.front();
       // tids/lvls must only have one value because filter loops only
       // corresponding to the one and only sparse tensor level.
       assert(isSparse && tids.size() == 1 && lvls.size() == 1);
-      OpOperand *t = &op->getOpOperand(tid);
+      OpOperand *t = &op->getOpOperand(*tid);
       auto enc = getSparseTensorEncoding(t->get().getType());
       // Retrieves the affine expression for the filter loop.
       // FIXME: `toOrigDim` is deprecated.
       AffineExpr a =
           op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, lvl));
       return env.emitter().enterFilterLoopOverTensorAtLvl(builder, loc, tid,
                                                           lvl, a, reduc);
     }
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
   env.merger().foreachTensorLoopId(
       p, /*simple=*/true,
       [&](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
           DimLevelType dlt, bool /*unused*/) {
         assert(ldx == env.merger().loop(b));
         Value clause;
         if (isCompressedDLT(dlt) || isSingletonDLT(dlt)) {
           assert(lvl.has_value());
-          const Value crd = env.emitter().getCoords()[tid][*lvl];
+          const Value crd = env.emitter().getCoords()[*tid][*lvl];
           const Value lvar = env.getLoopVar(ldx);
           clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                  crd, lvar);
         } else {
           assert(isDenseDLT(dlt) || isUndefDLT(dlt));
           clause = constantI1(builder, loc, true);
         }
         cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   return false;
 }
 
 static void genConstantDenseAddressFromLevel(CodegenEnv &env,
                                              OpBuilder &builder, TensorId tid,
                                              Level startLvl) {
   // TODO: Handle affine expression on output tensor.
   linalg::GenericOp op = env.op();
-  assert(tid < op.getNumDpsInputs());
-  OpOperand *input = op.getDpsInputOperands()[tid];
+  assert(*tid < op.getNumDpsInputs());
+  OpOperand *input = op.getDpsInputOperands()[*tid];
   const auto lvlExprs = op.getMatchingIndexingMap(input).getResults();
   const auto enc = getSparseTensorEncoding(input->get().getType());
   if (enc) {
     const Location loc = op.getLoc();
     const TensorId tid = env.makeTensorId(input->getOperandNumber());
     const Level lvlRank = enc.getLvlRank();
     assert(lvlExprs.size() == static_cast<size_t>(lvlRank));
     // FIXME: there is dim/lvl confusion here

 
 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.
   const TensorId hi = env.makeTensorId(env.op().getNumDpsInputs());
-  for (const TensorId tid : tensor_id::Range(0, hi))
+  for (const TensorId tid : TensorId::fromZeroUpTo(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) {
   const BitVector &simple = env.lat(li).simple;
   const TensorId outTid = env.merger().getOutTensorID();
   const std::optional<Level> outLvl = env.merger().getLvl(outTid, ldx);
 
   unsigned numloopCond = 0;
   bool hasNonUnique = false;
 
   env.merger().foreachTensorLoopId(
       li, [&, ldx](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
                    DimLevelType dlt, bool isIdxReduc) {
-        if (simple[b]) {
+        if (simple[*b]) {
           if (isIdxReduc) {
             tids.push_back(tid);
             lvls.push_back(*lvl);
             numloopCond++;
             return;
           }
           if (isUndefDLT(dlt)) {
             // An undefined dlt in the lattices, we probably mean to
env.merger().foreachTensorLoopId(
           lvls.push_back(*lvl);
           numloopCond++;
         } else if (isDenseDLT(dlt)) {
           tids.push_back(tid);
           lvls.push_back(*lvl);
         } else {
           assert(isUndefDLT(dlt));
           linalg::GenericOp op = env.op();
-          if (tid >= op.getNumDpsInputs())
+          if (*tid >= op.getNumDpsInputs())
             // We only handle affine expression on input tensors (for now).
             return;
-          OpOperand *operand = &op->getOpOperand(tid);
+          OpOperand *operand = &op->getOpOperand(*tid);
           const auto stt = getSparseTensorType(operand->get());
           // Non-annotated dense tensors requires no special handling.
           if (!stt.hasEncoding())
             return;
 
           ArrayRef<AffineExpr> affines =
               op.getMatchingIndexingMap(operand).getResults();
           const Level lvlRank = stt.getLvlRank();
if (at == env.topSortSize()) {
     const LoopId ldx = env.topSortAt(at - 1);
     Value rhs = genExp(env, rewriter, exp, ldx);
     genTensorStore(env, rewriter, exp, rhs);
     return;
   }
 
   // Construct iteration lattices for current loop index, with L0 at top.
   const LoopId idx = env.topSortAt(at);
-  const LoopId ldx = at == 0 ? ::mlir::sparse_tensor::detail::kInvalidId
-                             : env.topSortAt(at - 1);
+  const LoopId ldx = at == 0 ? LoopId::Invalid() : env.topSortAt(at - 1);
   const LatSetId lts =
       env.merger().optimizeSet(env.merger().buildLattices(exp, idx));
 
   // Start a loop sequence.
   bool needsUniv = startLoopSeq(env, rewriter, exp, at, idx, ldx, lts);
 
   // Emit a loop for every lattice point L0 >= Li in this loop sequence.
   //
for (OpOperand *t : env.op().getDpsInputOperands()) {
       auto dstEnc = SparseTensorEncodingAttr::get(
           getContext(), srcEnc.getDimLevelType(),
           permute(env, env.op().getMatchingIndexingMap(t)), // new order
           srcEnc.getHigherOrdering(), srcEnc.getPosWidth(),
           srcEnc.getCrdWidth());
       auto dstTp = RankedTensorType::get(srcTp.getShape(),
                                          srcTp.getElementType(), dstEnc);
       auto convert = rewriter.create<ConvertOp>(tval.getLoc(), dstTp, tval);
-      rewriter.updateRootInPlace(env.op(),
-                                 [&]() { env.op()->setOperand(tid, convert); });
+      rewriter.updateRootInPlace(
+          env.op(), [&]() { env.op()->setOperand(*tid, convert); });
       rewriter.setInsertionPointAfter(env.op());
       rewriter.create<bufferization::DeallocTensorOp>(tval.getLoc(), convert);
       return success();
     }
     // Cannot be resolved with a single conversion.
     // TODO: convert more than one?
     return failure();
   }

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

static ExpArity getExpArity(TensorExp::Kind k) {
   }
   llvm_unreachable("unexpected kind");
 }
 
 //===----------------------------------------------------------------------===//
 // Constructors.
 //===----------------------------------------------------------------------===//
 
-TensorExp::TensorExp(TensorExp::Kind k, unsigned x, ExprId y, Value v,
+constexpr TensorExp::TensorExp(TensorId t)
+    : kind(TensorExp::Kind::kTensor), tensor(t), val(), op(nullptr) {}
+
+constexpr TensorExp::TensorExp(LoopId i)
+    : kind(TensorExp::Kind::kLoopVar), loop(i), val(), op(nullptr) {}
+
+constexpr TensorExp::TensorExp(Value v)
+    : kind(TensorExp::Kind::kInvariant), children(), val(v), op(nullptr) {}
+
+constexpr TensorExp::TensorExp(Kind k, ExprId e0, ExprId e1, Value v,
                                Operation *o)
-    : kind(k), val(v), op(o) {
+    : kind(k), children(e0, e1), val(v), op(o) {
   switch (kind) {
   // Leaf.
   case TensorExp::Kind::kTensor:
-    assert(x != detail::kInvalidId && !y.isValid() && !v && !o);
-    tensor = x;
-    return;
   case TensorExp::Kind::kInvariant:
-    assert(x == detail::kInvalidId && !y.isValid() && v && !o);
-    return;
   case TensorExp::Kind::kLoopVar:
-    assert(x != detail::kInvalidId && !y.isValid() && !v && !o);
-    loop = x;
-    return;
+    assert(false && "expected non-leaf expression kind");
+    break;
   // Unary operations.
   case TensorExp::Kind::kAbsF:
   case TensorExp::Kind::kAbsC:
   case TensorExp::Kind::kAbsI:
   case TensorExp::Kind::kCeilF:
   case TensorExp::Kind::kFloorF:
   case TensorExp::Kind::kSqrtF:
   case TensorExp::Kind::kSqrtC:
   case TensorExp::Kind::kExpm1F:
   case TensorExp::Kind::kExpm1C:
   case TensorExp::Kind::kLog1pF:
   case TensorExp::Kind::kLog1pC:
   case TensorExp::Kind::kSinF:
   case TensorExp::Kind::kSinC:
   case TensorExp::Kind::kTanhF:
   case TensorExp::Kind::kTanhC:
   case TensorExp::Kind::kNegF:
   case TensorExp::Kind::kNegC:
   case TensorExp::Kind::kNegI:
   case TensorExp::Kind::kCIm:
   case TensorExp::Kind::kCRe:
-    assert(x != detail::kInvalidId && !y.isValid() && !v && !o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && !e1.isValid() && !v && !o);
     return;
   case TensorExp::Kind::kTruncF:
   case TensorExp::Kind::kExtF:
   case TensorExp::Kind::kCastFS:
   case TensorExp::Kind::kCastFU:
   case TensorExp::Kind::kCastSF:
   case TensorExp::Kind::kCastUF:
   case TensorExp::Kind::kCastS:
   case TensorExp::Kind::kCastU:
   case TensorExp::Kind::kCastIdx:
   case TensorExp::Kind::kTruncI:
   case TensorExp::Kind::kBitCast:
-    assert(x != detail::kInvalidId && !y.isValid() && v && !o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && !e1.isValid() && v && !o);
     return;
   case TensorExp::Kind::kBinaryBranch:
   case TensorExp::Kind::kSelect:
-    assert(x != detail::kInvalidId && !y.isValid() && !v && o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && !e1.isValid() && !v && o);
     return;
   case TensorExp::Kind::kUnary:
-    // No assertion on y can be made, as the branching paths involve both
+    // No assertion on `e1` can be made, as the branching paths involve both
     // a unary (`mapSet`) and binary (`disjSet`) pathway.
-    assert(x != detail::kInvalidId && !v && o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && !v && o);
     return;
   // Binary operations.
   case TensorExp::Kind::kMulF:
   case TensorExp::Kind::kMulC:
   case TensorExp::Kind::kMulI:
   case TensorExp::Kind::kDivF:
   case TensorExp::Kind::kDivC:
   case TensorExp::Kind::kDivS:
   case TensorExp::Kind::kDivU:
   case TensorExp::Kind::kAddF:
   case TensorExp::Kind::kAddC:
   case TensorExp::Kind::kAddI:
   case TensorExp::Kind::kSubF:
   case TensorExp::Kind::kSubC:
   case TensorExp::Kind::kSubI:
   case TensorExp::Kind::kAndI:
   case TensorExp::Kind::kOrI:
   case TensorExp::Kind::kXorI:
   case TensorExp::Kind::kShrS:
   case TensorExp::Kind::kShrU:
   case TensorExp::Kind::kShlI:
-    assert(x != detail::kInvalidId && y.isValid() && !v && !o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && e1.isValid() && !v && !o);
     return;
   case TensorExp::Kind::kBinary:
   case TensorExp::Kind::kReduce:
-    assert(x != detail::kInvalidId && y.isValid() && !v && o);
-    children.e0 = ExprId{x};
-    children.e1 = y;
+    assert(e0.isValid() && e1.isValid() && !v && o);
     return;
   }
   llvm_unreachable("unexpected kind");
 }
 
 Merger::Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
                unsigned numFilterLoops, unsigned maxLvlRank)
     : outTensor(numInputOutputTensors - 1),
       syntheticTensor(numInputOutputTensors),
-      numTensors(numInputOutputTensors + 1), numNativeLoops(numNativeLoops),
-      numLoops(numNativeLoops + numFilterLoops), hasSparseOut(false),
-      lvlTypes(numTensors,
-               std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
-      loopToLvl(numTensors,
-                std::vector<std::optional<Level>>(numLoops, std::nullopt)),
-      lvlToLoop(numTensors,
+      factory(numInputOutputTensors + 1, numNativeLoops + numFilterLoops),
+      numNativeLoops(numNativeLoops), hasSparseOut(false),
+      lvlTypes(getNumTensors(),
+               std::vector<DimLevelType>(getNumLoops(), DimLevelType::Undef)),
+      loopToLvl(getNumTensors(),
+                std::vector<std::optional<Level>>(getNumLoops(), std::nullopt)),
+      lvlToLoop(getNumTensors(),
                 std::vector<std::optional<LoopId>>(maxLvlRank, std::nullopt)),
-      loopToDependencies(numLoops, std::vector<std::optional<Level>>(
-                                       numTensors, std::nullopt)),
-      levelToDependentIdx(numTensors, std::vector<std::vector<LoopId>>(
+      loopToDependencies(getNumLoops(), std::vector<std::optional<Level>>(
+                                       getNumTensors(), std::nullopt)),
+      levelToDependentIdx(getNumTensors(), std::vector<std::vector<LoopId>>(
                                           maxLvlRank, std::vector<LoopId>())),
-      loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
+      loopBounds(getNumLoops(), std::make_pair(TensorId{getNumTensors()}, getNumLoops())) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addTensorExp(TensorId t) {
   assert(isValidTensorId(t));
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(TensorExp::Kind::kTensor, t, ExprId(), Value(),
-                          nullptr);
+  tensorExps.emplace_back(t);
   return eNew;
 }
 
 ExprId Merger::addLoopVarExp(LoopId i) {
   assert(isValidLoopId(i));
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(TensorExp::Kind::kLoopVar, i, ExprId(), Value(),
-                          nullptr);
+  tensorExps.emplace_back(i);
   return eNew;
 }
 
 ExprId Merger::addInvariantExp(Value v) {
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(TensorExp::Kind::kInvariant, detail::kInvalidId,
-                          ExprId(), v, nullptr);
+  tensorExps.emplace_back(v);
   return eNew;
 }
 
 ExprId Merger::addExp(TensorExp::Kind k, ExprId e0, ExprId e1, Operation *op) {
   assert(k > TensorExp::Kind::kLoopVar);
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(k, e0.value, e1, Value(), op);
+  tensorExps.emplace_back(k, e0, e1, Value(), op);
   return eNew;
 }
 
 ExprId Merger::addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op) {
   assert(k > TensorExp::Kind::kLoopVar);
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(k, e.value, ExprId(), v, op);
+  tensorExps.emplace_back(k, e, ExprId(), v, op);
   return eNew;
 }
 
 LatPointId Merger::addLat(TensorId t, LoopId i, ExprId e) {
   const LatPointId pNew(latPoints.size());
-  const unsigned size = numLoops * numTensors;
+  const unsigned size = getNumLoops() * getNumTensors();
   const TensorLoopId b = makeTensorLoopId(t, i);
   latPoints.emplace_back(size, e);
-  latPoints[pNew.value].bits.set(b);
+  latPoints[pNew.value].bits.set(b.value);
   return pNew;
 }
 
 LatPointId Merger::addLat(const BitVector &bits, ExprId e) {
-  assert(bits.size() == numLoops * numTensors);
+  assert(bits.size() == getNumLoops() * getNumTensors());
   const LatPointId pNew(latPoints.size());
   latPoints.emplace_back(bits, e);
   return pNew;
 }
 
 LatSetId Merger::addSet() {
   const LatSetId sNew(latSets.size());
   latSets.emplace_back();
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 (const TensorLoopId b : tensor_loop_id::Range(bitsj))
-      if (!bitsi[b])
+    for (const TensorLoopId b : TensorLoopId::Range(bitsj))
+      if (!bitsi[b.value])
         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 (const TensorLoopId b : tensor_loop_id::Range(bits)) {
+  for (const TensorLoopId b : TensorLoopId::Range(bits)) {
     const auto dlt = getDimLevelType(b);
     if (isCompressedDLT(dlt) || isSingletonDLT(dlt))
       return true;
   }
   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))
+  for (const TensorLoopId b : TensorLoopId::Range(bits))
+    if (isLvlWithNonTrivialIdxExp(b))
       return true;
   return false;
 }
 
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
void Merger::dumpExp(ExprId e) const {
   const auto &expr = exp(e);
   switch (expr.kind) {
   // Leaf.
   case TensorExp::Kind::kTensor:
     if (expr.tensor == syntheticTensor)
       llvm::dbgs() << "synthetic_";
     else if (expr.tensor == outTensor)
       llvm::dbgs() << "output_";
-    llvm::dbgs() << "tensor_" << expr.tensor;
+    llvm::dbgs() << "tensor_" << expr.tensor.value;
     break;
   case TensorExp::Kind::kInvariant:
     llvm::dbgs() << "invariant";
     break;
   case TensorExp::Kind::kLoopVar:
-    llvm::dbgs() << "loopvar_" << expr.loop;
+    llvm::dbgs() << "loopvar_" << expr.loop.value;
     break;
   // Unary operations.
   case TensorExp::Kind::kAbsF:
   case TensorExp::Kind::kAbsC:
   case TensorExp::Kind::kAbsI:
   case TensorExp::Kind::kCeilF:
   case TensorExp::Kind::kFloorF:
   case TensorExp::Kind::kSqrtF:
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 (const TensorLoopId b : tensor_loop_id::Range(bits)) {
-    const TensorId t = tensor(b);
-    const LoopId i = loop(b);
+  for (const TensorLoopId b : TensorLoopId::Range(bits)) {
+    const auto t = tensor(b).value;
+    const auto i = loop(b).value;
     const auto dlt = lvlTypes[t][i];
     if (isLvlWithNonTrivialIdxExp(b))
       llvm::dbgs() << " DEP_" << t << "_" << i;
     else
       llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
   }
 }
 

 
 std::optional<ExprId> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   if (auto arg = v.dyn_cast<BlockArgument>()) {
     const TensorId tid = makeTensorId(arg.getArgNumber());
     // Any argument of the generic op that is not marked as a scalar
     // argument is considered a tensor, indexed by the implicit loop
     // bounds. This includes rank-0 tensor arguments.
     if (arg.getOwner()->getParentOp() == op) {
-      OpOperand &t = op->getOpOperand(tid);
+      OpOperand &t = op->getOpOperand(tid.value);
       if (!op.isScalar(&t))
         return addTensorExp(tid);
       v = t.get(); // get scalar value
     }
     // Any other argument (marked as scalar argument for the generic op
     // or belonging to an enveloping op) is considered invariant.
     return addInvariantExp(v);
   }

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

void expectLatPoint(LatSetId s, unsigned lo, PatternRef pattern,
     EXPECT_TRUE(latPointWithinRange(s, lo, 1, pattern, bits, simple));
   }
 
   /// Converts a vector of (loop, tensor) pairs to a bitvector with the
   /// corresponding bits set.
   BitVector loopsToBits(const std::vector<std::pair<LoopId, TensorId>> &loops) {
     BitVector testBits = BitVector(merger.getNumTensors(), false);
     for (auto [loop, tensor] : loops)
-      testBits.set(merger.makeTensorLoopId(tensor, loop));
+      testBits.set(*(merger.makeTensorLoopId(tensor, loop)));
     return testBits;
   }
 
   /// Returns true if the bits of the `k`th point in set `s` matches
   /// the given `bits`.  If `simple` is true, then compares the `lat.simple`
   /// field instead, to test the result after optimization
   bool compareBits(LatSetId s, unsigned k, const BitVector &bits, bool simple) {
     const auto &point = merger.lat(merger.set(s)[k]);