[MLIR][LLVM] Fix inlining of a single block ending with unreachable #122615
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
wsmoses
requested review from
antiagainst,
kuhar,
dcaballe,
nicolasvasilache and
rengolin
as code owners
llvmbot
added
mlir:core
|
@llvm/pr-subscribers-flang-fir-hlfir @llvm/pr-subscribers-mlir Author: William Moses (wsmoses) ChangesFull diff: https://github.com/llvm/llvm-project/pull/122615.diff 10 Files Affected:
diff --git a/mlir/examples/toy/Ch4/mlir/Dialect.cpp b/mlir/examples/toy/Ch4/mlir/Dialect.cpp
index 6c6cdd934cea84..34478b48b127b2 100644
--- a/mlir/examples/toy/Ch4/mlir/Dialect.cpp
+++ b/mlir/examples/toy/Ch4/mlir/Dialect.cpp
@@ -73,7 +73,7 @@ struct ToyInlinerInterface : public DialectInlinerInterface {
/// Handle the given inlined terminator(toy.return) by replacing it with a new
/// operation as necessary.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
// Only "toy.return" needs to be handled here.
auto returnOp = cast<ReturnOp>(op);
diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index 88fc033a6ab7be..6e9f829d37fbbc 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -116,7 +116,7 @@ class DialectInlinerInterface
/// calls, these are directly replaced with the operands of the `return`
/// operation). The given 'op' will be removed by the caller, after this
/// function has been called.
- virtual void handleTerminator(Operation *op,
+ virtual void handleTerminator(Operation *op, OpBuilder &builder,
ValueRange valuesToReplace) const {
llvm_unreachable(
"must implement handleTerminator in the case of one inlined block");
@@ -212,7 +212,7 @@ class InlinerInterface
//===--------------------------------------------------------------------===//
virtual void handleTerminator(Operation *op, Block *newDest) const;
- virtual void handleTerminator(Operation *op, ValueRange valuesToRepl) const;
+ virtual void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const;
virtual Value handleArgument(OpBuilder &builder, Operation *call,
Operation *callable, Value argument,
diff --git a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
index 3328d58551bff1..05a8651605aac9 100644
--- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
@@ -67,7 +67,7 @@ struct FuncInlinerInterface : public DialectInlinerInterface {
/// Handle the given inlined terminator by replacing it with a new operation
/// as necessary.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
// Only return needs to be handled here.
auto returnOp = cast<ReturnOp>(op);
diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index b3bed5ab5f412f..fbae0f99efa7ba 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -746,8 +746,16 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
/// 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.
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
+ if (isa<LLVM::UnreachableOp>(op)) {
+ for (auto dst : valuesToRepl) {
+ auto repl = builder.create<LLVM::UndefOp>(op->getLoc(), dst.getType());
+ dst.replaceAllUsesWith(repl);
+ }
+ return;
+ }
+
+ // 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/Dialect/Linalg/IR/LinalgDialect.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
index 9e50c355c50417..4a4f54db3634e8 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
@@ -57,7 +57,7 @@ struct LinalgInlinerInterface : public DialectInlinerInterface {
}
// Handle the given inlined terminator by replacing it with a new operation
// as necessary. Required when the region has only one block.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {}
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {}
};
} // namespace
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 83ae79ce482669..3fd6dc59a170dc 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -51,7 +51,7 @@ struct SCFInlinerInterface : public DialectInlinerInterface {
}
// Handle the given inlined terminator by replacing it with a new operation
// as necessary. Required when the region has only one block.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
auto retValOp = dyn_cast<scf::YieldOp>(op);
if (!retValOp)
return;
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
index 48be287ef833b2..f903a4680d5efa 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
@@ -103,7 +103,7 @@ struct SPIRVInlinerInterface : public DialectInlinerInterface {
/// Handle the given inlined terminator by replacing it with a new operation
/// as necessary.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
// Only spirv.ReturnValue needs to be handled here.
auto retValOp = dyn_cast<spirv::ReturnValueOp>(op);
if (!retValOp)
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0db097d14cd3c7..a6a053b47b2d92 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -111,11 +111,11 @@ void InlinerInterface::handleTerminator(Operation *op, Block *newDest) const {
/// Handle the given inlined terminator by replacing it with a new operation
/// as necessary.
-void InlinerInterface::handleTerminator(Operation *op,
+void InlinerInterface::handleTerminator(Operation *op, OpBuilder &builder,
ValueRange valuesToRepl) const {
auto *handler = getInterfaceFor(op);
assert(handler && "expected valid dialect handler");
- handler->handleTerminator(op, valuesToRepl);
+ handler->handleTerminator(op, builder, valuesToRepl);
}
Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
@@ -304,7 +304,7 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
firstBlockTerminator->getOperands());
// Have the interface handle the terminator of this block.
- interface.handleTerminator(firstBlockTerminator, resultsToReplace);
+ interface.handleTerminator(firstBlockTerminator, builder, resultsToReplace);
firstBlockTerminator->erase();
// Merge the post insert block into the cloned entry block.
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index edaac4da0b044c..229777763a26d6 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -676,3 +676,18 @@ 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
+ %z = llvm.call @unreachable_func(%x) : (i32) -> (i32)
+ llvm.return %z : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
index 64add8cef36986..8d2b622f2a408c 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -316,7 +316,7 @@ struct TestInlinerInterface : public DialectInlinerInterface {
/// Handle the given inlined terminator by replacing it with a new operation
/// as necessary.
- void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {
+ void handleTerminator(Operation *op, OpBuilder &builder, ValueRange valuesToRepl) const final {
// Only handle "test.return" here.
auto returnOp = dyn_cast<TestReturnOp>(op);
if (!returnOp)
|
wsmoses
force-pushed
the
mlirterm
branch
2 times, most recently
from
95be43a to
e572297
Compare
llvmbot
added
flang
|
Certainly we shouldn't have a pass assert on valid IR. How does LLVM IR behave with unreachables? |
|
LLVM is similar: https://godbolt.org/z/ncM7enGh6 |
|
LLVM seems to be putting poison, we have that too https://mlir.llvm.org/docs/Dialects/LLVM/#llvmmlirpoison-llvmpoisonop. |
LLVM seems to split the callee block into a first block terminated by an unreachable operation and a second block that returns poison. This has the advantage that the unreachable is inlined and not lost (I believe the PR deletes the unreachable?). I wonder if we could implement the same behavior? One approach may be to overwrite |
|
I changed it to also return poison instead of undefop. @gysit I could potentially modify processInlinedBlocks to return a bool whether the fast path is permitted? But I'm otherwise hesitant to have too much unreachable-specific code in the generic interface for the inliner |
|
If I am not mistaken, Note: in that solution the |
|
Yeah we can override processInlingBlocks, but I worry that creating a second fictitious block will break other parts of the infrastructure (which for example, iterate through the previously computed list). my alternate proposal is I can change processInliningBlocks to return a book indicating whether the fast path is still permitted. Unfortunately this still requires a change in code (which now returns a bool). |
|
Right, it is not completely clear to me what users are supposed todo with that block list. Changing the MLIR core inlining code has been controversial in the past. You can try that as well. I really wonder if that optimized codepath in the inliner is worth the extra complexity... |
|
The alternate proposal I have is here: #122646 (separated to a new PR for ease). It does still require an interface change, but in a moved function that no users are impacted by. |
github-actions bot
pushed a commit
to arm/arm-toolchain
that referenced
this pull request
Jan 13, 2025
…reachable (#122646) alternate option to llvm/llvm-project#122615
kazutakahirata
pushed a commit
to kazutakahirata/llvm-project
that referenced
this pull request
Jan 13, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The inliner interface (for LLVM) assumes that a single-block fn can only end in a return.....asserting otherwise.