[mlir][sparse] moving kInvalidId into "detail" namespace

[wrengr]

Summary

In the next few commits I will be converting the various Merger identifier typedefs into newtypes; and once that's done, the kInvalidId constant will only be used internally and therefore does not need to be part of the public mlir::sparse_tensor namespace.

Depends On D146673

[aartbik]

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

1114

are you deliberately enforcing the ::mlir top level instead of relying on the using mlir::sparse_tensor with the shorter version of this?
(I assume yes, but just checking)

[wrengr]

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

1114

yeah it's deliberate, just for paranoia's sake really. I'm pretty sure just using sparse_tensor::detail::kInvalidId would work too. But in the next few CLs this'll be replaced with if (!e.isValid()) so I'm not woo worried about the verbiage in the interim

[wrengr]

rebase

[wrengr]

arc diff $(git merge-base HEAD origin) --update D146674

[wrengr]

rebase

[wrengr]

rebase

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

 /// and serve as unique identifiers for the corresponding `LatPoint` object.
 using LatPointId = unsigned;
 
 /// `LatSet` identifiers.  These are allocated by `Merger::addSet` (and
 /// by other methods calling that one), and serve as unique identifiers
 /// for the corresponding `SmallVector<LatPointId>` object.
 using LatSetId = unsigned;
 
+namespace detail {
 /// A constant serving as the canonically invalid identifier, regardless
 /// of the identifier type.
 static constexpr unsigned kInvalidId = -1u;
+} // namespace detail
 
