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

William Moses via flang-commits flang-commits at lists.llvm.org
Sun Jan 12 10:32:29 PST 2025


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

>From 675c1b71cddee5afa16dc2f348029c0d85b3d228 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

---
 flang/lib/Optimizer/Dialect/FIRDialect.cpp        |  2 +-
 mlir/examples/toy/Ch4/mlir/Dialect.cpp            |  3 ++-
 mlir/examples/toy/Ch5/mlir/Dialect.cpp            |  3 ++-
 mlir/examples/toy/Ch6/mlir/Dialect.cpp            |  3 ++-
 mlir/examples/toy/Ch7/mlir/Dialect.cpp            |  3 ++-
 mlir/include/mlir/Transforms/InliningUtils.h      |  5 +++--
 .../Dialect/Func/Extensions/InlinerExtension.cpp  |  3 ++-
 .../LLVMIR/Transforms/InlinerInterfaceImpl.cpp    | 13 +++++++++++--
 mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp      |  3 ++-
 mlir/lib/Dialect/SCF/IR/SCF.cpp                   |  3 ++-
 mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp        |  3 ++-
 mlir/lib/Transforms/Utils/InliningUtils.cpp       |  6 +++---
 mlir/test/Dialect/LLVMIR/inlining.mlir            | 15 +++++++++++++++
 .../lib/Dialect/Test/TestDialectInterfaces.cpp    |  3 ++-
 14 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/FIRDialect.cpp b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
index 4b1dadaac6728a..fa0c9c33cf8e61 100644
--- a/flang/lib/Optimizer/Dialect/FIRDialect.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
@@ -44,7 +44,7 @@ struct FIRInlinerInterface : public mlir::DialectInlinerInterface {
   /// We handle the return (a Fortran FUNCTION) by replacing the values
   /// previously returned by the call operation with the operands of the
   /// return.
-  void handleTerminator(mlir::Operation *op,
+  void handleTerminator(mlir::Operation *op, mlir::OpBuilder &builder,
                         mlir::ValueRange valuesToRepl) const final {
     auto returnOp = llvm::cast<mlir::func::ReturnOp>(op);
     assert(returnOp.getNumOperands() == valuesToRepl.size());
diff --git a/mlir/examples/toy/Ch4/mlir/Dialect.cpp b/mlir/examples/toy/Ch4/mlir/Dialect.cpp
index 6c6cdd934cea84..e543ee4a574496 100644
--- a/mlir/examples/toy/Ch4/mlir/Dialect.cpp
+++ b/mlir/examples/toy/Ch4/mlir/Dialect.cpp
@@ -73,7 +73,8 @@ 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/examples/toy/Ch5/mlir/Dialect.cpp b/mlir/examples/toy/Ch5/mlir/Dialect.cpp
index 72072f9188bf31..1412a1479bca8a 100644
--- a/mlir/examples/toy/Ch5/mlir/Dialect.cpp
+++ b/mlir/examples/toy/Ch5/mlir/Dialect.cpp
@@ -73,7 +73,8 @@ 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/examples/toy/Ch6/mlir/Dialect.cpp b/mlir/examples/toy/Ch6/mlir/Dialect.cpp
index 72072f9188bf31..1412a1479bca8a 100644
--- a/mlir/examples/toy/Ch6/mlir/Dialect.cpp
+++ b/mlir/examples/toy/Ch6/mlir/Dialect.cpp
@@ -73,7 +73,8 @@ 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/examples/toy/Ch7/mlir/Dialect.cpp b/mlir/examples/toy/Ch7/mlir/Dialect.cpp
index 7e030ffc5488c9..dab5617f18b367 100644
--- a/mlir/examples/toy/Ch7/mlir/Dialect.cpp
+++ b/mlir/examples/toy/Ch7/mlir/Dialect.cpp
@@ -79,7 +79,8 @@ 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..b5bab7c16adfa9 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,8 @@ 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..2b173ffba3cc2c 100644
--- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
@@ -67,7 +67,8 @@ 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..0a264de4213432 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -746,8 +746,17 @@ 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::PoisonOp>(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..341ba464ba96f3 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
@@ -57,7 +57,8 @@ 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..29bd73b4bed69d 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -51,7 +51,8 @@ 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..8c52da4a209cad 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
@@ -103,7 +103,8 @@ 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..bebd245fffd68b 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -316,7 +316,8 @@ 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)



More information about the flang-commits mailing list