diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp index 06c0af5445d5d..0e43480e82926 100644 --- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp +++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp @@ -725,6 +725,20 @@ struct LLVMInlinerInterface : public DialectInlinerInterface { })) return false; } + // Refuse to inline if any block in the callee ends with an op that does + // not have the terminator trait. The MLIR verifier conservatively accepts + // unregistered ops as potential terminators (via mightHaveTrait), but + // handleTerminator uses cast in the single-block path and + // would crash on such ops. Registered terminators from other dialects + // (e.g. cf.br) are safe: the multi-block path uses dyn_cast and skips + // non-llvm.return ops gracefully. + for (Block &block : funcOp.getBody()) { + if (!block.empty() && !block.back().hasTrait()) { + LDBG() << "Cannot inline " << funcOp.getSymName() + << ": block ends with non-terminator op"; + return false; + } + } return true; } diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir index 9a77c5e223110..70ce7ca20986b 100644 --- a/mlir/test/Dialect/LLVMIR/inlining.mlir +++ b/mlir/test/Dialect/LLVMIR/inlining.mlir @@ -709,3 +709,42 @@ func.func @llvm_ret(%arg0 : i32) -> i32 { // CHECK: return %[[R]] return %res : i32 } + +// ----- +// Regression test for https://github.com/llvm/llvm-project/issues/118766 +// A callee whose body block ends with an unregistered (non-terminator) op used +// to crash the inliner. The inliner should skip such callees instead. + +llvm.func @callee_malformed_terminator() -> i64 attributes {llvm.emit_c_interface} { + // Unregistered op acting as a fake terminator -- the function has no llvm.return. + "test.foo"() : () -> () +} + +// CHECK-LABEL: @caller_malformed_terminator +llvm.func @caller_malformed_terminator() -> i64 attributes {llvm.emit_c_interface} { + // CHECK: llvm.call @callee_malformed_terminator + %0 = llvm.call @callee_malformed_terminator() : () -> i64 + llvm.return %0 : i64 +} + +// ----- +// Complement to the above: a callee whose block ends with a registered +// non-LLVM terminator that genuinely has the IsTerminator trait (cf.br) does +// NOT trigger the guard. The call is inlined normally via the multi-block path +// (handleTerminator uses dyn_cast for non-llvm.return ops, so cf.br is left +// as-is and control flow reaches the llvm.return in the successor block). + +llvm.func @callee_cf_br_terminator(%arg0 : i64) -> i64 { + cf.br ^exit(%arg0 : i64) +^exit(%val : i64): + llvm.return %val : i64 +} + +// CHECK-LABEL: @caller_cf_br_terminator +llvm.func @caller_cf_br_terminator(%arg0 : i64) -> i64 { + // cf.br has IsTerminator, so isLegalToInline allows inlining. + // CHECK-NOT: llvm.call @callee_cf_br_terminator + // CHECK: llvm.return %arg0 + %0 = llvm.call @callee_cf_br_terminator(%arg0) : (i64) -> i64 + llvm.return %0 : i64 +}