-/// Children subexpressions of tensor operations.
-struct Children {
+/// Tensor expression. Represents an MLIR expression in tensor index notation.
+struct TensorExp final {
+  enum class Kind;
+
+  /// Child subexpressions for non-leaf expressions.
+  struct Children final {
     ExprId e0;
     ExprId e1;
   };
 
-/// Tensor expression. Represents a MLIR expression in tensor index notation.
-struct TensorExp {
-  enum class Kind;
-
   // 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
enum class TensorExp::Kind {
   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 {
+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 `(t,i)` as the only `TensorLoopId`,
   /// where `(t,i) < (numTensors,numLoops)`.
   LatPoint(unsigned numTensors, unsigned numLoops, TensorId t, LoopId i,
            ExprId e);
 
public:
   // gave the number of input tensors, instead of the current number of
   // input+output tensors.
   Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
          unsigned numFilterLoops);
 
   /// 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 = kInvalidId,
+  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, kInvalidId, v, op);
+    return addExp(k, e, detail::kInvalidId, v, op);
   }
   ExprId addExp(TensorExp::Kind k, Value v, Operation *op = nullptr) {
-    return addExp(k, kInvalidId, kInvalidId, v, op);
+    return addExp(k, detail::kInvalidId, detail::kInvalidId, v, op);
   }
 
   /// 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"
   /// of `LoopId` (effectively constructing a larger "intersection" of those
   /// loops) with a newly constructed tensor (sub)expression of given kind.
   /// Returns the identifier of the new lattice point.
public:
 
   /// 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();
   }
 
   /// Convenience getters to immediately access the stored nodes.
-  /// Typically it is inadvisible to keep the reference around, as in
-  /// `TensorExpr &te = merger.exp(e)`, since insertions into the merger
-  /// may cause data movement and invalidate the underlying memory address.
-  TensorExp &exp(ExprId e) { return tensorExps[e]; }
-  LatPoint &lat(LatPointId p) { return latPoints[p]; }
-  SmallVector<LatPointId> &set(LatSetId s) { return latSets[s]; }
+  /// 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]; }
+
+  /// Checks whether the given expression has an associated value.
+  bool hasExprValue(ExprId e) const {
+    return static_cast<bool>(tensorExps[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(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) {
+    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; }
 
 #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

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

public:
   genLoopBoundary(function_ref<
                   std::optional<Operation *>(MutableArrayRef<Value> parameters)>
                       callback);
 
   //
   // Merger delegates.
   //
 
-  TensorExp &exp(ExprId e) { return latticeMerger.exp(e); }
-  LatPoint &lat(LatPointId l) { return latticeMerger.lat(l); }
-  SmallVector<LatPointId> &set(LatSetId s) { return latticeMerger.set(s); }
+  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);
   }
 
   //
public:
   Value getExpandCount() const { return expCount; }
   void endExpand();
 
   //
   // Reduction methods.
   //
 
   void startReduc(ExprId exp, Value val);
-  bool isReduc() const { return redExp != kInvalidId; }
+  bool isReduc() const { return redExp != detail::kInvalidId; }
   void updateReduc(Value val);
   Value getReduc() const { return redVal; }
   Value endReduc();
   void setValidLexInsert(Value val);
   void clearValidLexInsert();
   Value getValidLexInsert() const { return redValidLexInsert; }
 
   void startCustomReduc(ExprId exp);
-  bool isCustomReduc() const { return redCustom != kInvalidId; }
+  bool isCustomReduc() const { return redCustom != detail::kInvalidId; }
   Value getCustomRedId();
   void endCustomReduc();
 
 private:
   // Linalg operation.
   linalg::GenericOp linalgOp;
 
   // Sparsification options.

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

 //===----------------------------------------------------------------------===//
 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
                        unsigned numFilterLoops)
     : linalgOp(linop), sparseOptions(opts),
       latticeMerger(numTensors, numLoops, numFilterLoops), loopEmitter(),
       topSort(), sparseOut(nullptr), outerParNest(-1u), insChain(), expValues(),
-      expFilled(), expAdded(), expCount(), redVal(), redExp(kInvalidId),
-      redCustom(kInvalidId), redValidLexInsert() {}
+      expFilled(), expAdded(), expCount(), redVal(), redExp(detail::kInvalidId),
+      redCustom(detail::kInvalidId), redValidLexInsert() {}
 
 LogicalResult CodegenEnv::initTensorExp() {
   // Builds the tensor expression for the Linalg operation in SSA form.
   std::optional<ExprId> optExp = latticeMerger.buildTensorExpFromLinalg(op());
   if (!optExp || !isAdmissibleTensorExp(*optExp))
     return failure();
 
   tensorExp = *optExp;
std::optional<Operation *> CodegenEnv::genLoopBoundary(
   }
   if (isExpand())
     params.push_back(expCount);
   if (insChain != nullptr)
     params.push_back(insChain);
   auto r = callback(params); // may update parameters
   unsigned i = 0;
   if (isReduc()) {
+    // FIXME: This requires `updateExprValue` to perform updates without
+    // checking for a previous value; but it's not clear whether that's
+    // by design or might be a potential source for bugs.
     updateReduc(params[i++]);
     if (redValidLexInsert)
       setValidLexInsert(params[i++]);
   }
   if (isExpand())
     updateExpandCount(params[i++]);
   if (insChain != nullptr)
     updateInsertionChain(params[i]);
void CodegenEnv::endExpand() {
   expValues = expFilled = expAdded = expCount = Value();
 }
 
 //===----------------------------------------------------------------------===//
 // Code generation environment reduction methods
 //===----------------------------------------------------------------------===//
 
 void CodegenEnv::startReduc(ExprId exp, Value val) {
-  assert(!isReduc() && exp != kInvalidId);
+  assert(!isReduc() && exp != detail::kInvalidId);
   redExp = exp;
   updateReduc(val);
 }
 
 void CodegenEnv::updateReduc(Value val) {
   assert(isReduc());
-  redVal = exp(redExp).val = val;
+  redVal = val;
+  // NOTE: `genLoopBoundary` requires that this performs a unilateral
+  // update without checking for a previous value first.  (It's not
+  // clear whether any other callsites also require that.)
+  latticeMerger.updateExprValue(redExp, val);
 }
 
 Value CodegenEnv::endReduc() {
+  assert(isReduc());
   Value val = redVal;
-  updateReduc(Value());
-  redExp = kInvalidId;
+  redVal = val;
+  latticeMerger.clearExprValue(redExp);
+  redExp = detail::kInvalidId;
   return val;
 }
 
 void CodegenEnv::setValidLexInsert(Value val) {
   assert(isReduc() && val);
   redValidLexInsert = val;
 }
 
 void CodegenEnv::clearValidLexInsert() {
   assert(!isReduc());
   redValidLexInsert = Value();
 }
 
 void CodegenEnv::startCustomReduc(ExprId exp) {
-  assert(!isCustomReduc() && exp != kInvalidId);
+  assert(!isCustomReduc() && exp != detail::kInvalidId);
   redCustom = exp;
 }
 
 Value CodegenEnv::getCustomRedId() {
   assert(isCustomReduc());
   return dyn_cast<sparse_tensor::ReduceOp>(exp(redCustom).op).getIdentity();
 }
 
 void CodegenEnv::endCustomReduc() {
   assert(isCustomReduc());
-  redCustom = kInvalidId;
+  redCustom = detail::kInvalidId;
 }

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

if (!rhs) {
       Value chain = env.getInsertionChain();
       scf::IfOp ifOp =
           builder.create<scf::IfOp>(loc, chain.getType(), rhs, /*else=*/true);
       builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
       // Existing value was preserved to be used here.
       assert(env.exp(exp).val);
       Value v0 = env.exp(exp).val;
       genInsertionStore(env, builder, t, v0);
-      env.exp(exp).val = Value();
+      env.merger().clearExprValue(exp);
       // Yield modified insertion chain along true branch.
       Value mchain = env.getInsertionChain();
       builder.create<scf::YieldOp>(op.getLoc(), mchain);
       // Yield original insertion chain along false branch.
       builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
       builder.create<scf::YieldOp>(loc, chain);
       // Done with if statement.
       env.updateInsertionChain(ifOp->getResult(0));

 }
 
 /// Recursively generates tensor expression.
 static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e,
                     LoopId ldx) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
 
-  if (e == kInvalidId)
+  if (e == ::mlir::sparse_tensor::detail::kInvalidId)
     return Value();
   const TensorExp &exp = env.exp(e);
   const auto kind = exp.kind;
   if (kind == TensorExp::Kind::kTensor)
     return genTensorLoad(env, rewriter, e);
   if (kind == TensorExp::Kind::kInvariant)
     return genInvariantValue(env, e);
   if (kind == TensorExp::Kind::kLoopVar)
if (ee &&
       (kind == TensorExp::Kind::kUnary || kind == TensorExp::Kind::kBinary ||
        kind == TensorExp::Kind::kBinaryBranch ||
        kind == TensorExp::Kind::kReduce || kind == TensorExp::Kind::kSelect))
     ee = relinkBranch(env, rewriter, ee.getParentBlock(), ee, ldx);
 
   if (kind == TensorExp::Kind::kReduce)
     env.endCustomReduc(); // exit custom
 
-  if (kind == TensorExp::Kind::kSelect) {
-    assert(!exp.val);
-    env.exp(e).val = v0; // Preserve value for later use.
-  }
+  if (kind == TensorExp::Kind::kSelect)
+    env.merger().setExprValue(e, v0); // Preserve value for later use.
 
   return ee;
 }
 
 /// 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 == kInvalidId)
