[Mlir-commits] [mlir] [mlir][affine] Add affine.for verifier and move arguments check earlier (PR #206685)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hocky Yudhiono (hockyy)

<details>
<summary>Changes</summary>

Fixes #<!-- -->206628 crash by adding an earlier body argument verifier for `affine.for`. This crash is caused because `LoopLikeOpInterface` verifier would call `getRegionIterArgs()` and assumed induction var of `affine.for` exists. 

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1) 
- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+9-7) 
- (modified) mlir/test/Dialect/Affine/invalid.mlir (+10) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 1e14f9f37288d..57cbbb06b9f65 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -343,6 +343,7 @@ def AffineForOp : Affine_Op<"for",
 
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
+  let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f095500495f18..0186cb6b6ffa5 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2203,19 +2203,21 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
                bodyBuilder);
 }
 
+LogicalResult AffineForOp::verify() {
+  auto *body = getBody();
+  if (body->getNumArguments() == 0 || !getInductionVar().getType().isIndex())
+    return emitOpError("expected body to have an index argument for the "
+                       "induction variable");
+
+  return success();
+}
+
 LogicalResult AffineForOp::verifyRegions() {
   // Step must be a strictly positive integer.
   if (getStepAsInt() <= 0)
     return emitOpError("expected step to be a positive integer, got ")
            << getStepAsInt();
 
-  // Check that the body defines as single block argument for the induction
-  // variable.
-  auto *body = getBody();
-  if (body->getNumArguments() == 0 || !body->getArgument(0).getType().isIndex())
-    return emitOpError("expected body to have a single index argument for the "
-                       "induction variable");
-
   // Verify that the bound operands are valid dimension/symbols.
   /// Lower bound.
   if (getLowerBoundMap().getNumInputs() > 0)
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index e31f2f00f96e2..ae209b4092d98 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -632,3 +632,13 @@ func.func @affine_for_zero_step_verifier() {
   }) : () -> ()
   return
 }
+
+// -----
+
+func.func @affine_for_missing_induction_var() {
+  // expected-error at +1 {{'affine.for' op expected body to have an index argument for the induction variable}}
+  "affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index, upperBoundMap = affine_map<() -> (2)>}> ({
+    "affine.yield"() : () -> ()
+  }) : () -> ()
+  return
+}

``````````

</details>


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