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

[William Moses]

https://github.com/wsmoses updated https://github.com/llvm/llvm-project/pull/122646

>From 9d264092b1dd50e4f149b970a0f98a32a2390736 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Sat, 11 Jan 2025 15:28:42 -0500
Subject: [PATCH] [MLIR][LLVM] Fix inlining of a single block ending with
 unreachable

---
 mlir/include/mlir/Transforms/InliningUtils.h  | 17 ++++++++++++-----
 .../Transforms/InlinerInterfaceImpl.cpp       |  8 ++++++++
 mlir/lib/Transforms/Utils/Inliner.cpp         |  3 ++-
 mlir/lib/Transforms/Utils/InliningUtils.cpp   | 19 +++++++++++++++++--
 mlir/test/Dialect/LLVMIR/inlining.mlir        | 16 ++++++++++++++++
 5 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index 88fc033a6ab7be..24954325237adc 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -176,6 +176,15 @@ class DialectInlinerInterface
   /// is invoked before inlined terminator operations have been processed.
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {}
+
+  /// Process a set of blocks that have been inlined. This callback is invoked
+  /// *before* inlined terminator operations have been processed. Returns true
+  /// if the inliner can assume a fast path of not creating a new block, if
+  /// there is only one block.
+  virtual bool
+  processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) const {
+    return true;
+  }
 };
 
 /// This interface provides the hooks into the inlining interface.
@@ -186,11 +195,6 @@ class InlinerInterface
 public:
   using Base::Base;
 
-  /// Process a set of blocks that have been inlined. This callback is invoked
-  /// *before* inlined terminator operations have been processed.
-  virtual void
-  processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) {}
-
   /// These hooks mirror the hooks for the DialectInlinerInterface, with default
   /// implementations that call the hook on the handler for the dialect 'op' is
   /// registered to.
@@ -211,6 +215,9 @@ class InlinerInterface
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
+  virtual bool
+  processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks);
+
   virtual void handleTerminator(Operation *op, Block *newDest) const;
   virtual void handleTerminator(Operation *op, ValueRange valuesToRepl) const;
 
diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index b3bed5ab5f412f..73bcb7526af4e4 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -743,6 +743,14 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
     op->erase();
   }
 
+  bool processInlinedBlocks(
+      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/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 8acfc96d2b611b..a153ab1182e9ae 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -382,7 +382,7 @@ struct InlinerInterfaceImpl : public InlinerInterface {
 
   /// Process a set of blocks that have been inlined. This callback is invoked
   /// *before* inlined terminator operations have been processed.
-  void
+  bool
   processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) final {
     // Find the closest callgraph node from the first block.
     CallGraphNode *node;
@@ -394,6 +394,7 @@ struct InlinerInterfaceImpl : public InlinerInterface {
 
     collectCallOps(inlinedBlocks, node, cg, symbolTable, calls,
                    /*traverseNestedCGNodes=*/true);
+    return InlinerInterface::processInlinedBlocks(inlinedBlocks);
   }
 
   /// Mark the given callgraph node for deletion.
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0db097d14cd3c7..6c3e51e7770387 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -118,6 +118,21 @@ void InlinerInterface::handleTerminator(Operation *op,
   handler->handleTerminator(op, valuesToRepl);
 }
 
+/// Process a set of blocks that have been inlined. This callback is invoked
+/// *before* inlined terminator operations have been processed. Returns true
+/// if the inliner can assume a fast path of not creating a new block, if
+/// there is only one block.
+bool InlinerInterface::processInlinedBlocks(
+    iterator_range<Region::iterator> inlinedBlocks) {
+  if (inlinedBlocks.empty()) {
+    return true;
+  } else {
+    auto *handler = getInterfaceFor(inlinedBlocks.begin()->getParentOp());
+    assert(handler && "expected valid dialect handler");
+    return handler->processInlinedBlocks(inlinedBlocks);
+  }
+}
+
 Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
                                        Operation *callable, Value argument,
                                        DictionaryAttr argumentAttrs) const {
@@ -292,10 +307,10 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   // Process the newly inlined blocks.
   if (call)
     interface.processInlinedCallBlocks(call, newBlocks);
-  interface.processInlinedBlocks(newBlocks);
+  bool singleBlockFastPath = interface.processInlinedBlocks(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
+}