+  if (exp == ::mlir::sparse_tensor::detail::kInvalidId)
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
     // Inspect tensor indices.
-    bool isAtLoop = ldx == kInvalidId;
+    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 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.
if (lhs == &t) {
         if (env.hasSparseOutput())
           env.setValidLexInsert(constantI1(builder, env.op().getLoc(), false));
       } else {
         genTensorStore(env, builder, exp, env.endReduc());
         env.clearValidLexInsert();
       }
     } else {
       // Start or end loop invariant hoisting of a tensor load.
-      env.exp(exp).val = atStart ? genTensorLoad(env, builder, exp) : Value();
+      if (atStart)
+        env.merger().setExprValue(exp, genTensorLoad(env, builder, exp));
+      else
+        env.merger().clearExprValue(exp);
     }
   } else if (env.exp(exp).kind != TensorExp::Kind::kInvariant &&
              env.exp(exp).kind != TensorExp::Kind::kLoopVar) {
     // Traverse into the binary operations. Note that we only hoist
     // tensor loads, since subsequent MLIR/LLVM passes know how to
     // deal with all other kinds of derived loop invariants.
     if (env.exp(exp).kind == TensorExp::Kind::kReduce)
       env.startCustomReduc(exp); // enter custom
if (isFor) {
     bool isInner = at == env.topSortSize() - 1;
     return genFor(env, builder, isOuter, isInner, idx, tids, lvls);
   }
   return genWhile(env, builder, idx, needsUniv, tids, lvls);
 }
 
 /// Generates the induction structure for a while-loop.
 static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder, LoopId idx,
-                            bool needsUniv, BitVector &induction,
-                            scf::WhileOp whileOp) {
+                            bool needsUniv, scf::WhileOp whileOp) {
   Location loc = env.op().getLoc();
   // Finalize each else branch of all if statements.
   if (env.isReduc() || env.isExpand() || env.getInsertionChain()) {
     while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
                builder.getInsertionBlock()->getParentOp())) {
       // Break on IfOp for slicing filtering.
       if (ifOp->getAttr(LoopEmitter::getLoopEmitterLoopAttrName()) ==
           StringAttr::get(ifOp->getContext(), "slice"))
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,
-                       BitVector &conditions) {
+                       const BitVector &conditions) {
   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);
     assert(ldx == env.merger().loop(b));
