[mlir] Fix FunctionOpInterface impl for external func

[ZenithalHourlyRate]

For function declarations (i.e. func op has no entry block), the FunctionOpInterface method insertArgument and eraseArgument will cause segfault. This PR guards against manipulation of empty entry block by checking whether func op is external.

An example can be seen in google/heir#1324

The segfault trace

 #1 0x0000560f1289d9db PrintStackTraceSignalHandler(void*) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:874:1
 #2 0x0000560f1289b116 llvm::sys::RunSignalHandlers() /proc/self/cwd/external/llvm-project/llvm/lib/Support/Signals.cpp:105:5
 #3 0x0000560f1289e145 SignalHandler(int) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:415:1
 #4 0x00007f829a3d9520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #5 0x0000560f1257f8bc void __gnu_cxx::new_allocator<mlir::BlockArgument>::construct<mlir::BlockArgument, mlir::BlockArgument>(mlir::BlockArgument*, mlir::BlockArgument&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/ext/new_allocator.h:162:23
 #6 0x0000560f1257f84d void std::allocator_traits<std::allocator<mlir::BlockArgument> >::construct<mlir::BlockArgument, mlir::BlockArgument>(std::allocator<mlir::BlockArgument>&, mlir::BlockArgument*, mlir::BlockArgument&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/alloc_traits.h:520:2
 #7 0x0000560f12580498 void std::vector<mlir::BlockArgument, std::allocator<mlir::BlockArgument> >::_M_insert_aux<mlir::BlockArgument>(__gnu_cxx::__normal_iterator<mlir::BlockArgument*, std::vector<mlir::BlockArgument, std::allocator<mlir::BlockArgument> > >, mlir::BlockArgument&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/vector.tcc:405:7
 #8 0x0000560f1257cf7e std::vector<mlir::BlockArgument, std::allocator<mlir::BlockArgument> >::insert(__gnu_cxx::__normal_iterator<mlir::BlockArgument const*, std::vector<mlir::BlockArgument, std::allocator<mlir::BlockArgument> > >, mlir::BlockArgument const&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/vector.tcc:154:6
 #9 0x0000560f1257b349 mlir::Block::insertArgument(unsigned int, mlir::Type, mlir::Location) /proc/self/cwd/external/llvm-project/mlir/lib/IR/Block.cpp:178:13
#10 0x0000560f123d2a1c mlir::function_interface_impl::insertFunctionArguments(mlir::FunctionOpInterface, llvm::ArrayRef<unsigned int>, mlir::TypeRange, llvm::ArrayRef<mlir::DictionaryAttr>, llvm::ArrayRef<mlir::Location>, unsigned int, mlir::Type) /proc/self/cwd/external/llvm-project/mlir/lib/Interfaces/FunctionInterfaces.cpp:232:11
#11 0x0000560f0be6b727 mlir::detail::FunctionOpInterfaceTrait<mlir::func::FuncOp>::insertArguments(llvm::ArrayRef<unsigned int>, mlir::TypeRange, llvm::ArrayRef<mlir::DictionaryAttr>, llvm::ArrayRef<mlir::Location>) /proc/self/cwd/bazel-out/k8-dbg/bin/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.h.inc:809:7
#12 0x0000560f0be6b536 mlir::detail::FunctionOpInterfaceTrait<mlir::func::FuncOp>::insertArgument(unsigned int, mlir::Type, mlir::DictionaryAttr, mlir::Location) /proc/self/cwd/bazel-out/k8-dbg/bin/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.h.inc:796:7

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Hongren Zheng (ZenithalHourlyRate)

Changes

For function declarations (i.e. func op has no entry block), the FunctionOpInterface method insertArgument and eraseArgument will cause segfault. This PR guards against manipulation of empty entry block by checking whether func op is external.

An example can be seen in google/heir#1324

