[Mlir-commits] [mlir] [mlir] Fix FunctionOpInterface impl for external func (PR #124693)

[Hongren Zheng]

https://github.com/ZenithalHourlyRate updated https://github.com/llvm/llvm-project/pull/124693

>From 77c9b3717ad05e51ff3be5ea96b7ac9faf7cf5d8 Mon Sep 17 00:00:00 2001
From: Zenithal <i at zenithal.me>
Date: Tue, 28 Jan 2025 04:18:48 +0000
Subject: [PATCH] [mlir] Fix FunctionOpInterface impl for external func

---
 mlir/lib/Interfaces/FunctionInterfaces.cpp | 26 ++++++++++++++--------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Interfaces/FunctionInterfaces.cpp b/mlir/lib/Interfaces/FunctionInterfaces.cpp
index 80f47a3f836768b..57a8668117c684e 100644
--- a/mlir/lib/Interfaces/FunctionInterfaces.cpp
+++ b/mlir/lib/Interfaces/FunctionInterfaces.cpp
@@ -199,8 +199,7 @@ void function_interface_impl::insertFunctionArguments(
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
-  // - Block arguments of entry block.
-  Block &entry = op->getRegion(0).front();
+  // - Block arguments of entry block, if not empty.
 
   // Update the argument attributes of the function.
   ArrayAttr oldArgAttrs = op.getArgAttrsAttr();
@@ -226,10 +225,15 @@ void function_interface_impl::insertFunctionArguments(
     setAllArgAttrDicts(op, newArgAttrs);
   }
 
-  // Update the function type and any entry block arguments.
+  // Update the function type.
   op.setFunctionTypeAttr(TypeAttr::get(newType));
-  for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
-    entry.insertArgument(argIndices[i] + i, argTypes[i], argLocs[i]);
+
+  // Update entry block arguments, if not empty.
+  if (!op.isExternal()) {
+    Block &entry = op->getRegion(0).front();
+    for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
+      entry.insertArgument(argIndices[i] + i, argTypes[i], argLocs[i]);
+  }
 }
 
 void function_interface_impl::insertFunctionResults(
@@ -279,8 +283,7 @@ void function_interface_impl::eraseFunctionArguments(
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
-  // - Block arguments of entry block.
-  Block &entry = op->getRegion(0).front();
+  // - Block arguments of entry block, if not empty.
 
   // Update the argument attributes of the function.
   if (ArrayAttr argAttrs = op.getArgAttrsAttr()) {
@@ -292,9 +295,14 @@ void function_interface_impl::eraseFunctionArguments(
     setAllArgAttrDicts(op, newArgAttrs);
   }
 
-  // Update the function type and any entry block arguments.
+  // Update the function type.
   op.setFunctionTypeAttr(TypeAttr::get(newType));
-  entry.eraseArguments(argIndices);
+
+  // Update entry block arguments, if not empty.
+  if (!op.isExternal()) {
+    Block &entry = op->getRegion(0).front();
+    entry.eraseArguments(argIndices);
+  }
 }
 
 void function_interface_impl::eraseFunctionResults(