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

Author: wsmoses

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;
 

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.

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.

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);

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
+}