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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/184589

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

>From a5b98f7594496021f21b882d939c4e768ab4b083 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 03:09:35 -0800
Subject: [PATCH] [mlir][Func] Use getMutableSuccessorOperands() in FuncOp
 verifier

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).
---
 mlir/lib/Dialect/Func/IR/FuncOps.cpp          | 27 +++++++++----------
 mlir/test/Dialect/Func/invalid.mlir           | 22 +++++++++++++++
 .../IR/test-region-branch-op-verifier.mlir    | 14 ++++++++++
 3 files changed, 49 insertions(+), 14 deletions(-)

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
+}