[mlir][sparse] use "current" and "curr" consistently

[aartbik]

Removes at in favor of curr; also makes method delegates consistent

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Aart Bik (aartbik)

Changes

Removes at in favor of curr; also makes method delegates consistent

---

Patch is 26.05 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/74656.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h (+1-1)
• (modified) mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp (+137-138)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
index 27dc407b493dc..a1947f48393ef 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
@@ -108,7 +108,7 @@ class CodegenEnv {
     return loopEmitter.unpackTensorLevelRange(std::forward<ContainerTy>(c));
   }
 
-  unsigned getLoopDepth() const { return loopEmitter.getCurrentDepth(); }
+  unsigned getCurrentDepth() const { return loopEmitter.getCurrentDepth(); }
 
   //
   // Code generation environment verify functions.
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 6637a26d0e5af..6c9adf9fa21a0 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -45,22 +45,22 @@ using namespace mlir::sparse_tensor;
 //===----------------------------------------------------------------------===//
 
 /// Returns true iff affine expression is invariant. Sets the
-/// parameter `isAtLoop` when expression just became invariant.
-static bool isInvariantAffine(AffineExpr a, LoopId at, bool &isAtLoop) {
+/// parameter `isCurrentLoop` when expression just became invariant.
+static bool isInvariantAffine(AffineExpr a, LoopId curr, bool &isCurrentLoop) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     const LoopId i = cast<AffineDimExpr>(a).getPosition();
-    if (i + 1 == at) {
-      isAtLoop = true;
+    if (i + 1 == curr) {
+      isCurrentLoop = true;
       return true; // becomes invariant at current loop
     }
-    return i < at; // invariant when already generated
+    return i < curr; // invariant when already generated
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
     auto binOp = cast<AffineBinaryOpExpr>(a);
-    return isInvariantAffine(binOp.getLHS(), at, isAtLoop) &&
-           isInvariantAffine(binOp.getRHS(), at, isAtLoop);
+    return isInvariantAffine(binOp.getLHS(), curr, isCurrentLoop) &&
+           isInvariantAffine(binOp.getRHS(), curr, isCurrentLoop);
   }
   default: {
     assert(isa<AffineConstantExpr>(a));
@@ -89,9 +89,8 @@ static bool findAffine(Merger &merger, TensorId tid, Level lvl, AffineExpr a,
     assert(isDenseLT(lt));
     if (auto binOp = dyn_cast<AffineBinaryOpExpr>(a)) {
       // We do not set dim level format for affine expression 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.
+      // 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(), lt, false) &&
              findAffine(merger, tid, lvl, binOp.getRHS(), lt, false);
     }
@@ -257,10 +256,10 @@ static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased) {
     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.
+    // 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
@@ -413,9 +412,8 @@ static void genInsertionStore(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
   if (!env.isExpand()) {
     const LoopId numLoops = op.getRank(t);
     // Retrieves the first `numLoop` induction variables.
-    SmallVector<Value> ivs = llvm::to_vector(
-        llvm::drop_end(env.emitter().getLoopIVsRange(),
-                       env.emitter().getCurrentDepth() - numLoops));
+    SmallVector<Value> ivs = llvm::to_vector(llvm::drop_end(
+        env.emitter().getLoopIVsRange(), env.getCurrentDepth() - numLoops));
     Value chain = env.getInsertionChain();
     if (!env.getValidLexInsert()) {
       env.updateInsertionChain(builder.create<InsertOp>(loc, rhs, chain, ivs));
@@ -655,7 +653,7 @@ static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e) {
 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
-                          LoopId at, bool atStart) {
+                          LoopId curr, bool isStart) {
   if (exp == ::mlir::sparse_tensor::detail::kInvalidId)
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
@@ -666,19 +664,19 @@ static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
     const auto stt = getSparseTensorType(t.get());
     const Level lvlRank = stt.getLvlRank();
     assert(static_cast<Level>(map.getNumResults()) == lvlRank);
-    bool isAtLoop = at == 0; // for scalar tensors
+    bool isCurrentLoop = curr == 0; // for scalar tensors
     for (Level l = 0; l < lvlRank; l++) {
       const AffineExpr a = map.getResult(l);
-      if (!isInvariantAffine(a, at, /*out*/ isAtLoop))
+      if (!isInvariantAffine(a, curr, /*out*/ isCurrentLoop))
         return; // still in play
     }
-    // All exhausted at this level (isAtLoop denotes exactly at this LoopId).
-    if (!isAtLoop)
+    // All exhausted at current level.
+    if (!isCurrentLoop)
       return;
     OpOperand *lhs = op.getDpsInitOperand(0);
     if (lhs == &t) {
       // Start or end a scalarized reduction.
-      if (atStart) {
+      if (isStart) {
         Value load = env.isCustomReduc() ? env.getCustomRedId()
                                          : genTensorLoad(env, builder, exp);
         env.startReduc(exp, load);
@@ -690,7 +688,7 @@ static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
       }
     } else {
       // Start or end loop invariant hoisting of a tensor load.
-      if (atStart)
+      if (isStart)
         env.merger().setExprValue(exp, genTensorLoad(env, builder, exp));
       else
         env.merger().clearExprValue(exp);
@@ -705,20 +703,20 @@ static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
       env.startCustomReduc(exp); // enter custom
     const ExprId e0 = env.exp(exp).children.e0;
     const ExprId e1 = env.exp(exp).children.e1;
-    genInvariants(env, builder, e0, at, atStart);
-    genInvariants(env, builder, e1, at, atStart);
+    genInvariants(env, builder, e0, curr, isStart);
+    genInvariants(env, builder, e1, curr, isStart);
     if (env.exp(exp).kind == TensorExp::Kind::kReduce)
       env.endCustomReduc(); // exit custom
   }
 }
 
 /// Generates an expanded access pattern in innermost dimension.
-static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId at,
-                      bool atStart) {
+static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId curr,
+                      bool isStart) {
   linalg::GenericOp op = env.op();
   OpOperand *lhs = op.getDpsInitOperand(0);
-  if (!env.atExpandLevel(lhs, op.getRank(lhs), at))
-    return; // not needed at this level
+  if (!env.atExpandLevel(lhs, op.getRank(lhs), curr))
+    return; // not needed at current level
   assert(!env.isReduc());
   // Generate start or end of an expanded access pattern. Note that because
   // an expansion does not rely on the ongoing contents of the sparse storage
@@ -727,7 +725,7 @@ static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId at,
   // needed, we will need to use the SSA value in the insertion chain instead.
   Value tensor = lhs->get();
   Location loc = op.getLoc();
-  if (atStart) {
+  if (isStart) {
     auto dynShape = {ShapedType::kDynamic};
     Type etp = cast<ShapedType>(tensor.getType()).getElementType();
     Type t1 = MemRefType::get(dynShape, etp);
@@ -740,7 +738,7 @@ static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId at,
                     r.getResult(3));
   } else {
     SmallVector<Value> indices;
-    for (LoopId i = 0; i < at; i++)
+    for (LoopId i = 0; i < curr; i++)
       indices.push_back(env.emitter().getLoopIV(i));
     Value values = env.getExpandValues();
     Value filled = env.getExpandFilled();
@@ -782,18 +780,18 @@ static bool isParallelFor(CodegenEnv &env, bool isOuter, bool isSparse) {
 
 /// Whether or not the current loop being generated should be parallized (if
 /// possible) according to the configuration.
-static bool shouldTryParallize(CodegenEnv &env, LoopId at,
+static bool shouldTryParallize(CodegenEnv &env, LoopId curr,
                                ArrayRef<TensorLevel> tidLvls) {
   linalg::GenericOp op = env.op();
   auto iteratorTypes = op.getIteratorTypesArray();
-  bool isSparse = llvm::any_of(tidLvls, [at, &env](TensorLevel tidLvl) {
+  bool isSparse = llvm::any_of(tidLvls, [curr, &env](TensorLevel tidLvl) {
     // Queries the LT based on the tensor and loop id, as requested by
     // `CodegenEnv::lt(TensorId, LoopId)`. The returned LT from CodegenEnv
     // should be consistent with the LT indexed by <TensorId, Level>.
-    const auto lt = env.lt(env.unpackTensorLevel(tidLvl).first, at);
+    const auto lt = env.lt(env.unpackTensorLevel(tidLvl).first, curr);
     return isCompressedLT(lt) || isSingletonLT(lt);
   });
-  return isParallelFor(env, /*isOuter=*/at == 0, isSparse);
+  return isParallelFor(env, /*isOuter=*/curr == 0, isSparse);
 }
 
 /// Emit a loop to coiterate over the list of tensor levels. The generated loop
@@ -814,9 +812,9 @@ static Operation *genCoIteration(CodegenEnv &env, OpBuilder &builder,
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
-static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopId at,
+static Operation *genLoop(CodegenEnv &env, OpBuilder &builder, LoopId curr,
                           bool needsUniv, ArrayRef<TensorLevel> tidLvls) {
-  bool tryParallel = shouldTryParallize(env, at, tidLvls);
+  bool tryParallel = shouldTryParallize(env, curr, tidLvls);
   return genCoIteration(env, builder, tidLvls, tryParallel, needsUniv);
 }
 
@@ -861,7 +859,7 @@ static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder,
 }
 
 /// Generates a single if-statement within a while-loop.
-static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId at,
+static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId curr,
                        LatPointId p) {
   Location loc = env.op().getLoc();
   SmallVector<Type> types;
@@ -879,13 +877,13 @@ static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId at,
           auto stt = getSparseTensorType(env.op().getInputs()[tid]);
           lt = stt.getLvlType(*lvl);
         }
-        assert(at == env.merger().loop(b));
+        assert(curr == env.merger().loop(b));
         Value clause;
         if (isCompressedLT(lt) || isSingletonLT(lt) ||
             isLooseCompressedLT(lt) || is2OutOf4LT(lt)) {
           assert(lvl.has_value());
           const Value crd = env.emitter().getCoords()[tid][*lvl];
-          const Value lvar = env.getLoopVar(at);
+          const Value lvar = env.getLoopVar(curr);
           clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                  crd, lvar);
         } else {
@@ -942,12 +940,12 @@ static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp,
 /// Starts a loop sequence at given level. Returns true if
 /// the universal loop index must be maintained at this level.
 static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
-                         LoopId at, LatSetId lts) {
-  assert(!env.getLoopVar(at));
+                         LoopId curr, LatSetId lts) {
+  assert(!env.getLoopVar(curr));
   // Emit invariants at this loop sequence level.
-  genInvariants(env, builder, exp, at, /*atStart=*/true);
+  genInvariants(env, builder, exp, curr, /*isStart=*/true);
   // Emit access pattern expansion for sparse tensor output.
-  genExpand(env, builder, at, /*atStart=*/true);
+  genExpand(env, builder, curr, /*isStart=*/true);
   // Emit further intitialization at this loop sequence level.
   const LatPointId l0 = env.set(lts)[0];
   bool needsUniv = false;
@@ -956,13 +954,12 @@ static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
                                            std::optional<Level> lvl,
                                            LevelType lt, bool isIdxReduc) {
-    assert(env.merger().loop(b) == at);
+    assert(env.merger().loop(b) == curr);
     if (isDenseLT(lt) || isUndefLT(lt)) {
       if (tid == env.merger().getSynTensorID()) {
         // Needs loop emitter to set up loop bounds for synthetic tensor too if
         // there is a loop condition imposed on the synthetic tensor.
-        tidLvls.push_back(
-            env.makeTensorLevel(tid, env.emitter().getCurrentDepth()));
+        tidLvls.push_back(env.makeTensorLevel(tid, env.getCurrentDepth()));
       }
       needsUniv = true;
     }
@@ -1025,87 +1022,89 @@ static void genInitConstantDenseAddress(CodegenEnv &env,
 
 /// Return true if the lattices bit can be iterated by a for loop.
 static bool translateBitsToTidLvlPairs(
-    CodegenEnv &env, LatPointId li, LoopId at,
+    CodegenEnv &env, LatPointId li, LoopId curr,
     SmallVectorImpl<TensorLevel> &tidLvls,
     SmallVectorImpl<std::pair<TensorLevel, AffineExpr>> &affineTidLvls) {
   const BitVector &simple = env.lat(li).simple;
   const TensorId outTid = env.merger().getOutTensorID();
-  const std::optional<Level> outLvl = env.merger().getLvl(outTid, at);
+  const std::optional<Level> outLvl = env.merger().getLvl(outTid, curr);
 
   unsigned numloopCond = 0;
   bool hasNonUnique = false;
-  env.merger().foreachTensorLoopId(li, [&, at](TensorLoopId b, TensorId tid,
-                                               std::optional<Level> lvl,
-                                               LevelType lt, bool isIdxReduc) {
-    if (simple[b]) {
-      if (isIdxReduc) {
-        tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
-        numloopCond++;
-        return;
-      }
-      if (isUndefLT(lt)) {
-        // An undefined lt in the lattices, we probably mean to
-        // generate a dense loop according to the synthetic tensor (for
-        // invariants and sparse output tensor).
-        if (env.merger().getSynTensorID() == tid) {
-          // Coiterating with an invariant
-          // e.g., out = prod(in[i][j] op invariant);
-          // or a broadcast
-          // e.g., out[i][j] = in[i] (j is undef for input)
-          //
-          // The level of the synthetic tensor is the current loop depth;
-          // the rank of the synthetic tensor equals to number of loops.
-          lvl = env.emitter().getCurrentDepth();
-        } else if (!lvl) {
-          // Skips invalid lvl (e.g., when this is a zero ranked tensor).
-          return;
-        }
-      }
-      hasNonUnique = !isUniqueLT(lt) || hasNonUnique;
-      tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
-      numloopCond++;
-    } else if (isDenseLT(lt) || isIdxReduc) {
-      tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
-    } else {
-      assert(isUndefLT(lt));
-      linalg::GenericOp op = env.op();
-      if (tid >= op.getNumDpsInputs())
-        // We only handle affine expression on input tensors (for now).
-        return;
-      OpOperand *operand = &op->getOpOperand(tid);
-      const auto stt = getSparseTensorType(operand->get());
-      // Non-annotated dense tensors requires no special handling.
-      if (!stt.hasEncoding())
-        return;
-
-      ArrayRef<AffineExpr> affines =
-          op.getMatchingIndexingMap(operand).getResults();
-      const Level lvlRank = stt.getLvlRank();
-      assert(affines.size() == static_cast<size_t>(lvlRank));
-      for (Level l = 0; l < lvlRank; l++) {
-        AffineExpr exp = affines[l];
-        // Skip simple affine expression and non-dense levels (which
-        // have their own filter loop).
-        if (isa<AffineDimExpr>(exp) || !stt.isDenseLvl(l))
-          continue;
-
-        // Constant affine expression are handled in genLoop.
-        if (!isa<AffineConstantExpr>(exp)) {
-          bool isAtLoop = false;
-          assert(at == env.getLoopDepth());
-          if (isInvariantAffine(exp, at + 1, /*out*/ isAtLoop) && isAtLoop) {
-            // If the compound affine is invariant and we are right at the
-            // level. We need to generate the address according to the
-            // affine expression. This is also the best place we can do it
-            // to avoid putting it inside inner loops.
-            affineTidLvls.emplace_back(env.makeTensorLevel(tid, l), exp);
+  env.merger().foreachTensorLoopId(
+      li, [&, curr](TensorLoopId b, TensorId tid, std::optional<Level> lvl,
+                    LevelType lt, bool isIdxReduc) {
+        if (simple[b]) {
+          if (isIdxReduc) {
+            tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
+            numloopCond++;
+            return;
+          }
+          if (isUndefLT(lt)) {
+            // An undefined lt in the lattices, we probably mean to
+            // generate a dense loop according to the synthetic tensor (for
+            // invariants and sparse output tensor).
+            if (env.merger().getSynTensorID() == tid) {
+              // Coiterating with an invariant
+              // e.g., out = prod(in[i][j] op invariant);
+              // or a broadcast
+              // e.g., out[i][j] = in[i] (j is undef for input)
+              //
+              // The level of the synthetic tensor is the current loop depth;
+              // the rank of the synthetic tensor equals to number of loops.
+              assert(curr == env.getCurrentDepth());
+              lvl = curr;
+            } else if (!lvl) {
+              // Skips invalid lvl (e.g., when this is a zero ranked tensor).
+              return;
+            }
+          }
+          hasNonUnique = !isUniqueLT(lt) || hasNonUnique;
+          tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
+          numloopCond++;
+        } else if (isDenseLT(lt) || isIdxReduc) {
+          tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
+        } else {
+          assert(isUndefLT(lt));
+          linalg::GenericOp op = env.op();
+          if (tid >= op.getNumDpsInputs())
+            // We only handle affine expression on input tensors (for now).
+            return;
+          OpOperand *operand = &op->getOpOperand(tid);
+          const auto stt = getSparseTensorType(operand->get());
+          // Non-annotated dense tensors requires no special handling.
+          if (!stt.hasEncoding())
+            return;
+
+          ArrayRef<AffineExpr> affines =
+              op.getMatchingIndexingMap(operand).getResults();
+          const Level lvlRank = stt.getLvlRank();
+          assert(affines.size() == static_cast<size_t>(lvlRank));
+          for (Level l = 0; l < lvlRank; l++) {
+            AffineExpr exp = affines[l];
+            // Skip simple affine expression and non-dense levels (which
+            // have their own filter loop).
+            if (isa<AffineDimExpr>(exp) || !stt.isDenseLvl(l))
+              continue;
+
+            // Constant affine expression are handled in genLoop.
+            if (!isa<AffineConstantExpr>(exp)) {
+              bool isCurrentLoop = false;
+              assert(curr == env.getCurrentDepth());
+              if (isInvariantAffine(exp, curr + 1, /*out*/ isCurrentLoop) &&
+                  isCurrentLoop) {
+                // If the compound affine is invariant and we are right at the
+                // level. We need to generate the address according to the
+                // affine expression. This is also the best place w...
[truncated]