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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+7-6) 
- (modified) mlir/test/Dialect/Affine/raise-memref.mlir (+1-1) 


``````````diff
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]]

``````````

</details>


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