[Mlir-commits] [mlir] Make `FunctionOpInterface` check `ReturnLike` (PR #112615)

[Tzung-Han Juang]

https://github.com/tzunghanjuang created https://github.com/llvm/llvm-project/pull/112615

**Description:** 

This PR is a follow-up for https://github.com/llvm/llvm-project/pull/110322. The previous PR makes `OneShotBufferize` check if a function operation has a terminator with the `ReturnLike` trait. However, `transform.named_sequence` does not expect `yieldOp` to have the trait. In that PR, the issue is circumvented by `--transform-interpreter="debug-payload-root-tag=payload" `.

This PR tries to tie `ReturnLike` to `FunctionOpInface`. However, I am not sure if it is correct to check `ReturnLike` in the tablegen for `FunctionOpInface`. I have not yet handled all the errors.  Just want to create this PR for further confirmation. 
We might also need to wait for this fix for `transform.named_sequence`  (https://github.com/llvm/llvm-project/pull/111408) to be merged?

>From 3e34c1d57ad0f0e9408c4589ed325ff63e0f34a7 Mon Sep 17 00:00:00 2001
From: Tzung-Han Juang <tzung-han.juang at mail.mcgill.ca>
Date: Tue, 15 Oct 2024 17:10:26 -0400
Subject: [PATCH 1/2] Fix error message for FunctionOpInterface

---
 .../Bufferization/Transforms/OneShotModuleBufferize.cpp       | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index a0e5c7fff7690f..bb1cf2fc3fd209 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -302,8 +302,8 @@ getFuncOpsOrderedByCalls(ModuleOp moduleOp,
       Operation *returnOp = getAssumedUniqueReturnOp(funcOp);
       if (!returnOp)
         return funcOp->emitError()
-               << "cannot bufferize a FuncOp with tensors and "
-                  "without a unique ReturnOp";
+               << "cannot bufferize a FunctionOpInterface with tensors and "
+                  "without a unique ReturnLike";
     }
 
     // Collect function calls and populate the caller map.

>From 1305c2d2f966b84799c575aec9321c75f4c8a3c1 Mon Sep 17 00:00:00 2001
From: Tzung-Han Juang <tzung-han.juang at mail.mcgill.ca>
Date: Wed, 16 Oct 2024 16:48:06 -0400
Subject: [PATCH 2/2] Add ReturnLike check for FunctionOpInterface

---
 .../mlir/Interfaces/FunctionInterfaces.h      |  1 +
 .../mlir/Interfaces/FunctionInterfaces.td     | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.h b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
index e10e9bd342702a..f121a6823711a0 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.h
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
@@ -20,6 +20,7 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
 
diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.td b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
index 697f951748c675..244c7b33f0119b 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.td
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
@@ -108,6 +108,29 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
         }
       }
 
+      // FunctionOpInterface is tied to a ReturnLike.
+      Operation *terminator = entryBlock.getTerminator();
+      if (terminator->hasTrait<OpTrait::ReturnLike>()) {
+        return $_op.emitOpError("The body of a FunctionOpInterface must")
+              << "have a ReturnLike terminator.";
+      }
+
+      // Match ReturnLike's operand types and FunctionOpInterface's 
+      // result types.
+      auto returnOperandTypes = terminator->getOperandTypes();
+      auto funcResultTypes =  $_op->getResultTypes();
+      if (funcResultTypes.size() != returnOperandTypes.size()) {
+        return $_op.emitOpError("The number of a FunctionOpInterface's")
+              << "result must match that of the ReturnLike operands.";
+      }
+
+      for (unsigned i = 0; i < funcResultTypes.size(); ++i) {
+        if (funcResultTypes[i] != returnOperandTypes[i]) {
+          return $_op.emitOpError("The result types of a FunctionOpInterface")
+              << "must match the operand types of the ReturnLike.";
+        }
+      }
+
       return success();
     }]>,
     InterfaceMethod<[{