[Mlir-commits] [mlir] Revert "[mlir][affine] allow iter args as valid dims" (PR #140808)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/140808

Reverts llvm/llvm-project#139069: the original change that got reverted by that PR 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.

>From 898669409278a7ae2000cc3cda2a6a7495d0f3af Mon Sep 17 00:00:00 2001
From: "Oleksandr \"Alex\" Zinenko" <ftynse at gmail.com>
Date: Tue, 20 May 2025 23:33:42 +0200
Subject: [PATCH] Revert "[mlir][affine] allow iter args as valid dims
 (#139069)"

This reverts commit 5f9fd475a03363db6da069b6ded1c503833a695c.
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   | 13 +++++++------
 mlir/test/Dialect/Affine/raise-memref.mlir |  2 +-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 4ba24c66e3b82..2364f8957992d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -294,10 +294,12 @@ 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 for an affine.for or affine.parallel.
+  // `AffineScope` trait or an induction var of an affine.for or
+  // affine.parallel.
+  if (isAffineInductionVar(value))
+    return true;
   auto *parentOp = llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
-  return parentOp && (parentOp->hasTrait<OpTrait::AffineScope>() ||
-                      isa<AffineForOp, AffineParallelOp>(parentOp));
+  return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
 }
 
 // Value can be used as a dimension id iff it meets one of the following
@@ -318,10 +320,9 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
 
   auto *op = value.getDefiningOp();
   if (!op) {
-    // This value has to be a block argument for an affine.for or an
+    // This value has to be an induction var for an affine.for or an
     // affine.parallel.
-    auto *parentOp = llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
-    return isa<AffineForOp, AffineParallelOp>(parentOp);
+    return isAffineInductionVar(value);
   }
 
   // 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 98c54e3998b73..8dc24d99db3e2 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:                  affine.store %[[lhs5]], %{{.*}}[%[[a1]], %arg4 + 1] :
+// CHECK:                  memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs6]]] :
 // CHECK:                  %[[lhs7:.*]] = "ab.v"
 // CHECK:                  memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs7]]] :
 // CHECK:                  affine.yield %[[lhs6]]