The segfault trace

 #<!-- -->1 0x0000560f1289d9db PrintStackTraceSignalHandler(void*) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:874:1
 #<!-- -->2 0x0000560f1289b116 llvm::sys::RunSignalHandlers() /proc/self/cwd/external/llvm-project/llvm/lib/Support/Signals.cpp:105:5
 #<!-- -->3 0x0000560f1289e145 SignalHandler(int) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:415:1
 #<!-- -->4 0x00007f829a3d9520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
 #<!-- -->5 0x0000560f1257f8bc void __gnu_cxx::new_allocator&lt;mlir::BlockArgument&gt;::construct&lt;mlir::BlockArgument, mlir::BlockArgument&gt;(mlir::BlockArgument*, mlir::BlockArgument&amp;&amp;) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/ext/new_allocator.h:162:23
 #<!-- -->6 0x0000560f1257f84d void std::allocator_traits&lt;std::allocator&lt;mlir::BlockArgument&gt; &gt;::construct&lt;mlir::BlockArgument, mlir::BlockArgument&gt;(std::allocator&lt;mlir::BlockArgument&gt;&amp;, mlir::BlockArgument*, mlir::BlockArgument&amp;&amp;) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/alloc_traits.h:520:2
 #<!-- -->7 0x0000560f12580498 void std::vector&lt;mlir::BlockArgument, std::allocator&lt;mlir::BlockArgument&gt; &gt;::_M_insert_aux&lt;mlir::BlockArgument&gt;(__gnu_cxx::__normal_iterator&lt;mlir::BlockArgument*, std::vector&lt;mlir::BlockArgument, std::allocator&lt;mlir::BlockArgument&gt; &gt; &gt;, mlir::BlockArgument&amp;&amp;) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/vector.tcc:405:7
 #<!-- -->8 0x0000560f1257cf7e std::vector&lt;mlir::BlockArgument, std::allocator&lt;mlir::BlockArgument&gt; &gt;::insert(__gnu_cxx::__normal_iterator&lt;mlir::BlockArgument const*, std::vector&lt;mlir::BlockArgument, std::allocator&lt;mlir::BlockArgument&gt; &gt; &gt;, mlir::BlockArgument const&amp;) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/vector.tcc:154:6
 #<!-- -->9 0x0000560f1257b349 mlir::Block::insertArgument(unsigned int, mlir::Type, mlir::Location) /proc/self/cwd/external/llvm-project/mlir/lib/IR/Block.cpp:178:13
#<!-- -->10 0x0000560f123d2a1c mlir::function_interface_impl::insertFunctionArguments(mlir::FunctionOpInterface, llvm::ArrayRef&lt;unsigned int&gt;, mlir::TypeRange, llvm::ArrayRef&lt;mlir::DictionaryAttr&gt;, llvm::ArrayRef&lt;mlir::Location&gt;, unsigned int, mlir::Type) /proc/self/cwd/external/llvm-project/mlir/lib/Interfaces/FunctionInterfaces.cpp:232:11
#<!-- -->11 0x0000560f0be6b727 mlir::detail::FunctionOpInterfaceTrait&lt;mlir::func::FuncOp&gt;::insertArguments(llvm::ArrayRef&lt;unsigned int&gt;, mlir::TypeRange, llvm::ArrayRef&lt;mlir::DictionaryAttr&gt;, llvm::ArrayRef&lt;mlir::Location&gt;) /proc/self/cwd/bazel-out/k8-dbg/bin/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.h.inc:809:7
#<!-- -->12 0x0000560f0be6b536 mlir::detail::FunctionOpInterfaceTrait&lt;mlir::func::FuncOp&gt;::insertArgument(unsigned int, mlir::Type, mlir::DictionaryAttr, mlir::Location) /proc/self/cwd/bazel-out/k8-dbg/bin/external/llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.h.inc:796:7

---

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

1 Files Affected:

• (modified) mlir/lib/Interfaces/FunctionInterfaces.cpp (+11-5)

diff --git a/mlir/lib/Interfaces/FunctionInterfaces.cpp b/mlir/lib/Interfaces/FunctionInterfaces.cpp
index 80f47a3f836768..1d565eac2aeb88 100644
--- a/mlir/lib/Interfaces/FunctionInterfaces.cpp
+++ b/mlir/lib/Interfaces/FunctionInterfaces.cpp
@@ -199,7 +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 arguments of entry block, if not empty.
   Block &entry = op->getRegion(0).front();
 
   // Update the argument attributes of the function.
@@ -228,8 +228,11 @@ void function_interface_impl::insertFunctionArguments(
 
   // Update the function type and any entry block arguments.
   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 the block arguments of the entry block when it is not external.
+  if (!op.isExternal())
+    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,7 +282,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 arguments of entry block, if not empty.
   Block &entry = op->getRegion(0).front();
 
   // Update the argument attributes of the function.
@@ -294,7 +297,10 @@ void function_interface_impl::eraseFunctionArguments(
 
   // Update the function type and any entry block arguments.
   op.setFunctionTypeAttr(TypeAttr::get(newType));
-  entry.eraseArguments(argIndices);
+
+  // Update the block arguments of the entry block when it is not external.
+  if (!op.isExternal())
+    entry.eraseArguments(argIndices);
 }
 
 void function_interface_impl::eraseFunctionResults(

[ZenithalHourlyRate]

Cc @fabianschuiki, @River707 for review

[ZenithalHourlyRate]

Ping for review

[ZenithalHourlyRate]

Comments addressed and tested insertArgument for op declaration in downstream project.

[ZenithalHourlyRate]

Ping for review

[ZenithalHourlyRate]

Ping for review