[mlir][sparse] Updating the `Merger::{exp,lat,set}` methods to return const

[wrengr]

Summary

This helps the Merger maintain invariants, as well as clarifying the immutability of the underlying objects (with the one exception of TensorExp::val).

Depends On: D146559

[wrengr]

Split D146559 off into a separate commit

[wrengr]

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

133

If any reviewers know the answer to this, please let me know.

[wrengr]

rebase

[wrengr]

rebase

[aartbik]

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

516

I would remove L468-471 as a note to self (unless you have an immediate solution)
mainly because I want to make it very clear that L452-464 apply to all three methods here

550

would it be an option to enforce that update is only used when (1) old value IS set, and (2) new value IS set?
That way, we have a clear set->update->update->clear behavioral chain?

(not saying we should do this, but the question comes to mind)

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

133

Ah, that answers my question above ;-)

Even for the reductions, we should have a clear set->update->update->clear chain
if that is not the case, we must have a second look at this

[wrengr]

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

516

I think your line numbers are off vs the current version (e.g., L468 is the definition of setLoopDependentTensorLevel...). If you mean the FIXME comment attached to the set method, then I can remove it (and stash it among my local notes-to-self). Just let me know if that's the lines you mean

Afaict, the problem occurs because the entire inner vector is unpacked and stored directly within the outer vector, hence whenever the outer vector is resized then it copies the contents of the inner vectors over to the new buffer. So the obvious solution is to instead have the outer vector only store a pointer to the data of the inner vector (perhaps along with other header info like the vector length); so that resizing the outer vector only copies the pointers themselves around, but leaves the data in place. As for how to implement that, I'm not sure if switching the inner vector to std::vector would be sufficient to do that, or not. If not, then we'd have to do some manual code to explicitly allocate things on the heap and free them at the end. Definitely doable, but not the highest priority at the moment

550

Originally I'd been hoping to have this method do (the equivalent of) { clearExprValue(e); setExprValue(e, v); }, but that didn't work out for CodegenEnv::genLoopBoundary... So then I was hoping to have the method do (the equivalent of) { clearExprValue(e); if (v) setExprValue(e, v); }, but that also didn't work for CodegenEnv::genLoopBoundary...

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

133

Yeah. Personally I think the implementation of this method is buggy, precisely because this updateExprValue has to do a unilateral change rather than either of the well-behaved protocols I discussed above (clear+set, or clear+if(v)set).

[aartbik]

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

516

Ah, yes, I meant remove the FIXME, so that the NOTE: above clearly applies to all three

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

133

Yeah, just leave this FIXME and I will have a look later

[wrengr]

Addressing nits, and rebasing

[wrengr]

rebase

[wrengr]

rebase

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

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);
   }
 
   //

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

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::startReduc(ExprId exp, Value val) {
   assert(!isReduc() && exp != 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());
+  redVal = val;
+  latticeMerger.clearExprValue(redExp);
   redExp = kInvalidId;
   return val;
 }
 
 void CodegenEnv::setValidLexInsert(Value val) {
   assert(isReduc() && val);
   redValidLexInsert = val;
 }

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));
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 (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;
   }
 
static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp,
   const LoopId ldx = at == 0 ? 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);