env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
     }
   });
 
   env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tids, lvls);
 
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
-    unsigned lsize = env.set(lts).size();
-    for (unsigned i = 1; i < lsize; i++) {
-      const LatPointId li = env.set(lts)[i];
+    for (const LatPointId li : env.set(lts).drop_front())
       if (!env.merger().hasAnySparse(env.lat(li).simple) &&
           !env.merger().hasSparseIdxReduction(env.lat(li).simple))
         return true;
-    }
-  }
   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();
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
   return std::make_pair(loop, isSingleCond);
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
 static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
                     LoopId idx, LatPointId li, bool needsUniv) {
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
-    finalizeWhileOp(env, rewriter, idx, needsUniv, env.lat(li).bits, whileOp);
+    finalizeWhileOp(env, rewriter, idx, needsUniv, whileOp);
   } else if (auto forOp = dyn_cast<scf::ForOp>(loop)) {
     // Any iteration of a reduction for-loop creates a valid lex insert.
     if (env.isReduc() && env.getValidLexInsert())
       env.setValidLexInsert(constantI1(rewriter, env.op().getLoc(), true));
   } else {
     needsUniv = false;
   }
 
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 ? kInvalidId : env.topSortAt(at - 1);
+  const LoopId ldx = at == 0 ? ::mlir::sparse_tensor::detail::kInvalidId
+                             : 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.
-  unsigned lsize = env.set(lts).size();
+  //
+  // NOTE: We cannot change this to `for (const LatPointId li : env.set(lts))`
+  // because the loop body causes data-movement which invalidates
+  // the iterator.
+  const unsigned lsize = env.set(lts).size();
   for (unsigned i = 0; i < lsize; i++) {
-    // Start a loop.
     const LatPointId li = env.set(lts)[i];
+    // Start a loop.
     auto [loop, isSingleCond] = startLoop(env, rewriter, at, li, needsUniv);
 
     // Visit all lattices points with Li >= Lj to generate the
     // loop-body, possibly with if statements for coiteration.
     Value redInput = env.getReduc();
     Value cntInput = env.getExpandCount();
     Value insInput = env.getInsertionChain();
+    // NOTE: We cannot change this to `for (const LatPointId lj : env.set(lts))`
+    // 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);
           genStmt(env, rewriter, ej, at + 1);

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

 // Constructors.
 //===----------------------------------------------------------------------===//
 
 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 != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     tensor = x;
-    break;
+    return;
   case TensorExp::Kind::kInvariant:
-    assert(x == kInvalidId && y == kInvalidId && v && !o);
-    break;
+    assert(x == detail::kInvalidId && y == detail::kInvalidId && v && !o);
+    return;
   case TensorExp::Kind::kLoopVar:
-    assert(x != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     loop = x;
-    break;
+    return;
   // 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 != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    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 != kInvalidId && y == kInvalidId && v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && v && !o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    return;
   case TensorExp::Kind::kBinaryBranch:
   case TensorExp::Kind::kSelect:
-    assert(x != kInvalidId && y == kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    return;
   case TensorExp::Kind::kUnary:
     // No assertion on y can be made, as the branching paths involve both
     // a unary (`mapSet`) and binary (`disjSet`) pathway.
-    assert(x != kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    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 != kInvalidId && y != kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && !o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    return;
   case TensorExp::Kind::kBinary:
   case TensorExp::Kind::kReduce:
-    assert(x != kInvalidId && y != kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
-    break;
+    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);

 
 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;
 }
 
+LatPointId Merger::addLat(const BitVector &bits, ExprId e) {
+  assert(bits.size() == numLoops * numTensors);
+  const LatPointId p = latPoints.size();
+  latPoints.emplace_back(bits, e);
+  return p;
+}
+
 LatSetId Merger::addSet() {
   const LatSetId s = latSets.size();
   latSets.emplace_back();
   return s;
 }
 
 LatPointId Merger::conjLat(TensorExp::Kind kind, LatPointId p0, LatPointId p1,
                            Operation *op) {

 }
 
 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);
-    latPoints.emplace_back(latPoints[p].bits, e);
-    latSets[s].push_back(latPoints.size() - 1);
+    latSets[s].push_back(addLat(latPoints[p].bits, e));
   }
   return s;
 }
 
 LatSetId Merger::optimizeSet(LatSetId s0) {
   const LatSetId s = addSet();
   assert(!latSets[s0].empty());
   const LatPointId p0 = latSets[s0][0];