[Mlir-commits] [mlir] [mlir][func] Fix ReturnOp issue 112146 (PR #112385)

[Fabian Mora]

https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/112385

This patches fixes issue https://github.com/llvm/llvm-project/issues/112146, where an assertion was being triggered by `func::ReturnOp::getSuccessorRegions` and `func::FuncOp` not implementing `RegionBranchOpInterface`.

>From 5181d96e367b46aa42309036207d801aadf83e4a Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Tue, 15 Oct 2024 15:08:49 +0000
Subject: [PATCH] [mlir][func] Fix ReturnOp issue 112146

This patches fixes issue 112146, where an assertion was being triggered by
`func::ReturnOp::getSuccessorRegions` and `func::FuncOp` not implementing
`RegionBranchOpInterface`.
---
 mlir/include/mlir/Dialect/Func/IR/FuncOps.td | 6 ++++--
 mlir/lib/Dialect/Func/IR/FuncOps.cpp         | 6 ++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

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
 //===----------------------------------------------------------------------===//