[Mlir-commits] [mlir] 5f9fd47 - [mlir][affine] allow iter args as valid dims (#139069)

[llvmlistbot at llvm.org]

Author: Mikhail Goncharov
Date: 2025-05-08T16:12:06+02:00
New Revision: 5f9fd475a03363db6da069b6ded1c503833a695c

URL: https://github.com/llvm/llvm-project/commit/5f9fd475a03363db6da069b6ded1c503833a695c
DIFF: https://github.com/llvm/llvm-project/commit/5f9fd475a03363db6da069b6ded1c503833a695c.diff

LOG: [mlir][affine] allow iter args as valid dims (#139069)

that is effectivevely a revert of
7aabf47522625e227433cc9603e0b6858c5dd66d 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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/raise-memref.mlir

Removed: 
    


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

diff  --git a/mlir/test/Dialect/Affine/raise-memref.mlir b/mlir/test/Dialect/Affine/raise-memref.mlir
index 8dc24d99db3e2..98c54e3998b73 100644
--- a/mlir/test/Dialect/Affine/raise-memref.mlir
+++ b/mlir/test/Dialect/Affine/raise-memref.mlir
@@ -112,7 +112,7 @@ func.func @symbols(%N : index) {
 // CHECK:                  %[[lhs5:.*]] = arith.addf %[[lhs]], %[[lhs4]]
 // CHECK:                  %[[lhs6:.*]] = arith.addi %[[a4]], %[[cst1]]
 // CHECK:                  affine.store %[[lhs5]], %{{.*}}[%[[a1]], symbol(%arg0) + 1] :
-// CHECK:                  memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs6]]] :
+// CHECK:                  affine.store %[[lhs5]], %{{.*}}[%[[a1]], %arg4 + 1] :
 // CHECK:                  %[[lhs7:.*]] = "ab.v"
 // CHECK:                  memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs7]]] :
 // CHECK:                  affine.yield %[[lhs6]]