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

[llvmlistbot at llvm.org]

Author: Hongren Zheng
Date: 2025-02-19T12:49:40+08:00
New Revision: 8b284dc31070b9d1d99c593146da6248a5ca545d

URL: https://github.com/llvm/llvm-project/commit/8b284dc31070b9d1d99c593146da6248a5ca545d
DIFF: https://github.com/llvm/llvm-project/commit/8b284dc31070b9d1d99c593146da6248a5ca545d.diff

LOG: [mlir] Fix FunctionOpInterface impl for external func (#124693)

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 https://github.com/google/heir/pull/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
```

Added: 
    

Modified: 
    mlir/lib/Interfaces/FunctionInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Interfaces/FunctionInterfaces.cpp b/mlir/lib/Interfaces/FunctionInterfaces.cpp
index 80f47a3f83676..57a8668117c68 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(