[Mlir-commits] [mlir] b39c5cb - [MLIR][LLVM] Fix inlining of a single block ending with unreachable (#122646)

[llvmlistbot at llvm.org]

Author: William Moses
Date: 2025-01-13T12:37:16-06:00
New Revision: b39c5cb6977f35ad727d86b2dd6232099734ffd3

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

LOG: [MLIR][LLVM] Fix inlining of a single block ending with unreachable (#122646)

alternate option to https://github.com/llvm/llvm-project/pull/122615

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/InliningUtils.h
    mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
    mlir/lib/Transforms/Utils/InliningUtils.cpp
    mlir/test/Dialect/LLVMIR/inlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index 88fc033a6ab7be..becfe9b047ef45 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -176,6 +176,13 @@ class DialectInlinerInterface
   /// is invoked before inlined terminator operations have been processed.
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {}
+
+  /// Returns true if the inliner can assume a fast path of not creating a new
+  /// block, if there is only one block.
+  virtual bool allowSingleBlockOptimization(
+      iterator_range<Region::iterator> inlinedBlocks) const {
+    return true;
+  }
 };
 
 /// This interface provides the hooks into the inlining interface.
@@ -223,6 +230,9 @@ class InlinerInterface
 
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const;
+
+  virtual bool allowSingleBlockOptimization(
+      iterator_range<Region::iterator> inlinedBlocks) const;
 };
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index 233cadebeec026..79dd3e30696482 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -744,6 +744,14 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
     op->erase();
   }
 
+  bool allowSingleBlockOptimization(
+      iterator_range<Region::iterator> inlinedBlocks) const final {
+    if (!inlinedBlocks.empty() &&
+        isa<LLVM::UnreachableOp>(inlinedBlocks.begin()->getTerminator()))
+      return false;
+    return true;
+  }
+
   /// Handle the given inlined return by replacing the uses of the call with the
   /// operands of the return. This overload is called when the inlined region
   /// only contains one block.

diff  --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0db097d14cd3c7..0cae63c58ca7be 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -118,6 +118,18 @@ void InlinerInterface::handleTerminator(Operation *op,
   handler->handleTerminator(op, valuesToRepl);
 }
 
+/// Returns true if the inliner can assume a fast path of not creating a
+/// new block, if there is only one block.
+bool InlinerInterface::allowSingleBlockOptimization(
+    iterator_range<Region::iterator> inlinedBlocks) const {
+  if (inlinedBlocks.empty()) {
+    return true;
+  }
+  auto *handler = getInterfaceFor(inlinedBlocks.begin()->getParentOp());
+  assert(handler && "expected valid dialect handler");
+  return handler->allowSingleBlockOptimization(inlinedBlocks);
+}
+
 Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
                                        Operation *callable, Value argument,
                                        DictionaryAttr argumentAttrs) const {
@@ -294,8 +306,10 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
     interface.processInlinedCallBlocks(call, newBlocks);
   interface.processInlinedBlocks(newBlocks);
 
+  bool singleBlockFastPath = interface.allowSingleBlockOptimization(newBlocks);
+
   // Handle the case where only a single block was inlined.
-  if (std::next(newBlocks.begin()) == newBlocks.end()) {
+  if (singleBlockFastPath && std::next(newBlocks.begin()) == newBlocks.end()) {
     // Run the result attribute handler on the terminator operands.
     Operation *firstBlockTerminator = firstNewBlock->getTerminator();
     builder.setInsertionPoint(firstBlockTerminator);

diff  --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index edaac4da0b044c..eb249a47717534 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -676,3 +676,19 @@ llvm.func @caller(%x : i32) -> i32 {
   %z = llvm.call @private_func(%x) : (i32) -> (i32)
   llvm.return %z : i32
 }
+
+// -----
+
+llvm.func @unreachable_func(%a : i32) -> i32 {
+  "llvm.intr.trap"() : () -> ()
+  llvm.unreachable
+}
+
+// CHECK-LABEL: func @caller
+llvm.func @caller(%x : i32) -> i32 {
+  // CHECK-NOT: llvm.call @unreachable_func
+  // CHECK: llvm.intr.trap
+  // CHECK: llvm.unreachable
+  %z = llvm.call @unreachable_func(%x) : (i32) -> (i32)
+  llvm.return %z : i32
+}