[mlir][sparse] Preliminary code changes for ExprId, LatPointId, LatSetId newtypes

[wrengr]

Summary

This commit contains several code changes which are ultimately required for converting the varions Merger identifiers from typedefs to newtypes. The actual implementation of the newtypes themselves has been split off into separate commits, in hopes of simplifying the review process.

Depends On D146561

[wrengr]

rebase & git-clang-format

[wrengr]

rebase

[aartbik]

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

31

Can you leave this in the same place for now (and not introduce a new MergerNewtypes.h). I am not convinced yes we really need two headers in the end, and even if so, we can have an intermediate revision that just moves this out. That way, I can focus on the real changes in the CL.

[wrengr]

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

31

I think having a separate header makes a lot of sense, since by the end (i.e., after D146693) that header comes to ~500 lines of code, all of which is boilerplate stuff that distracts from "the 'real' API" presented in the Merger header. (Fwiw, I did start on some templating stuff to try to reduce the amount of boilerplate; but that's shelved for the time being, since I was running into issues getting the public/private/friend access controls to work right.)

Though I can postpone and split up the move, so that:

• D146690 moves the ExprId, LatPointId, LatSetId types over (i.e., the same CL that converts those typedefs into newtypes)
• D146691 moves the TensorId, LoopId, TensorLoopId types over (to avoid cyclic dependency issues)

Does that work for you?

[aartbik]

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

31

Yeah, postponing the split works for me, so that I can decide in the other revision if it is really worth the code size explosion.

But the changes in this revision all look reasonable, so for the sake of quick progress, leaving the typedefs in header for now will speed up this review.

[wrengr]

Postponing the introduction of MergerNewtypes.h until D146690 (and D146691, D146693)

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

 namespace mlir {
 namespace sparse_tensor {
 
 // TODO: These type aliases currently only serve to make the code more
 // self-documenting, however because they are not type-checked they can
 // do nothing to prevent mixups.  We should really change them from mere
 // aliases to actual struct definitions, so that we can type-check them.
 
 /// Tensor identifiers.  The valid set of identifiers is defined by the
 /// first argument passed to the `Merger` ctor.
 using TensorId = unsigned;
 
 /// Loop identifiers.  The valid set of identifiers is defined by the
 /// second two arguments to the `Merger` ctor.
 ///
 /// These identifiers serve as proxies for the `$dim` argument to
 /// `linalg::IndexOp`, however the numerical value of a `LoopId` should
enum class TensorExp::Kind {
   kXorI,
   kShrS, // signed
   kShrU, // unsigned
   kShlI,
   kBinary, // semiring binary op
   kReduce, // semiring reduction op
 };
 
+//===----------------------------------------------------------------------===//
 /// Lattice point.  Each lattice point consists of a formal conjunction
 /// of `TensorLoopId`s, together with the identifier of the corresponding
 /// tensor expression.  The formal conjunction is represented as a set of
 /// `TensorLoopId`, where that set is implemented as a `BitVector`.
 struct LatPoint final {
-  /// Construct the lattice point from a given set of `TensorLoopId`s.
-  LatPoint(const BitVector &bits, ExprId e);
+  /// Construct a lattice point with the empty set of `TensorLoopId`s.
+  LatPoint(unsigned size, ExprId e) : bits(size, false), exp(e) {}
 
-  /// Construct a lattice point with `(t,i)` as the only `TensorLoopId`,
-  /// where `(t,i) < (numTensors,numLoops)`.
-  LatPoint(unsigned numTensors, unsigned numLoops, TensorId t, LoopId i,
-           ExprId e);
+  /// Construct a lattice point from the given set of `TensorLoopId`s.
+  LatPoint(const BitVector &bits, ExprId e) : bits(bits), exp(e) {}
 
   /// Conjunction of all `TensorLoopId`s involved in the tensor expression.
   BitVector bits;
 
   /// Simplified conjunction of `TensorLoopId` as bitvector.  This
   /// represents a simplified condition under which this tensor expression
   /// must execute. Pre-computed during codegen to avoid repeated eval.
   BitVector simple;
 
   /// Identifier of the tensor expression.
   ExprId exp;
 };
 
+//===----------------------------------------------------------------------===//
 /// A class to handle all iteration lattice operations. This class abstracts
 /// away from some implementation details of storing iteration lattices and
 /// tensor expressions. This allows for fine-tuning performance characteristics
 /// independently from the basic algorithm if bottlenecks are identified.
 class Merger {
 public:
   /// Constructs a merger for the given number of tensors, native loops, and
   /// filter loops. The user supplies the number of tensors involved in the
public:
   // the coordinate.
   //
   // TODO: would be cleaner to understand/document if the first argument
   // gave the number of input tensors, instead of the current number of
   // input+output tensors.
   Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
          unsigned numFilterLoops, unsigned maxLvlRank);
 
-  /// Constructs a new tensor expression, and returns its identifier.
-  /// The type of the `e0` argument varies according to the value of the
-  /// `k` argument, as described by the `TensorExp` ctor.
-  ExprId addExp(TensorExp::Kind k, unsigned e0, ExprId e1 = detail::kInvalidId,
-                Value v = Value(), Operation *op = nullptr);
-  ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr) {
-    return addExp(k, e, detail::kInvalidId, v, op);
-  }
-  ExprId addExp(TensorExp::Kind k, Value v, Operation *op = nullptr) {
-    return addExp(k, detail::kInvalidId, detail::kInvalidId, v, op);
+  //
+  // 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;
+  }
+
+  /// Safely converts the argument to a loop identifier.
+  constexpr LoopId makeLoopId(unsigned i) const {
+    assert(isValidLoopId(i));
+    return 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;
   }
 
+  //
+  // Allocating new expressions, points, and sets.
+  //
+
+  /// Constructs a new tensor expression, and returns its identifier.
+  ExprId addTensorExp(TensorId t);
+  /// Constructs a new loop-variable expression, and returns its identifier.
+  ExprId addLoopVarExp(LoopId i);
+  /// Constructs a new invariant expression, and returns its identifier.
+  ExprId addInvariantExp(Value v);
+  /// Constructs a new unary or binary expression, and returns its identifier.
+  ExprId addExp(TensorExp::Kind k, ExprId e0, ExprId e1 = detail::kInvalidId,
+                Operation *op = nullptr);
+  /// Constructs a new sesquinary expression, and returns its identifier.
+  /// Currently no sesquinary `Kind` allows specifying the `op`, but we
+  /// allow it anyways because `mapSet` is designed to allow it.
+  ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr);
+
   /// Constructs a new iteration lattice point, and returns its identifier.
   LatPointId addLat(TensorId t, LoopId i, ExprId e);
   LatPointId addLat(const BitVector &bits, ExprId e);
 
   /// Constructs a new (initially empty) set, and returns its identifier.
   LatSetId addSet();
 
   /// Computes a single conjunction of two lattice points by taking the "union"
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`.
-  TensorId tensor(TensorLoopId b) const { return b % numTensors; }
+  constexpr TensorId tensor(TensorLoopId b) const { return b % numTensors; }
   /// Gets the loop-identifier of the `TensorLoopId`.
-  LoopId loop(TensorLoopId b) const { return b / numTensors; }
+  constexpr LoopId loop(TensorLoopId b) const { return b / numTensors; }
 
   /// Get the total number of tensors (including the output-tensor and
-  /// synthetic-tensor).  The result is given the type `TensorId` since
-  /// the result is primarily used as an upper bound for `TensorId`s.
-  TensorId getNumTensors() const { return numTensors; }
+  /// synthetic-tensor).
+  constexpr unsigned getNumTensors() const { return numTensors; }
 
   /// Get the total number of loops (native loops + filter loops).
-  /// The result is given the type `LoopId` since the result will
-  /// generally be used as a for-loop upper bound.
-  LoopId getNumLoops() const { return numLoops; }
-  /// Get the number of native loops.  The result is given the type
-  /// `LoopId` since the result will generally be used as a for-loop
-  /// upper bound.
-  LoopId getNumNativeLoops() const { return numNativeLoops; }
-  /// Get the number of filter loops.  The result is given the type
-  /// `LoopId` since the result will generally be used as a for-loop
-  /// upper bound.
-  LoopId getNumFilterLoops() const { return numLoops - numNativeLoops; }
+  constexpr unsigned getNumLoops() const { return numLoops; }
+  /// Get the number of native loops.
+  constexpr unsigned getNumNativeLoops() const { return numNativeLoops; }
+  /// Get the number of filter loops.
+  constexpr unsigned getNumFilterLoops() const {
+    return numLoops - numNativeLoops;
+  }
   /// Get the identifier of the first filter-loop.
-  LoopId getStartingFilterLoopId() const { return getNumNativeLoops(); }
+  constexpr LoopId getStartingFilterLoopId() const {
+    return getNumNativeLoops();
+  }
 
   /// Returns true if `b` is the `i`th loop of the output tensor.
-  bool isOutTensor(TensorLoopId b, LoopId i) const {
-    assert(i < numLoops);
-    return b == numTensors * i + outTensor;
+  constexpr bool isOutTensor(TensorLoopId b, LoopId i) const {
+    return b == makeTensorLoopId(outTensor, i);
   }
 
   /// Get the output tensor's identifier.
-  TensorId getOutTensorID() const { return outTensor; }
+  constexpr TensorId getOutTensorID() const { return outTensor; }
   /// Get the synthetic tensor's identifier (used for all invariant
   /// tensor expressions).
-  TensorId getSynTensorID() const { return syntheticTensor; }
+  constexpr TensorId getSynTensorID() const { return syntheticTensor; }
 
-  bool isFilterLoop(LoopId i) const {
-    assert(i < numLoops);
+  constexpr bool isFilterLoop(LoopId i) const {
+    assert(isValidLoopId(i));
     return i >= numNativeLoops;
   }
 
   /// Returns true if the expression is `(kTensor t)`.
   bool expIsTensor(ExprId e, TensorId t) const {
-    return tensorExps[e].kind == TensorExp::Kind::kTensor &&
-           tensorExps[e].tensor == t;
+    const auto &expr = exp(e);
+    return expr.kind == TensorExp::Kind::kTensor && expr.tensor == t;
   }
 
   /// Returns true if the expression contains the tensor as an operand.
   bool expContainsTensor(ExprId e, TensorId t) const;
 
   /// Returns true if the expression contains a negation on output tensor.
   /// I.e., `- outTensor` or `exp - outputTensor`
   /// NOTE: this is an trivial tests in that it does not handle recursive
public:
   bool hasAnySparse(const BitVector &bits) const;
 
   /// 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(t < numTensors && i < numLoops);
+    assert(isValidTensorId(t) && isValidLoopId(i));
     return lvlTypes[t][i];
   }
 
   /// 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(t < numTensors && lvl < lvlToLoop[t].size());
+    assert(isValidLevel(t, lvl));
     return lvlToLoop[t][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(t < numTensors && i < numLoops);
+    assert(isValidTensorId(t) && isValidLoopId(i));
     return loopToLvl[t][i];
   }
   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(t < numTensors && i < numLoops && lvl < lvlToLoop[t].size() &&
-           isValidDLT(dlt));
+    assert(isValidLevel(t, lvl) && isValidLoopId(i) && isValidDLT(dlt));
     lvlTypes[t][i] = dlt;
     loopToLvl[t][i] = lvl;
     lvlToLoop[t][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, function_ref<void(TensorLoopId, TensorId,
-                                      std::optional<Level>, DimLevelType, bool)>
-                        callback) {
-    for (const TensorLoopId b : latPoints[p].bits.set_bits()) {
-      TensorId t = tensor(b);
+  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()) {
+      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(!getLvl(b).has_value());
+        assert(!optLvl.has_value());
         // Slice the tid along the dependent level to iterate current loop.
-        callback(b, t, loopToDependencies[loop(b)][t], getDimLevelType(b),
+        callback(b, t, loopToDependencies[loop(b)][t], lvlTp,
                  /*isIdxReduc=*/true);
       } else {
-        callback(b, t, getLvl(b), getDimLevelType(b), /*isIdxReduc=*/false);
+        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(lvl < numLoops);
+    assert(isValidLoopId(i) && isValidLevel(t, lvl));
     loopToDependencies[i][t] = lvl;
     levelToDependentIdx[t][lvl].push_back(i);
   }
 
   /// Whether the loop has dependent slice.
-  bool hasDependentLvl(LoopId i, TensorId tid) {
-    return loopToDependencies[i][tid].has_value();
+  bool hasDependentLvl(LoopId i, TensorId t) {
+    assert(isValidTensorId(t) && isValidLoopId(i));
+    return loopToDependencies[i][t].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];
   }
 
   /// Returns the defining [tid, lvl] for the loop.
   std::pair<TensorId, Level> getLoopDefiningLvl(LoopId i) const {
+    assert(isValidLoopId(i));
     return loopBounds[i];
   }
 
   /// Checks whether the TensorLoopId represents a tensor level with
   /// non-trivial index expression on it.
   bool isLvlWithNonTrivialIdxExp(TensorLoopId b) const {
-    return loopToDependencies[loop(b)][tensor(b)].has_value();
+    const TensorId t = tensor(b);
+    const LoopId i = loop(b);
+    assert(isValidTensorId(t) && isValidLoopId(i));
+    return loopToDependencies[i][t].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.
   ///
   /// NOTE: It is inadvisable to keep the reference alive for a long
   /// time (e.g., as in `TensorExpr &te = merger.exp(e)`), since insertions
   /// into the merger can cause data movement which will invalidate the
   /// underlying memory address.  This isn't just a problem with the `&`
   /// references, but also applies to the `ArrayRef`.  In particular,
   /// using `for (LatPointId p : merger.set(s))` will run into the same
   /// dangling-reference problems if the loop body inserts new sets.
-  const TensorExp &exp(ExprId e) const { return tensorExps[e]; }
-  const LatPoint &lat(LatPointId p) const { return latPoints[p]; }
-  ArrayRef<LatPointId> set(LatSetId s) const { return latSets[s]; }
+  const TensorExp &exp(ExprId e) const {
+    assert(isValidExprId(e));
+    return tensorExps[e];
+  }
+  const LatPoint &lat(LatPointId p) const {
+    assert(isValidLatPointId(p));
+    return latPoints[p];
+  }
+  ArrayRef<LatPointId> set(LatSetId s) const {
+    assert(isValidLatSetId(s));
+    return latSets[s];
+  }
 
   /// Checks whether the given expression has an associated value.
-  bool hasExprValue(ExprId e) const {
-    return static_cast<bool>(tensorExps[e].val);
-  }
+  bool hasExprValue(ExprId e) const { return static_cast<bool>(exp(e).val); }
 
   /// Sets the expression to have the associated value.  Asserts that
   /// the new value is defined, and that the expression does not already
   /// have a value.  If you want to overwrite a previous associated value,
   /// use `updateExprValue` instead.
   void setExprValue(ExprId e, Value v) {
+    assert(isValidExprId(e));
     assert(v && "Got an undefined value");
     auto &val = tensorExps[e].val;
     assert(!val && "Expression already has an associated value");
     val = v;
   }
 
   /// Clears the value associated with the expression.  Asserts that the
   /// expression does indeed have an associated value before clearing it.
   /// If you don't want to check for a previous associated value first,
   /// then use `updateExprValue` instead.
   void clearExprValue(ExprId e) {
+    assert(isValidExprId(e));
     auto &val = tensorExps[e].val;
     assert(val && "Expression does not have an associated value to clear");
     val = Value();
   }
 
   /// Unilaterally updates the expression to have the associated value.
   /// That is, unlike `setExprValue` and `clearExprValue`, this method
   /// does not perform any checks on whether the expression had a
   /// previously associated value nor whether the new value is defined.
   //
   // TODO: The unilateral update semantics are required by the
   // current implementation of `CodegenEnv::genLoopBoundary`; however,
   // that implementation seems a bit dubious.  We would much rather have
   // the semantics `{ clearExprValue(e); setExprValue(e, v); }` or
   // `{ clearExprValue(e); if (v) setExprValue(e, v); }` since those
   // provide better invariants.
-  void updateExprValue(ExprId e, Value v) { tensorExps[e].val = v; }
+  void updateExprValue(ExprId e, Value v) {
+    assert(isValidExprId(e));
+    tensorExps[e].val = v;
+  }
 
 #ifndef NDEBUG
   /// Print methods (for debugging).
   void dumpExp(ExprId e) const;
   void dumpLat(LatPointId p) const;
   void dumpSet(LatSetId s) const;
   void dumpBits(const BitVector &bits) const;
 #endif
 
   /// Builds the iteration lattices in a bottom-up traversal given the
   /// remaining tensor (sub)expression and the next loop in the iteration
   /// graph.  Returns the identifier of the root set.
   LatSetId buildLattices(ExprId e, LoopId i);
 
   /// Builds a tensor expression from the given Linalg operation.
   /// On success, returns the identifier of the root expression.
   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 isValidLoopId(LoopId i) const {
+    return i != detail::kInvalidId && i < numLoops;
+  }
+  bool isValidLevel(TensorId t, Level lvl) const {
+    return isValidTensorId(t) && lvl < lvlToLoop[t].size();
+  }
+  bool isValidExprId(ExprId e) const {
+    return e != detail::kInvalidId && e < tensorExps.size();
+  }
+  bool isValidLatPointId(LatPointId p) const {
+    return p != detail::kInvalidId && p < latPoints.size();
+  }
+  bool isValidLatSetId(LatSetId s) const {
+    return s != detail::kInvalidId && s < latSets.size();
+  }
   bool maybeZero(ExprId e) const;
-  bool isInvariant(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;

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

public:
   genLoopBoundary(function_ref<
                   std::optional<Operation *>(MutableArrayRef<Value> parameters)>
                       callback);
 
   //
   // Merger delegates.
   //
 
+  constexpr TensorId makeTensorId(unsigned t) const {
+    return latticeMerger.makeTensorId(t);
+  }
+  constexpr LoopId makeLoopId(unsigned i) const {
+    return latticeMerger.makeLoopId(i);
+  }
+  constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const {
+    return latticeMerger.makeTensorLoopId(t, i);
+  }
   const TensorExp &exp(ExprId e) const { return latticeMerger.exp(e); }
   const LatPoint &lat(LatPointId l) const { return latticeMerger.lat(l); }
   ArrayRef<LatPointId> set(LatSetId s) const { return latticeMerger.set(s); }
   DimLevelType dlt(TensorId t, LoopId i) const {
     return latticeMerger.getDimLevelType(t, i);
   }
   DimLevelType dlt(TensorLoopId b) const {
     return latticeMerger.getDimLevelType(b);

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

void CodegenEnv::startEmit() {
   // addition is assoicative, the order how we evaluate the expression does
   // not matter. However, to support multiplication, the order of the loop
   // index should match the evaluation order to the affine expression AST.
 
   // Initialize loop emitter.
   SmallVector<Value> tensors; // input tensors passed to loop emitter
   for (OpOperand &t : linalgOp->getOpOperands()) {
     tensors.push_back(t.get());
-    Level rank = linalgOp.getMatchingIndexingMap(&t).getNumResults();
-    for (Level lvl = 0; lvl < rank; lvl++) {
-      sortArrayBasedOnOrder(
-          latticeMerger.getDependentLoops(t.getOperandNumber(), lvl), topSort);
-    }
+    const TensorId tid = makeTensorId(t.getOperandNumber());
+    const Level lvlRank = linalgOp.getMatchingIndexingMap(&t).getNumResults();
+    const auto enc = getSparseTensorEncoding(t.get().getType());
+    (void)enc;
+    assert(!enc || lvlRank == enc.getLvlRank());
+    for (Level lvl = 0; lvl < lvlRank; lvl++)
+      sortArrayBasedOnOrder(latticeMerger.getDependentLoops(tid, lvl), topSort);
   }
 
   loopEmitter.initialize(
       tensors,
       StringAttr::get(linalgOp.getContext(),
                       linalg::GenericOp::getOperationName()),
       /*hasOutput=*/true,
       /*isSparseOut=*/sparseOut != nullptr, topSort,
for (utils::IteratorType it : linalgOp.getIteratorTypesArray()) {
     if (it == utils::IteratorType::reduction) {
       if (latticeMerger.hasNegateOnOut(exp))
         return false;
       break;
     }
   }
 
   OpOperand *lhs = linalgOp.getDpsInitOperand(0);
-  // That the operand number is a valid `TensorId` will be verified
-  // by the call to `isSingleCondition` below; though we may want to add
-  // assertions to check it here, in order to give better error messages.
-  const TensorId tensor = lhs->getOperandNumber();
+  const TensorId tensor = makeTensorId(lhs->getOperandNumber());
   // An non-annotated output tensor is assumed dense, and becomes a random
   // access n-dim memref. Admissible since insertions cannot occur.
   if (getSparseTensorType(lhs->get()).isAllDense())
     return true;
 
   // A tensor expression with a sparse output tensor that changes its values
   // but not its nonzero structure, an operation called "simply dynamic" in
   // [Bik96,Ch9], is also admissible without special env.

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

private:
 
   /// Generates a predicate to determine whether the tranformed coordinates are
   /// 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);
 
-  TensorId getNumTensors() const { return tensors.size(); }
+  unsigned getNumTensors() const { return tensors.size(); }
 
   bool isOutputTensor(TensorId tid) const {
-    return hasOutput && tid == static_cast<TensorId>(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();

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

 void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
                              bool isSparseOut, ArrayRef<LoopId> topSort,
                              DependentLvlGetter dimGetter) {
   // First initialize the top-level type of the fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
 
-  const TensorId numTensors = ts.size();
+  const unsigned numTensors = ts.size();
   this->tensors.assign(ts.begin(), ts.end());
   this->lvlTypes.assign(numTensors, std::vector<DimLevelType>());
   this->lvlSizes.assign(numTensors, std::vector<Value>());
   this->highs.assign(numTensors, std::vector<Value>());
   this->segHi.assign(numTensors, std::vector<Value>());
   this->posits.assign(numTensors, std::vector<Value>());
   this->coords.assign(numTensors, std::vector<Value>());
   this->positionsBuffers.assign(numTensors, std::vector<Value>());

 
 Value LoopEmitter::genAffine(OpBuilder &builder, Location loc, AffineExpr a) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     // FIXME: since the one callsite in Sparsification passes in a
     // level-expression, the `getPosition` must in fact be a `Dimension`.
     // However, elsewhere we have been lead to expect that `loopIdToOrd`
     // should be indexed by `LoopId`...
-    const LoopId i = a.cast<AffineDimExpr>().getPosition();
-    return loopStack[loopIdToOrd[i]].iv;
+    const auto loopId = a.cast<AffineDimExpr>().getPosition();
+    assert(loopId < loopIdToOrd.size());
+    return loopStack[loopIdToOrd[loopId]].iv;
   }
   case AffineExprKind::Add: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
     return builder.create<arith::AddIOp>(
         loc, genAffine(builder, loc, binOp.getLHS()),
         genAffine(builder, loc, binOp.getRHS()));
   }
   case AffineExprKind::Mul: {
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // 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)) {
-    unsigned tid = t; // Why `t` can not be captured by lambda?
+    const TensorId tid = t; // Why `t` can not be captured by lambda?
     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()));
       for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         if (!isUniqueDLT(lvlTypes[tid][reassoc[i]])) {
           // Links the SSA chain for segHi.
           segHi[tid][reassoc[i]] = after->getArgument(o++);
if (isSingletonDLT(lvlTp)) {
       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 theSegHi = segHi[tid][srcLvl - 1];
-      highs[tid][srcLvl] = (!isUniqueDLT(lvlTypes[tid][srcLvl - 1]) && theSegHi)
-                               ? theSegHi
-                               : builder.create<arith::AddIOp>(loc, pLo, c1);
+      const auto parentSegHi = segHi[tid][srcLvl - 1];
+      highs[tid][srcLvl] =
+          (!isUniqueDLT(lvlTypes[tid][srcLvl - 1]) && parentSegHi)
+              ? parentSegHi
       return;
     }
   }
 
   llvm_unreachable("Unrecognized level-type!");
 }
 
 void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder,

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;
-      loopEmitter.enterNewLoopSeq(rewriter, loc, 0, ld);
+      const SmallVector<TensorId, 1> tids{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, 0, l, reduc);
+      loopEmitter.enterLoopOverTensorAtLvl(rewriter, loc, tids, l, reduc);
     }
 
     SmallVector<Value> lcvs;
     lcvs.reserve(lvlRank);
     loopEmitter.getLoopIVs(lcvs);
 
     if (op.getOrder()) {
       // FIXME: There is some dim/lvl confusion here since `dimRank != lvlRank`

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

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

 /// 3rd, when a constant operand is used in the non-trivial index expression.
 ///
 /// TODO: constant should be easy to handle.
 static bool findDepIdxSet(Merger &merger, TensorId tensor, Level lvl,
                           AffineExpr a, DimLevelType dlt,
                           bool isSubExp = false) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
-    LoopId ldx = a.cast<AffineDimExpr>().getPosition();
+    const LoopId ldx = merger.makeLoopId(a.cast<AffineDimExpr>().getPosition());
     if (!isUndefDLT(merger.getDimLevelType(tensor, ldx)))
       return false; // used more than once, e.g., A[i][i]
 
     // TODO: Generalizes the following two cases. A[i] (with trivial index
     // expression) can be treated as a special affine index expression. We do
     // not necessarily need to differentiate them.
     if (!isSubExp)
       merger.setLevelAndType(tensor, ldx, lvl, dlt);

 /// gradually removed from the codebase.
 /// When using dependent index reducton-based approach, it currently only
 /// supports affine addition index expression.
 static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased) {
   bool annotated = false;
   // `filterLdx` may be mutated by `findAffine`.
   LoopId filterLdx = env.merger().getStartingFilterLoopId();
   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();
     assert(!enc || lvlRank == enc.getLvlRank());
     assert(static_cast<Level>(env.op().getRank(&t)) == lvlRank);
 
     // We only need to do index reduction if there is at least one non-trivial
     // index expression on sparse levels.
     // If all non-trivial index expression is on dense levels, we can
     // efficiently rely on the random access to locate the element.
     bool needIdxReduc =
         enc && getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) != 0;
     // If then current tensor being inspected requires affine index, it need
     // to be sliced.
     for (Level l = 0; l < lvlRank; l++) {
-      const TensorId tid = t.getOperandNumber();
-      AffineExpr a = map.getResult(toOrigDim(enc, l));
-      DimLevelType dlt = enc.getLvlType(l);
+      // FIXME: `toOrigDim` is deprecated.
+      const AffineExpr a = map.getResult(toOrigDim(enc, l));
+      const DimLevelType dlt = enc.getLvlType(l);
       if (idxReducBased && needIdxReduc) {
         if (!findDepIdxSet(env.merger(), tid, l, a, dlt))
           return false; // inadmissible affine expression
       } else {
         if (!findAffine(env.merger(), tid, l, a, dlt, filterLdx))
           return false; // inadmissible affine expression
       }
     }
   }
   assert(filterLdx == env.merger().getNumLoops());
   return annotated;
 }
 
 /// A helper to compute a topological sort. O(n^2) time complexity
 /// as we use adj matrix for the graph.
 /// The sorted result will put the first Reduction iterator to the
-/// latest possible index.
-/// FIXME(wrengr): correct the above "index"
+/// latest possible `LoopOrd`.
 ///
 /// The `inDegree` is indexed by `LoopId`, and the `adjM` is indexed by
 /// `(LoopId,LoopId)`.
-static bool topSortOptimal(CodegenEnv &env, LoopId n,
+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 (LoopId i = 0; i < n; i++) {
+  const LoopId numLoops = env.merger().getNumLoops();
+  for (LoopId i = 0; i < numLoops; i++) {
     if (inDegree[i] == 0) {
       if (env.merger().isFilterLoop(i))
         filterIt.push_back(i);
       else if (linalg::isReductionIterator(iteratorTypes[i]))
         redIt.push_back(i);
       else
         parIt.push_back(i);
     }
while (!redIt.empty() || !parIt.empty() || !filterIt.empty()) {
     //    if (xxx)
     //      for (1 to M)
     //        O(X) computation  => O(NK+NMX) time complexity
     auto &it = !filterIt.empty() ? filterIt : (!parIt.empty() ? parIt : redIt);
     auto src = it.back();
     env.topSortPushBack(src);
     it.pop_back();
     // Update in-degree, and push 0-degree node into worklist.
-    for (LoopId dst = 0; dst < n; dst++) {
+    for (LoopId dst = 0; dst < numLoops; dst++) {
       if (adjM[src][dst] && --inDegree[dst] == 0) {
         if (env.merger().isFilterLoop(dst))
           filterIt.push_back(dst);
         else if (linalg::isReductionIterator(iteratorTypes[dst]))
           redIt.push_back(dst);
         else
           parIt.push_back(dst);
       }
     }
   }
-  return env.topSortSize() == n;
+  return env.topSortSize() == numLoops;
 }
 
 /// Helper method to add all constraints from the indices in one affine
 /// expression before all indices in the other affine expression. For
 /// example i0+i1 < i2+i3+1 yields i0<i2, i0<i3, i1<i2, and i1<i3.
 /// The affine expression `a` is empty iff `fidx` have a value, leading to
 /// b = (i0 + i1) < fidx => i0 < fidx, i1 < fidx.
 /// The affine expression `b` is empty iff `tidx` have a value, leading to
if (!adjM[f][t]) {
       inDegree[t]++;
     }
     return;
   }
   // Picks an affine expression and expand (recurse into) it.
   const auto toExpand = a ? a : b;
   switch (toExpand.getKind()) {
   case AffineExprKind::DimId: {
-    std::optional<LoopId> idx = toExpand.cast<AffineDimExpr>().getPosition();
+    const std::optional<LoopId> idx{
+        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) {
     }
   }
 }
 
 static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
                                           OpOperand *skip, SortMask mask,
                                           std::vector<std::vector<bool>> &adjM,
                                           std::vector<unsigned> &inDegree) {
-  // Get map and encoding.
-  auto map = env.op().getMatchingIndexingMap(&t);
-  auto enc = getSparseTensorEncoding(t.get().getType());
+  // Get map, encoding, and tensor-identifier.
+  const auto map = env.op().getMatchingIndexingMap(&t);
+  const auto enc = getSparseTensorEncoding(t.get().getType());
+  const TensorId tid = env.makeTensorId(t.getOperandNumber());
 
   // Each tensor expression and optional dimension ordering (row-major
   // by default) puts an ordering constraint on the loop indices. For
   // example, the tensor expresion A_ijk forces the ordering i < j < k
   // on the loop indices if no explicit dimension ordering is given.
-  for (Level l = 0, rank = map.getNumResults(); l < rank; l++) {
-    AffineExpr ta = map.getResult(toOrigDim(enc, l));
-    std::optional<LoopId> tldx =
-        env.merger().getLoopId(t.getOperandNumber(), l);
+  const Level lvlRank = map.getNumResults();
+  assert(!enc || lvlRank == enc.getLvlRank());
+  for (Level lvl = 0; lvl < lvlRank; lvl++) {
+    // FIXME: `toOrigDim` is deprecated.
+    AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
+    std::optional<LoopId> tldx = env.merger().getLoopId(tid, lvl);
     // Filter loops should be constructed after all the dependent loops,
     // i.e., d0 + d1 < filter_loop(d0 + d1)
     if (tldx && env.merger().isFilterLoop(*tldx)) {
-      assert(!ta.isa<AffineDimExpr>() && !isDenseDLT(enc.getDimLevelType()[l]));
+      assert(!ta.isa<AffineDimExpr>() &&
+             !isDenseDLT(enc.getDimLevelType()[lvl]));
       addAffineOrderings(adjM, inDegree, ta, AffineExpr(), std::nullopt, tldx);
       // Now that the ordering of affine expression is captured by filter
       // loop idx, we only need to ensure the affine ordering against filter
       // loop. Thus, we reset the affine express to nil here to mark it as
       // resolved.
       ta = AffineExpr();
     }
 
     // Skip tensor during cycle resolution, though order between filter loop
     // and dependent loops need to be guaranteed unconditionally.
     if (&t == skip)
       continue;
 
-    if (l > 0) {
-      AffineExpr fa = map.getResult(toOrigDim(enc, l - 1));
-      std::optional<LoopId> fldx =
-          env.merger().getLoopId(t.getOperandNumber(), l - 1);
+    if (lvl > 0) {
+      // FIXME: `toOrigDim` is deprecated.
+      AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
+      std::optional<LoopId> fldx = env.merger().getLoopId(tid, lvl - 1);
 
       // Applying order constraints on every pair of dimExpr between two
       // compound affine expressions can sometime too strict:
       // E.g, for [dense, dense] -> (d0 + d1, d2 + d3).
       // It is totally fine to have loop sequence d0->d2->d1->d3 instead of
       // requiring d0 < d2, d1 < d2, d0 < d3, d1 < d3.
       // We also relax the affine constraint when use slice-based algorithm
       // as there is no filter loop for affine index on sparse dimension.
static void addFilterLoopBasedConstraints(CodegenEnv &env, OpOperand &t,
   }
 }
 
 static void addSliceBasedConstraints(CodegenEnv &env, OpOperand &t,
                                      OpOperand *skip, SortMask mask,
                                      std::vector<std::vector<bool>> &adjM,
                                      std::vector<unsigned> &inDegree) {
   // Get map and encoding.
-  auto map = env.op().getMatchingIndexingMap(&t);
-  auto enc = getSparseTensorEncoding(t.get().getType());
+  const auto map = env.op().getMatchingIndexingMap(&t);
+  const auto enc = getSparseTensorEncoding(t.get().getType());
 
   // No special treatment for simple indices.
   if (getNumNonTrivialIdxExpOnSparseLvls(map, t.get()) == 0)
     return addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
 
   // Skip tensor during cycle resolution, though order between filter loop
   // and dependent loops need to be guaranteed unconditionally.
   if (&t == skip)
     return;
 
   AffineDimFinder finder(env.op());
   finder.setPickedIterType(utils::IteratorType::reduction);
   // To compute iteration graph for tensor[d0 + d1 + d3, d4 + d5 + d6],
   // we requires there exist d_x \in {d0, d1, d3} and d_y \in {d4, d5, d6},
   // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
-  for (Level lvl = 1, rank = map.getNumResults(); lvl < rank; lvl++) {
-    AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
-    AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
+  const Level lvlRank = map.getNumResults();
+  assert(!enc || lvlRank == enc.getLvlRank());
+  for (Level lvl = 1; lvl < lvlRank; lvl++) {
+    // FIXME: `toOrigDim` is deprecated.
+    const AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
+    const AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
 
     // This is a heuristic, we pick an abitrary reduction loop from lhs and
     // rhs and use them as d_x and d_y.
     finder.walkPostOrder(fa);
-    AffineDimExpr fexp = finder.getDimExpr();
-    LoopId fldx = fexp.getPosition();
+    const AffineDimExpr fexp = finder.getDimExpr();
+    const LoopId fldx = env.makeLoopId(fexp.getPosition());
 
     finder.walkPostOrder(ta);
-    AffineDimExpr texp = finder.getDimExpr();
-    LoopId tldx = texp.getPosition();
+    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]++;
     }
 
     AffineDimCollector fCollector;
     fCollector.walkPostOrder(fa);
     AffineDimCollector tCollector;
     tCollector.walkPostOrder(ta);
 
     // make sure dx and dy is the last;
     for (auto fd : fCollector.dims) {
-      LoopId f = fd.getPosition();
+      const LoopId f = env.makeLoopId(fd.getPosition());
       if (f == fldx)
         continue;
       if (!adjM[f][fldx]) {
         adjM[f][fldx] = true;
         inDegree[fldx]++;
       }
     }
     for (auto td : tCollector.dims) {
-      LoopId t = td.getPosition();
+      const LoopId t = env.makeLoopId(td.getPosition());
       if (t == tldx)
         continue;
       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) {
-      LoopId f = fd.getPosition();
+      const LoopId f = env.makeLoopId(fd.getPosition());
       if (f == fldx) // skip d_x
         continue;
 
       for (auto td : tCollector.dims) {
-        LoopId t = td.getPosition();
+        const LoopId t = env.makeLoopId(td.getPosition());
         if (t == tldx) // skip d_y
           continue;
         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
 /// essential for sparse storage formats since these only support access
 /// along fixed dimensions. Even for dense storage formats, however, the natural
 /// index order yields innermost unit-stride access with better spatial
 /// locality.
 static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
                                   OpOperand *skip, bool idxReducBased = false) {
   // Set up an n x n from/to adjacency matrix of the iteration graph
   // for the implicit loop indices i_0 .. i_n-1.
-  const LoopId n = env.merger().getNumLoops();
-  std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
-  std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
+  const unsigned numLoops = env.merger().getNumLoops();
+  std::vector<std::vector<bool>> adjM(numLoops,
+                                      std::vector<bool>(numLoops, false));
+  std::vector<unsigned> inDegree(numLoops, 0); // in-degree of each node.
   const auto iteratorTypes = env.op().getIteratorTypesArray();
   // Iterate over the indexing maps of every tensor in the tensor expression.
   for (OpOperand &t : env.op()->getOpOperands()) {
     // Get map and encoding.
     const auto enc = getSparseTensorEncoding(t.get().getType());
     assert(env.op().getMatchingIndexingMap(&t).getNumDims() +
                getNumNonTrivialIdxExpOnSparseLvls(env.op()) ==
-           n);
+           numLoops);
 
     // Skips dense inputs/outputs when not requested.
     const bool isDenseInput = !enc && env.op().isDpsInput(&t);
     const bool isDenseOutput = !enc && !isDenseInput;
     if ((isDenseInput && !includesDenseInput(mask)) ||
         (isDenseOutput && !includesDenseOutput(mask)))
       continue;
 
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     // TODO: Do we really need this?
     if (includesUndef(mask)) {
-      const TensorId tensor = t.getOperandNumber();
-      for (LoopId i = 0; i < n; i++) {
-        const auto dltI = env.dlt(tensor, i);
+      const TensorId tid = env.makeTensorId(t.getOperandNumber());
+      for (LoopId i = 0; i < numLoops; i++) {
+        const auto dltI = env.dlt(tid, i);
         if (isCompressedDLT(dltI) || isSingletonDLT(dltI)) {
-          for (LoopId j = 0; j < n; j++)
-            if (isUndefDLT(env.dlt(tensor, j))) {
+          for (LoopId j = 0; j < numLoops; j++)
+            if (isUndefDLT(env.dlt(tid, 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.
     if (idxReducBased)
       addSliceBasedConstraints(env, t, skip, mask, adjM, inDegree);
     else
       addFilterLoopBasedConstraints(env, t, skip, mask, adjM, inDegree);
   }
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
-  env.topSortClear(n);
-  return topSortOptimal(env, n, iteratorTypes, inDegree, adjM);
+  env.topSortClear(numLoops);
+  return topSortOptimal(env, iteratorTypes, inDegree, adjM);
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
 /// Local bufferization of all dense and sparse data structures.
 static void genBuffers(CodegenEnv &env, OpBuilder &builder) {
env.emitter().initializeLoopEmit(
       });
 }
 
 /// Generates index for load/store on sparse tensor.
 // FIXME: It's not entirely clear what "index" means here (i.e., is it
 // a "coordinate", or "Ldx", or what).  So the function should be renamed
 // and/or the documentation expanded in order to clarify.
 static Value genIndex(CodegenEnv &env, OpOperand *t) {
-  auto map = env.op().getMatchingIndexingMap(t);
+  const auto map = env.op().getMatchingIndexingMap(t);
   const auto stt = getSparseTensorType(t->get());
   const Level lvlRank = stt.getLvlRank();
   assert(static_cast<Level>(map.getNumResults()) == lvlRank);
   // 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.
-  AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), lvlRank - 1));
+  const AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), lvlRank - 1));
   assert(a.getKind() == AffineExprKind::DimId);
-  const LoopId idx = a.cast<AffineDimExpr>().getPosition();
+  const LoopId idx = env.makeLoopId(a.cast<AffineDimExpr>().getPosition());
   return env.getLoopVar(idx);
 }
 
 /// Generates subscript for load/store on a dense or sparse tensor.
 static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
                           SmallVectorImpl<Value> &args) {
   const Location loc = env.op().getLoc();
-  const TensorId tid = t->getOperandNumber();
+  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();
     assert(pos);
     args.push_back(pos);
   } else {

 /// analysis has verified that all computations are "local" to the inlined
 /// branch or otherwise invariantly defined outside the loop nest, with the
 /// exception of index computations, which need to be relinked to actual
 /// inlined cloned code.
 static Value relinkBranch(CodegenEnv &env, RewriterBase &rewriter, Block *block,
                           Value e, LoopId ldx) {
   if (Operation *def = e.getDefiningOp()) {
     if (auto indexOp = dyn_cast<linalg::IndexOp>(def))
-      return env.getLoopVar(indexOp.getDim());
+      return env.getLoopVar(env.makeLoopId(indexOp.getDim()));
     if (def->getBlock() == block) {
       for (unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
         rewriter.updateRootInPlace(def, [&]() {
           def->setOperand(
               i, relinkBranch(env, rewriter, block, def->getOperand(i), ldx));
         });
       }
     }
static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
                           LoopId ldx, bool atStart) {
   if (exp == ::mlir::sparse_tensor::detail::kInvalidId)
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
     // Inspect tensor indices.
     bool isAtLoop = ldx == ::mlir::sparse_tensor::detail::kInvalidId;
     linalg::GenericOp op = env.op();
     OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
-    auto map = op.getMatchingIndexingMap(&t);
+    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.
-      AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), l));
-      const auto sldx = env.merger().getLoopId(t.getOperandNumber(), l);
+      const AffineExpr a = map.getResult(toOrigDim(stt.getEncoding(), l));
+      const auto sldx =
+          env.merger().getLoopId(env.makeTensorId(t.getOperandNumber()), l);
       if (sldx && env.merger().isFilterLoop(*sldx)) {
         if (!env.getLoopVar(*sldx))
           // The filter loops has not been constructed.
           return;
         if (*sldx == ldx)
           isAtLoop = true;
       } else if (!isInvariantAffine(env, a, ldx, isAtLoop))
         return; // still in play
while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
       builder.setInsertionPointAfter(ifOp);
     }
   }
   builder.setInsertionPointToEnd(&whileOp.getAfter().front());
 }
 
 /// Generates a single if-statement within a while-loop.
 static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
-                       const BitVector &conditions) {
+                       LatPointId p) {
   Location loc = env.op().getLoc();
   SmallVector<Type> types;
   Value cond;
-  for (TensorLoopId b = 0, be = conditions.size(); b < be; b++) {
-    if (!conditions[b])
-      continue;
-    const TensorId tid = env.merger().tensor(b);
+  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;
-    const auto dlt = env.dlt(b);
         if (isCompressedDLT(dlt) || isSingletonDLT(dlt)) {
-      const Level lvl = *env.merger().getLvl(tid, ldx);
-      const Value crd = env.emitter().getCoords()[tid][lvl];
+          assert(lvl.has_value());
+          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);
+          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;
-  }
+      });
   if (env.isReduc()) {
     types.push_back(env.getReduc().getType());
     if (env.getValidLexInsert())
       types.push_back(env.getValidLexInsert().getType());
   }
   if (env.isExpand())
     types.push_back(builder.getIndexType());
   if (env.getInsertionChain())
static void genConstantDenseAddressFromLevel(CodegenEnv &env,
   // TODO: Handle affine expression on output tensor.
   linalg::GenericOp op = env.op();
   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 = input->getOperandNumber();
+    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
     for (Level l = startLvl; l < lvlRank; l++) {
       // FIXME: `toOrigDim` is deprecated.
       AffineExpr lvlExpr = lvlExprs[toOrigDim(enc, l)];
       if (enc.isDenseLvl(l) && lvlExpr.isa<AffineConstantExpr>())
         env.emitter().genDenseAffineAddress(builder, loc, tid, l, lvlExpr);
static bool translateBitsToTidLvlPairs(
   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.test(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
static bool translateBitsToTidLvlPairs(
   assert(numloopCond > 0);
   // If we just need to one loop conditions and the conditions is not imposed on
   // non-unique level, the loop can be generated by a for loop.
   return numloopCond == 1 && !hasNonUnique;
 }
 
 /// Starts a single loop in current sequence.
 static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
-                                              OpBuilder &builder, unsigned at,
-                                              unsigned li, bool needsUniv) {
+                                              OpBuilder &builder, LoopOrd at,
+                                              LatPointId li, bool needsUniv) {
   // The set of tensors + lvls to generate loops on
   SmallVector<TensorId> tids, affineTids;
   SmallVector<Level> lvls, affineLvls;
   // The set of dense tensors with non-trivial affine expression that just
   // becomes invariant and the address shall now be generated at the current
   // level.
   SmallVector<AffineExpr> affines;
   bool isSingleCond = translateBitsToTidLvlPairs(
for (unsigned i = 0; i < lsize; i++) {
     // because the loop body causes data-movement which invalidates the
     // iterator.
     for (unsigned j = 0; j < lsize; j++) {
       const LatPointId lj = env.set(lts)[j];
       const ExprId ej = env.lat(lj).exp;
       if (li == lj || env.merger().latGT(li, lj)) {
         // Recurse into body of each branch.
         if (!isSingleCond) {
-          scf::IfOp ifOp = genIf(env, rewriter, idx, env.lat(lj).simple);
+          scf::IfOp ifOp = genIf(env, rewriter, idx, lj);
           genStmt(env, rewriter, ej, at + 1);
           endIf(env, rewriter, ifOp, loop, redInput, cntInput, insInput);
         } else {
           genStmt(env, rewriter, ej, at + 1);
         }
       }
     }
 

 
 private:
   // Last resort cycle resolution.
   LogicalResult resolveCycle(CodegenEnv &env, PatternRewriter &rewriter) const {
     // Compute topological sort while leaving out every
     // sparse input tensor in succession until an acylic
     // iteration graph results.
     for (OpOperand *t : env.op().getDpsInputOperands()) {
-      const TensorId tid = t->getOperandNumber();
+      const TensorId tid = env.makeTensorId(t->getOperandNumber());
       Value tval = t->get();
       auto srcEnc = getSparseTensorEncoding(tval.getType());
       if (!srcEnc || !computeIterationGraph(env, SortMask::kSparseOnly, t))
         continue;
       // Found an input tensor that resolves the cycle by inserting a
       // conversion into a sparse tensor that adheres to the iteration
       // graph order. Also releases the temporary sparse tensor.
       //

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

static ExpArity getExpArity(TensorExp::Kind k) {
   }
   llvm_unreachable("unexpected kind");
 }
 
 //===----------------------------------------------------------------------===//
 // Constructors.
 //===----------------------------------------------------------------------===//
 
-TensorExp::TensorExp(Kind k, unsigned x, ExprId y, Value v, Operation *o)
+TensorExp::TensorExp(TensorExp::Kind k, unsigned x, ExprId y, Value v,
+                     Operation *o)
     : kind(k), val(v), op(o) {
   switch (kind) {
   // Leaf.
   case TensorExp::Kind::kTensor:
     assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     tensor = x;
     return;
   case TensorExp::Kind::kInvariant:
case TensorExp::Kind::kReduce:
     assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
     return;
   }
   llvm_unreachable("unexpected kind");
 }
 
-LatPoint::LatPoint(const BitVector &bits, ExprId e) : bits(bits), exp(e) {}
-
-LatPoint::LatPoint(unsigned numTensors, unsigned numLoops, TensorId t, LoopId i,
-                   ExprId e)
-    : bits(numLoops * numTensors, false), exp(e) {
-  assert(t < numTensors && i < numLoops);
-  const TensorLoopId b = numTensors * i + t;
-  bits.set(b);
-}
-
 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,
                 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>>(
                                           maxLvlRank, std::vector<LoopId>())),
       loopBounds(numLoops, std::make_pair(numTensors, numLoops)) {}
 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
-ExprId Merger::addExp(TensorExp::Kind k, unsigned x, ExprId y, Value v,
-                      Operation *op) {
-  const ExprId e = tensorExps.size();
-  assert((k != TensorExp::Kind::kTensor || x < numTensors) &&
-         (k != TensorExp::Kind::kLoopVar || x < numLoops));
-  tensorExps.emplace_back(k, x, y, v, op);
-  return e;
+ExprId Merger::addTensorExp(TensorId t) {
+  assert(isValidTensorId(t));
+  const ExprId eNew(tensorExps.size());
+  tensorExps.emplace_back(TensorExp::Kind::kTensor, t, detail::kInvalidId,
+                          Value(), nullptr);
+  return eNew;
+}
+
+ExprId Merger::addLoopVarExp(LoopId i) {
+  assert(isValidLoopId(i));
+  const ExprId eNew(tensorExps.size());
+  tensorExps.emplace_back(TensorExp::Kind::kLoopVar, i, detail::kInvalidId,
+                          Value(), nullptr);
+  return eNew;
+}
+
+ExprId Merger::addInvariantExp(Value v) {
+  const ExprId eNew(tensorExps.size());
+  tensorExps.emplace_back(TensorExp::Kind::kInvariant, detail::kInvalidId,
+                          detail::kInvalidId, v, nullptr);
+  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, 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, detail::kInvalidId, v, op);
+  return eNew;
 }
 
 LatPointId Merger::addLat(TensorId t, LoopId i, ExprId e) {
-  assert(t < numTensors && i < numLoops);
-  const LatPointId p = latPoints.size();
-  latPoints.emplace_back(numTensors, numLoops, t, i, e);
-  return p;
+  const LatPointId pNew(latPoints.size());
+  const unsigned size = numLoops * numTensors;
+  const TensorLoopId b = makeTensorLoopId(t, i);
+  latPoints.emplace_back(size, e);
+  latPoints[pNew].bits.set(b);
+  return pNew;
 }
 
 LatPointId Merger::addLat(const BitVector &bits, ExprId e) {
   assert(bits.size() == numLoops * numTensors);
-  const LatPointId p = latPoints.size();
+  const LatPointId pNew(latPoints.size());
   latPoints.emplace_back(bits, e);
-  return p;
+  return pNew;
 }
 
 LatSetId Merger::addSet() {
-  const LatSetId s = latSets.size();
+  const LatSetId sNew(latSets.size());
   latSets.emplace_back();
-  return s;
+  return sNew;
 }
 
 LatPointId Merger::conjLat(TensorExp::Kind kind, LatPointId p0, LatPointId p1,
                            Operation *op) {
-  const LatPointId p = latPoints.size();
-  BitVector bits(latPoints[p0].bits);
-  bits |= latPoints[p1].bits;
-  const ExprId e =
-      addExp(kind, latPoints[p0].exp, latPoints[p1].exp, Value(), op);
+  const LatPointId pNew(latPoints.size());
+  const auto &point0 = lat(p0);
+  const auto &point1 = lat(p1);
+  BitVector bits(point0.bits);
+  bits |= point1.bits;
+  const ExprId e = addExp(kind, point0.exp, point1.exp, op);
   latPoints.emplace_back(bits, e);
-  return p;
+  return pNew;
 }
 
 LatSetId Merger::conjSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                          Operation *op) {
-  const LatSetId s = addSet();
-  for (const LatPointId p0 : latSets[s0])
-    for (const LatPointId p1 : latSets[s1])
-      latSets[s].push_back(conjLat(kind, p0, p1, op));
-  return s;
+  const LatSetId sNew = addSet();
+  auto &setNew = latSets[sNew];
+  for (const LatPointId p0 : set(s0))
+    for (const LatPointId p1 : set(s1))
+      setNew.push_back(conjLat(kind, p0, p1, op));
+  return sNew;
 }
 
 LatSetId Merger::disjSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                          Operation *op) {
-  const LatSetId s = conjSet(kind, s0, s1, op);
+  const LatSetId sNew = conjSet(kind, s0, s1, op);
   // Followed by all in s0.
-  for (const LatPointId p : latSets[s0])
-    latSets[s].push_back(p);
+  latSets[sNew].append(latSets[s0]);
   // Map binary 0-y to unary -y.
   // TODO: move this if-else logic into buildLattices
   if (kind == TensorExp::Kind::kSubF)
     s1 = mapSet(TensorExp::Kind::kNegF, s1);
   else if (kind == TensorExp::Kind::kSubC)
     s1 = mapSet(TensorExp::Kind::kNegC, s1);
   else if (kind == TensorExp::Kind::kSubI)
     s1 = mapSet(TensorExp::Kind::kNegI, s1);
   // Followed by all in s1.
-  for (const LatPointId p : latSets[s1])
-    latSets[s].push_back(p);
-  return s;
+  latSets[sNew].append(latSets[s1]);
+  return sNew;
 }
 
 LatSetId Merger::combiSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                           Operation *orig, bool includeLeft,
                           TensorExp::Kind ltrans, Operation *opleft,
                           bool includeRight, TensorExp::Kind rtrans,
                           Operation *opright) {
-  const LatSetId s = conjSet(kind, s0, s1, orig);
+  const LatSetId sNew = conjSet(kind, s0, s1, orig);
   // Left Region.
   if (includeLeft) {
     if (opleft)
       s0 = mapSet(ltrans, s0, Value(), opleft);
-    for (const LatPointId p : latSets[s0])
-      latSets[s].push_back(p);
+    latSets[sNew].append(latSets[s0]);
   }
   // Right Region.
   if (includeRight) {
     if (opright)
       s1 = mapSet(rtrans, s1, Value(), opright);
-    for (const LatPointId p : latSets[s1])
-      latSets[s].push_back(p);
+    latSets[sNew].append(latSets[s1]);
   }
-  return s;
+  return sNew;
 }
 
 LatSetId Merger::mapSet(TensorExp::Kind kind, LatSetId s0, Value v,
                         Operation *op) {
   assert(TensorExp::Kind::kAbsF <= kind && kind <= TensorExp::Kind::kSelect);
-  const LatSetId s = addSet();
-  for (const LatPointId p : latSets[s0]) {
-    const ExprId e = addExp(kind, latPoints[p].exp, v, op);
-    latSets[s].push_back(addLat(latPoints[p].bits, e));
+  const LatSetId sNew = addSet();
+  auto &setNew = latSets[sNew];
+  for (const LatPointId p : set(s0)) {
+    const auto &point = latPoints[p];
+    setNew.push_back(addLat(point.bits, addExp(kind, point.exp, v, op)));
   }
-  return s;
+  return sNew;
 }
 
 LatSetId Merger::optimizeSet(LatSetId s0) {
-  const LatSetId s = addSet();
-  assert(!latSets[s0].empty());
-  const LatPointId p0 = latSets[s0][0];
-  for (const LatPointId p1 : latSets[s0]) {
+  const LatSetId sNew = addSet();
+  auto &setNew = latSets[sNew];
+  const auto &set0 = set(s0);
+  assert(!set0.empty());
+  const LatPointId p0 = set0[0];
+  for (const LatPointId p1 : set0) {
     bool add = true;
     if (p0 != p1) {
       // Check whether this is a straightforward copy.
-      const ExprId e = latPoints[p1].exp;
-      if (expIsTensor(e, outTensor))
+      if (expIsTensor(latPoints[p1].exp, outTensor))
         continue;
       // Check whether this conjunction is already covered.
-      for (const LatPointId p2 : latSets[s]) {
+      for (const LatPointId p2 : setNew) {
         assert(!latGT(p1, p2)); // Lj => Li would be bad
         if (onlyDenseDiff(p2, p1)) {
           add = false;
           break;
         }
       }
       assert(!add || latGT(p0, p1));
     }
     if (add)
-      latSets[s].push_back(p1);
+      setNew.push_back(p1);
   }
-  for (const LatPointId p : latSets[s])
-    latPoints[p].simple = simplifyCond(s, p);
-  return s;
+  for (const LatPointId p : setNew)
+    latPoints[p].simple = simplifyCond(sNew, p);
+  return sNew;
 }
 
 BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
   // First determine if this lattice point is a *singleton*, i.e.,
   // the last point in a lattice, no other is less than this one.
   bool isSingleton = true;
-  for (const LatPointId p1 : latSets[s0]) {
+  for (const LatPointId p1 : set(s0)) {
     if (p0 != p1 && latGT(p0, p1)) {
       isSingleton = false;
       break;
     }
   }
 
-  BitVector simple(latPoints[p0].bits);
+  BitVector simple(lat(p0).bits);
   bool reset =
       isSingleton && (hasAnySparse(simple) || hasSparseIdxReduction(simple));
-  const TensorLoopId be = simple.size();
-  TensorLoopId offset = 0; // relative to the end
+  // `be`, `b`, and `offset` are `TensorLoopId` in spirit; but we avoid
+  // using that class in this function because we need to do a bunch of
+  // arithmetic on them, so using the newtype would introduce too much
+  // boilerplate.
+  const unsigned be = simple.size();
+  unsigned offset = 0; // relative to the end
   if (!reset)
     // Starts resetting from a dense level, so that the first bit (if kept)
     // is not undefined level-type.
-    for (TensorLoopId b = 0; b < be; b++) {
-      if (simple[b] && isDenseDLT(getDimLevelType(b))) {
+    for (unsigned b = 0; b < be; b++) {
+      if (simple[b] && isDenseDLT(getDimLevelType(TensorLoopId{b}))) {
         offset = be - b - 1; // relative to the end
         break;
       }
     }
 
   // Now apply the two basic rules. We also iterate the bits reversely to always
   // keep the rightmost bit (which could possibly be a synthetic tensor).
-  for (TensorLoopId b = be - 1 - offset, i = 0; i < be;
+  for (unsigned b = be - 1 - offset, i = 0; i < be;
        b = b == 0 ? be - 1 : b - 1, i++) {
     // FIXME: better name? also slice on dense level has locate property as
     // well. Handle it correctly!
-    if (simple[b] && !isLvlWithNonTrivialIdxExp(b)) {
-      const auto dlt = getDimLevelType(b);
+    if (simple[b] && !isLvlWithNonTrivialIdxExp(TensorLoopId{b})) {
+      const auto dlt = getDimLevelType(TensorLoopId{b});
       if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt)) {
         if (reset)
           simple.reset(b);
         reset = true;
       }
     }
   }
   return simple;
 }
 
 bool Merger::latGT(LatPointId i, LatPointId j) const {
-  const BitVector &bitsi = latPoints[i].bits;
-  const BitVector &bitsj = latPoints[j].bits;
+  const BitVector &bitsi = lat(i).bits;
+  const BitVector &bitsj = lat(j).bits;
   assert(bitsi.size() == bitsj.size());
   if (bitsi.count() > bitsj.count()) {
     for (TensorLoopId b = 0, be = bitsj.size(); b < be; b++)
       if (bitsj[b] && !bitsi[b])
         return false;
     return true;
   }
   return false;
 }
 
 bool Merger::onlyDenseDiff(LatPointId i, LatPointId j) const {
-  BitVector tmp(latPoints[j].bits);
-  tmp ^= latPoints[i].bits;
+  BitVector tmp(lat(j).bits);
+  tmp ^= lat(i).bits;
   return !hasAnySparse(tmp) && !hasSparseIdxReduction(tmp);
 }
 
 bool Merger::expContainsTensor(ExprId e, TensorId t) const {
-  if (tensorExps[e].kind == TensorExp::Kind::kTensor)
+