[Mlir-commits] [mlir] 220f91a - [MLIR] Fix crash in inliner when return arity mismatches call results (#185037)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-06T16:50:04Z
New Revision: 220f91a05b3d5cb2c81d38afff83892f78d8a586

URL: https://github.com/llvm/llvm-project/commit/220f91a05b3d5cb2c81d38afff83892f78d8a586
DIFF: https://github.com/llvm/llvm-project/commit/220f91a05b3d5cb2c81d38afff83892f78d8a586.diff

LOG: [MLIR] Fix crash in inliner when return arity mismatches call results (#185037)

The `handleTerminator` implementation in the test dialect's inliner
interface was asserting that the number of `test.return` operands equals
the number of values to replace. This assertion fires when inlining a
callee whose body uses `test.return` with values into a call site that
expects zero results (e.g., a void `llvm.func` calling a function whose
implementation uses `test.return` with operands).

Replace the assertion with a conditional early return so the inliner
gracefully skips replacement instead of crashing.

Fixes #108376

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/test/Transforms/inlining.mlir
    mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index d8e10aa4212ba..63166e16b27b4 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -375,3 +375,25 @@ func.func @inline_with_complex_ops() -> complex<f32> {
   %r = call @double_square_complex(%c) : (complex<f32>) -> (complex<f32>)
   return %r : complex<f32>
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/108376:
+// Inlining a function whose test.return arity doesn't match the call result
+// count (e.g., mixing llvm.func returning void with test.return returning
+// values) should not crash with an assertion failure.
+
+// CHECK-LABEL: llvm.func @inline_test_return_arity_mismatch
+llvm.func @inline_test_return_arity_mismatch(%arg0: f16, %arg1: f16) attributes {llvm.emit_c_interface} {
+  // CHECK: test.op_with_bitcast_type
+  // CHECK: test.op_with_bitcast_type
+  // CHECK: llvm.return
+  // CHECK-NOT: llvm.call
+  llvm.call @inline_test_return_arity_mismatch_callee(%arg0, %arg1) : (f16, f16) -> ()
+  llvm.return
+}
+llvm.func @inline_test_return_arity_mismatch_callee(%arg0: f16, %arg1: f16) {
+  %0 = "test.op_with_bitcast_type"(%arg0) : (f16) -> tensor<4xf32>
+  %1 = "test.op_with_bitcast_type"(%arg1) : (f16) -> tensor<2xi32>
+  "test.return"(%0, %1) : (tensor<4xf32>, tensor<2xi32>) -> ()
+}

diff  --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
index 3d4aa23ebe78a..7ccbe49b2c62e 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -350,8 +350,11 @@ struct TestInlinerInterface : public DialectInlinerInterface {
     if (!returnOp)
       return;
 
-    // Replace the values directly with the return operands.
-    assert(returnOp.getNumOperands() == valuesToRepl.size());
+    // Replace the values directly with the return operands. Skip if the
+    // number of operands doesn't match (e.g., when inlining into a call
+    // with a 
diff erent result arity due to invalid IR mixing dialects).
+    if (returnOp.getNumOperands() != valuesToRepl.size())
+      return;
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }