[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.