[mlir][func] Fix ReturnOp issue #112146

[fabianmcg]

This patches fixes issue #112146, where an assertion was being triggered by func::ReturnOp::getSuccessorRegions and func::FuncOp not implementing RegionBranchOpInterface.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

Changes

This patches fixes issue #112146, where an assertion was being triggered by func::ReturnOp::getSuccessorRegions and func::FuncOp not implementing RegionBranchOpInterface.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Func/IR/FuncOps.td (+4-2)
• (modified) mlir/lib/Dialect/Func/IR/FuncOps.cpp (+6)

diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index 22efe15aa83a50..8f20ca64602e4c 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -338,8 +338,10 @@ def FuncOp : Func_Op<"func", [
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : Func_Op<"return", [Pure, HasParent<"FuncOp">,
-                                MemRefsNormalizable, ReturnLike, Terminator]> {
+def ReturnOp : Func_Op<"return", [
+    Pure, HasParent<"FuncOp">, MemRefsNormalizable, ReturnLike, Terminator,
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface, ["getSuccessorRegions"]>]
+  > {
   let summary = "Function return operation";
   let description = [{
     The `func.return` operation represents a return operation within a function.
diff --git a/mlir/lib/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index a490b4c3c4ab43..06e0172a0fab6a 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -306,6 +306,12 @@ LogicalResult ReturnOp::verify() {
   return success();
 }
 
+void ReturnOp::getSuccessorRegions(ArrayRef<Attribute> operands,
+                                   SmallVectorImpl<RegionSuccessor> &regions) {
+  // Return control back to func::FuncOp.
+  regions.push_back(RegionSuccessor());
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

[llvmbot]

@llvm/pr-subscribers-mlir-func

Author: Fabian Mora (fabianmcg)

Changes

This patches fixes issue #112146, where an assertion was being triggered by func::ReturnOp::getSuccessorRegions and func::FuncOp not implementing RegionBranchOpInterface.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Func/IR/FuncOps.td (+4-2)
• (modified) mlir/lib/Dialect/Func/IR/FuncOps.cpp (+6)

diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index 22efe15aa83a50..8f20ca64602e4c 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -338,8 +338,10 @@ def FuncOp : Func_Op<"func", [
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : Func_Op<"return", [Pure, HasParent<"FuncOp">,
-                                MemRefsNormalizable, ReturnLike, Terminator]> {
+def ReturnOp : Func_Op<"return", [
+    Pure, HasParent<"FuncOp">, MemRefsNormalizable, ReturnLike, Terminator,
+    DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface, ["getSuccessorRegions"]>]
+  > {
   let summary = "Function return operation";
   let description = [{
     The `func.return` operation represents a return operation within a function.
diff --git a/mlir/lib/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index a490b4c3c4ab43..06e0172a0fab6a 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -306,6 +306,12 @@ LogicalResult ReturnOp::verify() {
   return success();
 }
 
+void ReturnOp::getSuccessorRegions(ArrayRef<Attribute> operands,
+                                   SmallVectorImpl<RegionSuccessor> &regions) {
+  // Return control back to func::FuncOp.
+  regions.push_back(RegionSuccessor());
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

[joker-eph]

Can this be triggered in a test?

[fabianmcg]

Can this be triggered in a test?

On a LIT test, no, only by creating an unit test. There are very few places upstream using this interface, dataflow being one of them.
The reason dataflow doesn't trigger it is because dataflow always makes sure the parent op is RegionBranchOp.

I stumbled upon this assert when I was testing the interface for building the CFG, and I didn't check if the parent was valid.

[joker-eph]

I'm just concerned about code coverage upstream: in general we aim to have all this kind of behavior tested.
There has to be tests around the Regionbranchopinterface in tree?

[fabianmcg]

Yes there are tests: https://github.com/llvm/llvm-project/blob/main/mlir/unittests/Interfaces/ControlFlowInterfacesTest.cpp for the interface.

This issue is specific to the implementation of func::FuncOp and func::ReturnOp. I would also say it's quite easy to overlook because ReturnLike is the one attaching the region branch interface (we should change the trait's name). One option to spot similar errors would be adding a verifier to RegionBranchTerminatorOpInterface and make sure the parent is always valid. However, it's likely we would undo that change in the future with early exit.