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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From bf84e8bc6b5ab17a4116c785d63d017dff8d9efe Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 5 Mar 2026 07:34:51 -0800
Subject: [PATCH] [MLIR] Fix crash in inliner when return arity mismatches call
 results

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
---
 mlir/test/Transforms/inlining.mlir            | 22 +++++++++++++++++++
 .../Dialect/Test/TestDialectInterfaces.cpp    |  7 ++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

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 different 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());
   }