[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122646

wsmoses · 2025-01-12T19:33:20Z

alternate option to #122615

llvmbot · 2025-01-12T19:33:51Z

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: William Moses (wsmoses)

Changes

alternate option to #122615

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

5 Files Affected:

(modified) mlir/include/mlir/Transforms/InliningUtils.h (+12-5)
(modified) mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp (+9-1)
(modified) mlir/lib/Transforms/Utils/Inliner.cpp (+2-1)
(modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+17-2)
(modified) mlir/test/Dialect/LLVMIR/inlining.mlir (+16)

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..1635809ce160c8 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -743,11 +743,19 @@ 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.
   void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
-    // Return will be the only terminator present.
+    // Otherwise return will be the only terminator present.
     auto returnOp = cast<LLVM::ReturnOp>(op);
 
     // Replace the values directly with the return operands.
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
+}

llvmbot · 2025-01-12T19:33:52Z

@llvm/pr-subscribers-mlir-llvm

Author: William Moses (wsmoses)

Changes

alternate option to #122615

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

5 Files Affected:

(modified) mlir/include/mlir/Transforms/InliningUtils.h (+12-5)
(modified) mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp (+9-1)
(modified) mlir/lib/Transforms/Utils/Inliner.cpp (+2-1)
(modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+17-2)
(modified) mlir/test/Dialect/LLVMIR/inlining.mlir (+16)

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..1635809ce160c8 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -743,11 +743,19 @@ 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.
   void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
-    // Return will be the only terminator present.
+    // Otherwise return will be the only terminator present.
     auto returnOp = cast<LLVM::ReturnOp>(op);
 
     // Replace the values directly with the return operands.
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
+}

gysit · 2025-01-13T09:04:20Z

I believe I prefer this solution over the other PR (especially with regards to the generated code).

Should we add a separate function to the inlining interface called allowSingleBlockOptimization or similar. The function could have a default implementation that returns true. Then only the LLVM inlining interface would have to overwrite this function. Adding boolean result to processInlinedBlocks works, but it feels a bit like the result flag would be unrelated to the actual purpose of the function.

Personally, my preferred solution would be to remove the inlining "fast path" entirely and rely on region simplify to cleanup the control flow. I am not sure how easy it is to find consensus for such a change though.

@ftynse do you have an opinion here?

wsmoses · 2025-01-13T16:15:18Z

Removing the fast past is reasonable, but I feel that should be a follow up PR and/or RFC given the wider implication of the change.

I'm happy to change this to a dedicated fn

wsmoses · 2025-01-13T16:35:25Z

@gysit okay pushed a version with the dedicated function, if you can give it another look

ftynse · 2025-01-13T17:04:11Z

I'm onboard with the current proposition (allowSingleBlockOptimization), looks like the least intrusive change IMO. We can always revise if we get consensus for a bigger change. I'm rather undecided on having to run region simplification as a cleanup. Same for canonicalization, btw. On one hand, I don't want to duplicate simplification code everywhere when we have dedicated logic. On another hand, I don't think it's reasonable to run canonicalizer/CSE/DCE/CFG-simplify after every single pass. We can circulate the idea of removing this optimization completely on the forum, but beware of it devolving into that discussion.

River707 · 2025-01-13T17:25:59Z

Just coming out of OOO, sorry for not participating on the other review. In general one thing you need to keep in mind is the fact that some dialects only have single blocks (i.e. no branching control flow), in those cases you can't have multiple blocks in any scenario. These kinds of handling should always be left to the dialect. I'm fine with this PR as-is, though for dialects with single blocks and unreachable something else will likely need to happen.

mlir/lib/Transforms/Utils/InliningUtils.cpp

…lvm#122646) alternate option to llvm#122615

wsmoses requested review from ftynse, definelicht and gysit January 12, 2025 19:33

llvmbot added mlir:core MLIR Core Infrastructure mlir:llvm mlir labels Jan 12, 2025

wsmoses mentioned this pull request Jan 12, 2025

[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122615

Closed

wsmoses force-pushed the mlirterm2 branch from 5b2c3de to 9d26409 Compare January 12, 2025 21:43

wsmoses mentioned this pull request Jan 13, 2025

WIP: adapt to sroa jll EnzymeAD/Reactant.jl#521

Merged

wsmoses force-pushed the mlirterm2 branch from 9d26409 to fd1cbf0 Compare January 13, 2025 16:35

ftynse approved these changes Jan 13, 2025

View reviewed changes

River707 reviewed Jan 13, 2025

View reviewed changes

mlir/lib/Transforms/Utils/InliningUtils.cpp Outdated Show resolved Hide resolved

[MLIR][LLVM] Fix inlining of a single block ending with unreachable

ee9de2e

wsmoses force-pushed the mlirterm2 branch from fd1cbf0 to ee9de2e Compare January 13, 2025 17:29

wsmoses merged commit b39c5cb into llvm:main Jan 13, 2025
6 of 7 checks passed

wsmoses deleted the mlirterm2 branch January 13, 2025 18:37

kazutakahirata pushed a commit to kazutakahirata/llvm-project that referenced this pull request Jan 13, 2025

[MLIR][LLVM] Fix inlining of a single block ending with unreachable (l…

a0db95c

…lvm#122646) alternate option to llvm#122615

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122646

[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122646

Uh oh!

wsmoses commented Jan 12, 2025

Uh oh!

llvmbot commented Jan 12, 2025 •

edited

Loading

Uh oh!

llvmbot commented Jan 12, 2025

Uh oh!

gysit commented Jan 13, 2025

Uh oh!

wsmoses commented Jan 13, 2025

Uh oh!

wsmoses commented Jan 13, 2025

Uh oh!

ftynse commented Jan 13, 2025

Uh oh!

River707 commented Jan 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122646

[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122646

Uh oh!

Conversation

wsmoses commented Jan 12, 2025

Uh oh!

llvmbot commented Jan 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jan 12, 2025

Uh oh!

gysit commented Jan 13, 2025

Uh oh!

wsmoses commented Jan 13, 2025

Uh oh!

wsmoses commented Jan 13, 2025

Uh oh!

ftynse commented Jan 13, 2025

Uh oh!

River707 commented Jan 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Jan 12, 2025 •

edited

Loading