From 8b284dc31070b9d1d99c593146da6248a5ca545d Mon Sep 17 00:00:00 2001 From: Hongren Zheng Date: Wed, 19 Feb 2025 12:49:40 +0800 Subject: [PATCH] [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::construct(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 >::construct(std::allocator&, 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 >::_M_insert_aux(__gnu_cxx::__normal_iterator > >, mlir::BlockArgument&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/vector.tcc:405:7 #8 0x0000560f1257cf7e std::vector >::insert(__gnu_cxx::__normal_iterator > >, 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, mlir::TypeRange, llvm::ArrayRef, llvm::ArrayRef, unsigned int, mlir::Type) /proc/self/cwd/external/llvm-project/mlir/lib/Interfaces/FunctionInterfaces.cpp:232:11 #11 0x0000560f0be6b727 mlir::detail::FunctionOpInterfaceTrait::insertArguments(llvm::ArrayRef, mlir::TypeRange, llvm::ArrayRef, llvm::ArrayRef) /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::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 ``` --- 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 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(