[Mlir-commits] [mlir] [mlir][inliner] Ask the terminator dialect before the single-block fast path (PR #237936)

[VeeraBabu Jakkula]

https://github.com/JakkulaVeerababu created https://github.com/llvm/llvm-project/pull/237936

### What was wrong

MLIR's inliner was crashing when trying to inline a `func.func` that ended with `llvm.unreachable` or `ub.unreachable`. Before taking the single-block fast path, it only asked the dialect of the function (`func`), not the dialect of the last instruction. The last instruction is then handled by its own dialect's inliner interface, which could not handle the fast path:

- `LLVMInlinerInterface::handleTerminator` casts the terminator to `llvm.return`, so `llvm.unreachable` failed the cast assertion (#237592).
- `UBInlinerInterface` does not implement `handleTerminator`, so `ub.unreachable` hit the default `llvm_unreachable` (#237590).

### What I changed

I updated the inliner to check both the function's dialect and the last instruction's dialect before taking the fast path. The LLVM dialect already disables the fast path for `llvm.unreachable`, so this alone fixes #237592.

I also added proper handling for `ub.unreachable` in the UB dialect: it disables the fast path for it and keeps it in place when the inlined region has more than one block, the same way the LLVM dialect handles `llvm.unreachable`.

Finally, I added two tests based on both reported reproducers. Both crashes are fixed, the new tests fail without the change, and all 122 related inliner tests are passing.

Fixes #237592
Fixes #237590

This change was made with the help of Claude Code (AI assistant). I reviewed it and can explain every part of it.

>From 838cbdff867c47e3c3a4350f98702322e4bdcfb5 Mon Sep 17 00:00:00 2001
From: VeeraBabu Jakkula <squadpanda447 at gmail.com>
Date: Mon, 2 Nov 2026 12:02:33 +0000
Subject: [PATCH] [mlir][inliner] Ask the terminator dialect before the
 single-block fast path

The inliner decides on the single-block fast path by asking the dialect
of the inlined region's parent operation, but the terminator is then
handled by the interface of its own dialect. A func.func callee that
ends in llvm.unreachable passed the check through the func dialect, and
LLVMInlinerInterface::handleTerminator then cast the unreachable to
llvm.return and asserted. With ub.unreachable, the default UB hook hit
llvm_unreachable.

Also ask the dialect of the terminator, and let the UB dialect disable
the fast path for ub.unreachable and keep it in place when the inlined
region has more than one block, as the LLVM dialect already does for
llvm.unreachable.

Fixes #237592
Fixes #237590

Assisted-by: Claude Code
---
 mlir/lib/Dialect/UB/IR/UBOps.cpp            | 15 +++++++++++++++
 mlir/lib/Transforms/Utils/InliningUtils.cpp | 14 +++++++++++++-
 mlir/test/Dialect/LLVMIR/inlining.mlir      | 17 +++++++++++++++++
 mlir/test/Transforms/inlining.mlir          | 16 ++++++++++++++++
 4 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/UB/IR/UBOps.cpp b/mlir/lib/Dialect/UB/IR/UBOps.cpp
index ee523f9522953d..64bdd4797dd633 100644
--- a/mlir/lib/Dialect/UB/IR/UBOps.cpp
+++ b/mlir/lib/Dialect/UB/IR/UBOps.cpp
@@ -29,6 +29,21 @@ struct UBInlinerInterface : public DialectInlinerInterface {
   bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
+
+  /// `ub.unreachable` does not return to the caller, so it stays in place when
+  /// the inlined region has more than one block.
+  void handleTerminator(Operation *op, Block *newDest) const final {
+    assert(isa<UnreachableOp>(op) && "unexpected UB terminator");
+  }
+
+  /// The single-block fast path erases the terminator and continues the
+  /// inlined block with the operations after the call, which is wrong for
+  /// `ub.unreachable`.
+  bool allowSingleBlockOptimization(
+      iterator_range<Region::iterator> inlinedBlocks) const final {
+    return inlinedBlocks.empty() ||
+           !isa<UnreachableOp>(inlinedBlocks.begin()->getTerminator());
+  }
 };
 } // namespace
 
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index ce4d6ad71b6d3a..a0c12e417e067a 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -151,7 +151,19 @@ bool InlinerInterface::allowSingleBlockOptimization(
   }
   auto *handler = getInterfaceFor(inlinedBlocks.begin()->getParentOp());
   assert(handler && "expected valid dialect handler");
-  return handler->allowSingleBlockOptimization(inlinedBlocks);
+  if (!handler->allowSingleBlockOptimization(inlinedBlocks))
+    return false;
+
+  // The terminator is handled by the interface of its own dialect, which may
+  // differ from the dialect of the parent operation, so it must also allow the
+  // fast path.
+  Block &firstBlock = *inlinedBlocks.begin();
+  if (!firstBlock.mightHaveTerminator())
+    return true;
+  auto *terminatorHandler = getInterfaceFor(firstBlock.getTerminator());
+  if (!terminatorHandler || terminatorHandler == handler)
+    return true;
+  return terminatorHandler->allowSingleBlockOptimization(inlinedBlocks);
 }
 
 Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 30616eaa715359..3c5203b72637f3 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -794,3 +794,20 @@ llvm.func @caller_cf_br_terminator(%arg0 : i64) -> i64 {
   %0 = llvm.call @callee_cf_br_terminator(%arg0) : (i64) -> i64
   llvm.return %0 : i64
 }
+
+// -----
+
+// A func.func callee that ends in llvm.unreachable must not take the
+// single-block fast path, which only handles llvm.return.
+
+func.func @func_unreachable_callee() {
+  llvm.unreachable
+}
+
+// CHECK-LABEL: func.func @func_unreachable_caller
+// CHECK-NOT: call @func_unreachable_callee
+// CHECK: llvm.unreachable
+func.func @func_unreachable_caller() {
+  call @func_unreachable_callee() : () -> ()
+  return
+}
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 5b3b1d8cd211f4..5c041da8ebf598 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -420,3 +420,19 @@ func.func @inline_functional_region_multiblock(%arg0: i32) -> i32 {
   %0 = call_indirect %fn(%arg0) : (i32) -> i32
   return %0 : i32
 }
+
+// A func.func callee that ends in ub.unreachable must not take the single-block
+// fast path, which would erase the terminator.
+func.func @func_with_ub_unreachable() {
+  ub.unreachable
+}
+
+// CHECK-LABEL: func @inline_ub_unreachable(
+func.func @inline_ub_unreachable() {
+  // CHECK-NOT:  call
+  // CHECK:      ub.unreachable
+  // CHECK-NEXT: ^{{.*}}:
+  // CHECK-NEXT: return
+  call @func_with_ub_unreachable() : () -> ()
+  return
+}