[mlir][affine] allow iter args as valid dims

[metaflow]

that is effectivevely a revert of
7aabf47 for
mlir/lib/Dialect/Affine/IR/AffineOps.cpp

there are situations when iter args can be used as a dims. For example in
https://github.com/google/heir/blob/main/lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp#L1036

rootExp and batchSize are iter args that are being used as dims and from the point of internal loops
they are fixed.

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Mikhail Goncharov (metaflow)

Changes

that is effectivevely a revert of
7aabf47 for
mlir/lib/Dialect/Affine/IR/AffineOps.cpp

there are situations when iter args can be used as a dims. For example in
https://github.com/google/heir/blob/main/lib/Dialect/Polynomial/Conversions/PolynomialToModArith/PolynomialToModArith.cpp#L1036

rootExp and batchSize are iter args that are being used as dims and from the point of internal loops
they are fixed.

---

Full diff: https://github.com/llvm/llvm-project/pull/139069.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+6-7)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 65f85444e70db..eb23403a68813 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -294,12 +294,10 @@ bool mlir::affine::isValidDim(Value value) {
     return isValidDim(value, getAffineScope(defOp));
 
   // This value has to be a block argument for an op that has the
-  // `AffineScope` trait or an induction var of an affine.for or
-  // affine.parallel.
-  if (isAffineInductionVar(value))
-    return true;
+  // `AffineScope` trait or for an affine.for or affine.parallel.
   auto *parentOp = llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
-  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
+  return parentOp && (parentOp->hasTrait<OpTrait::AffineScope>() ||
+                      isa<AffineForOp, AffineParallelOp>(parentOp));
 }
 
 // Value can be used as a dimension id iff it meets one of the following
@@ -318,9 +316,10 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
 
   auto *op = value.getDefiningOp();
   if (!op) {
-    // This value has to be an induction var for an affine.for or an
+    // This value has to be a block argument for an affine.for or an
     // affine.parallel.
-    return isAffineInductionVar(value);
+    auto *parentOp = llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
+    return isa<AffineForOp, AffineParallelOp>(parentOp);
   }
 
   // Affine apply operation is ok if all of its operands are ok.

[metaflow]

need to update the test, hold on..

[metaflow]

cc @oowekyala please take a look

[metaflow]

alright, I will merge as is to unblock builds, please feel free to follow up

[oowekyala]

@metaflow thanks for fixing this so quickly. I wasn't aware of a use case for this.

[asraa]

Hey, HEIR maintainer here - I don't really know whether what we are doing (updating the upper bound using an operand from the iter_args) counts as an "affine" loop.

In that fastNTT algorithm, all the accessors are affine expressions of the induction var - so by that definition it seems like it might be an affine loop.

But the upper bound's expression uses loop-carried vars as operands and not just induction vars and constants. I don't know whether the loop-carried vars are considered to be regular enough that they can be used as operands to the bound maps.

I was thinking that if this change is valid to make reduce the scope of an affine for loop, I would have to rewrite this at a lower level, which I can do. But I don't know the "right" answer.

[ftynse]

Haven't dug too deep into this, but IMO the original change that got reverted was in fact correct. Since we don't know how secondary iteration args evolve with loop iterations (e.g., the loop may be yielding the result of a function call or the state of a PRNG), we cannot statically represent them as an integer set bound by affine constraints, unlike primary iteration args where the evolution is clear.

[asraa]

I think that's fair - I'm happy to update our code to avoid using affine loops. In the end, this code is destined to lower to LLVM IR, and it really shouldn't matter if I use a lower level abstraction.

[ftynse]

#140808. @asraa please let me know when you are done updating downstream.