[Mlir-commits] [mlir] [mlir][Func] Use getMutableSuccessorOperands() in FuncOp verifier (PR #184589)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When verifying return-like terminators, use getMutableSuccessorOperands() instead of getNumOperands() so that only the operands passed to the parent region are checked against the function result types. This handles terminators that implement RegionBranchTerminatorOpInterface and carry additional operands for other successor regions (e.g. loop back-edges).

Add tests using test.loop_block_term, which has both an iter operand (passed back to the region) and an exit operand (passed to the parent).

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


3 Files Affected:

- (modified) mlir/lib/Dialect/Func/IR/FuncOps.cpp (+13-14) 
- (modified) mlir/test/Dialect/Func/invalid.mlir (+22) 
- (modified) mlir/test/IR/test-region-branch-op-verifier.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index 8ee2c9956d2d6..2749ad5262f2b 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -298,21 +298,20 @@ LogicalResult FuncOp::verify() {
     auto returnOp = dyn_cast<RegionBranchTerminatorOpInterface>(&block.back());
     if (!returnOp)
       continue;
-
-    if (returnOp->getNumOperands() != resultTypes.size())
+    auto operands =
+        returnOp.getMutableSuccessorOperands(RegionSuccessor::parent());
+    if (operands.size() != resultTypes.size())
       return returnOp->emitOpError("has ")
-             << returnOp->getNumOperands()
-             << " operands, but enclosing function (@" << getName()
-             << ") returns " << resultTypes.size();
-
-    for (auto [i, opType] :
-         llvm::enumerate(llvm::zip(returnOp->getOperandTypes(), resultTypes))) {
-      auto [opTy, resTy] = opType;
-      if (opTy != resTy)
-        return returnOp->emitError()
-               << "type of return operand " << i << " (" << opTy
-               << ") doesn't match function result type (" << resTy
-               << ") in function @" << getName();
+             << operands.size() << " operands, but enclosing function (@"
+             << getName() << ") returns " << resultTypes.size();
+
+    for (auto [i, opType] : llvm::enumerate(llvm::zip(operands, resultTypes))) {
+      auto [operand, resTy] = opType;
+      if (operand.get().getType() != resTy)
+        return returnOp->emitError() << "type of return operand " << i << " ("
+                                     << operand.get().getType()
+                                     << ") doesn't match function result type ("
+                                     << resTy << ") in function @" << getName();
     }
   }
 
diff --git a/mlir/test/Dialect/Func/invalid.mlir b/mlir/test/Dialect/Func/invalid.mlir
index a30aa448a1b09..2620d284c1d06 100644
--- a/mlir/test/Dialect/Func/invalid.mlir
+++ b/mlir/test/Dialect/Func/invalid.mlir
@@ -191,3 +191,25 @@ func.func private @invalid_func_result_attr() -> (i1 {non_dialect_attr = 10})
 // -----
 
 func.func @foo() {} // expected-error {{expected non-empty function body}}
+
+// -----
+
+// test.loop_block_term implements RegionBranchTerminatorOpInterface.
+// getMutableSuccessorOperands(parent) returns only the exit operand (f32).
+// The function returns i32, so the type check must fail.
+func.func @region_branch_term_type_mismatch(%arg: i32) -> i32 {
+  %0 = "test.constant"() { value = 5.3 : f32 } : () -> f32
+  // expected-error @+1 {{type of return operand 0 ('f32') doesn't match function result type ('i32') in function @region_branch_term_type_mismatch}}
+  test.loop_block_term iter %arg exit %0
+}
+
+// -----
+
+// test.loop_block_term has one exit operand (f32) but the function returns
+// nothing. getMutableSuccessorOperands(parent) returns 1 operand while the
+// function has 0 results, so the count check must fail.
+func.func @region_branch_term_count_mismatch(%arg: i32) {
+  %0 = "test.constant"() { value = 5.3 : f32 } : () -> f32
+  // expected-error @+1 {{'test.loop_block_term' op has 1 operands, but enclosing function (@region_branch_term_count_mismatch) returns 0}}
+  test.loop_block_term iter %arg exit %0
+}
diff --git a/mlir/test/IR/test-region-branch-op-verifier.mlir b/mlir/test/IR/test-region-branch-op-verifier.mlir
index b94f6beb9796f..1c7a87c4f0356 100644
--- a/mlir/test/IR/test-region-branch-op-verifier.mlir
+++ b/mlir/test/IR/test-region-branch-op-verifier.mlir
@@ -21,3 +21,17 @@ func.func @test_no_terminator(%arg: index) {
   }
   return
 }
+
+// -----
+
+// test.loop_block_term has two operands: iter (i32, passed back to the region)
+// and exit (f32, passed to the parent). getMutableSuccessorOperands(parent)
+// returns only the exit operand. The function returns f32, matching the exit
+// operand type, so verification must succeed.
+//
+// A verifier using getNumOperands() instead would incorrectly report "has 2
+// operands, but enclosing function returns 1".
+func.func @func_with_region_branch_terminator(%arg: i32) -> f32 {
+  %0 = "test.constant"() { value = 5.3 : f32 } : () -> f32
+  test.loop_block_term iter %arg exit %0
+}

``````````

</details>


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