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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: VeeraBabu Jakkula  (JakkulaVeerababu)

<details>
<summary>Changes</summary>

### 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.

---
Full diff: https://github.com/llvm/llvm-project/pull/237936.diff


4 Files Affected:

- (modified) mlir/lib/Dialect/UB/IR/UBOps.cpp (+15) 
- (modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+13-1) 
- (modified) mlir/test/Dialect/LLVMIR/inlining.mlir (+17) 
- (modified) mlir/test/Transforms/inlining.mlir (+16) 


``````````diff
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
+}

``````````

</details>


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