Add method getInlineBlockAndPoint for DialectInlinerInterface

Author: nikita-kud

mlir/include/mlir/Transforms/InliningUtils.h

@@ -172,6 +172,14 @@ class DialectInlinerInterface
     return result;
   }
 
+  /// Hook to cleanup IR before erase call op
+  virtual void eraseCall(Operation *call) const;
+
+  /// Hook to get proper place where callable region will be inlined
+  /// By default returns block of the call operation 
+  virtual std::tuple<Block*, Block::iterator> 
+            getInlineBlockAndPoint(Operation *call) const;
+
   /// Process a set of blocks that have been inlined for a call. This callback
   /// is invoked before inlined terminator operations have been processed.
   virtual void processInlinedCallBlocks(

mlir/lib/Dialect/Arith/IR/CMakeLists.txt

@@ -22,6 +22,7 @@ add_mlir_dialect_library(MLIRArithDialect
   MLIRArithOpsInterfacesIncGen
 
   LINK_LIBS PUBLIC
+  MLIRTransformUtils
   MLIRCastInterfaces
   MLIRDialect
   MLIRInferIntRangeCommon

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

@@ -28,6 +28,7 @@ add_mlir_dialect_library(MLIRLLVMDialect
   Core
 
   LINK_LIBS PUBLIC
+  MLIRTransformUtils
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces

mlir/lib/Transforms/Utils/Inliner.cpp

@@ -688,7 +688,8 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
     useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
 
     // then erase the call.
-    call.erase();
+    const auto *callInterface = inlinerIface.getInterfaceFor(call->getDialect());
+    callInterface->eraseCall(call);
 
     // If we inlined in place, mark the node for deletion.
     if (inlineInPlace) {

mlir/lib/Transforms/Utils/InliningUtils.cpp

@@ -69,6 +69,19 @@ static void remapInlinedOperands(iterator_range<Region::iterator> inlinedBlocks,
     block.walk(remapOperands);
 }
 
+//===----------------------------------------------------------------------===//
+// DialectInlinerInterface
+//===----------------------------------------------------------------------===//
+
+void DialectInlinerInterface::eraseCall(Operation *call) const {
+  call->erase();
+}
+
+std::tuple<Block*, Block::iterator> 
+      DialectInlinerInterface::getInlineBlockAndPoint(Operation *call) const {
+  return std::make_tuple(call->getBlock(), std::next(call->getIterator()));
+}
+
 //===----------------------------------------------------------------------===//
 // InlinerInterface
 //===----------------------------------------------------------------------===//
@@ -530,9 +543,11 @@ LogicalResult mlir::inlineCall(InlinerInterface &interface,
   if (!interface.isLegalToInline(call, callable, shouldCloneInlinedRegion))
     return cleanupState();
 
+  auto [inlineBlock, inlinePoint] = callInterface->getInlineBlockAndPoint(call);
+
   // Attempt to inline the call.
-  if (failed(inlineRegionImpl(interface, src, call->getBlock(),
-                              ++call->getIterator(), mapper, callResults,
+  if (failed(inlineRegionImpl(interface, src, inlineBlock,
+                              inlinePoint, mapper, callResults,
                               callableResultTypes, call.getLoc(),
                               shouldCloneInlinedRegion, call)))
     return cleanupState();