[mlir][inliner] Add doClone and canHandleMultipleBlocks callbacks to Inliner

mlir/include/mlir/Transforms/Inliner.h

@@ -90,18 +90,40 @@ class Inliner {
   /// this hook's interface might need to be extended in future.
   using ProfitabilityCallbackTy = std::function<bool(const ResolvedCall &)>;
 
+  /// Type of the callback that determines if the inliner can inline a function
+  /// containing multiple blocks into a region that requires a single block. By
+  /// default, it is not allowed.
+  /// If this function return true, the static member function doClone()
+  /// should perform the actual transformation with its support.
+  using canHandleMultipleBlocksCbTy = std::function<bool()>;
+
+  using CloneCallbackTy =
+      std::function<void(OpBuilder &builder, Region *src, Block *inlineBlock,
+                         Block *postInsertBlock, IRMapping &mapper,
+                         bool shouldCloneInlinedRegion)>;
+
   Inliner(Operation *op, CallGraph &cg, Pass &pass, AnalysisManager am,
           RunPipelineHelperTy runPipelineHelper, const InlinerConfig &config,
-          ProfitabilityCallbackTy isProfitableToInline)
+          ProfitabilityCallbackTy isProfitableToInline,
+          canHandleMultipleBlocksCbTy canHandleMultipleBlocks)
       : op(op), cg(cg), pass(pass), am(am),
         runPipelineHelper(std::move(runPipelineHelper)), config(config),
-        isProfitableToInline(std::move(isProfitableToInline)) {}
+        isProfitableToInline(std::move(isProfitableToInline)),
+        canHandleMultipleBlocks(std::move(canHandleMultipleBlocks)) {}
   Inliner(Inliner &) = delete;
   void operator=(const Inliner &) = delete;
 
   /// Perform inlining on a OpTrait::SymbolTable operation.
   LogicalResult doInlining();
 
+  /// This function provides a callback mechanism to clone the instructions and
+  /// other information from the callee function into the caller function.
+  static CloneCallbackTy &doClone();
+
+  /// Set the clone callback function.
+  /// The provided function "func" will be invoked by Inliner::doClone().
+  void setCloneCallback(CloneCallbackTy func) { doClone() = func; }
+
 private:
   /// An OpTrait::SymbolTable operation to run the inlining on.
   Operation *op;
@@ -119,10 +141,14 @@ class Inliner {
   /// Returns true, if it is profitable to inline the callable operation
   /// at the call site.
   ProfitabilityCallbackTy isProfitableToInline;
+  /// Return true, if functions with multiple blocks can be inlined
+  /// into a region with the SingleBlock trait.
+  canHandleMultipleBlocksCbTy canHandleMultipleBlocks;
 
   /// Forward declaration of the class providing the actual implementation.
   class Impl;
 };
+
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_INLINER_H

mlir/lib/Transforms/InlinerPass.cpp

@@ -142,9 +142,13 @@ void InlinerPass::runOnOperation() {
     return isProfitableToInline(call, inliningThreshold);
   };
 
+  // By default, prevent inlining a functon containing multiple blocks into a
+  // region that requires a single block.
+  auto canHandleMultipleBlocksCb = [=]() { return false; };
+
   // Get an instance of the inliner.
   Inliner inliner(op, cg, *this, getAnalysisManager(), runPipelineHelper,
-                  config, profitabilityCb);
+                  config, profitabilityCb, canHandleMultipleBlocksCb);
 
   // Run the inlining.
   if (failed(inliner.doInlining()))

mlir/lib/Transforms/Utils/Inliner.cpp

@@ -344,6 +344,28 @@ static void collectCallOps(iterator_range<Region::iterator> blocks,
     }
   }
 }
+//===----------------------------------------------------------------------===//
+// Inliner
+//===----------------------------------------------------------------------===//
+// Initialize doClone function with the default implementation
+Inliner::CloneCallbackTy &Inliner::doClone() {
+  static Inliner::CloneCallbackTy doWork =
+      [](OpBuilder &builder, Region *src, Block *inlineBlock,
+         Block *postInsertBlock, IRMapping &mapper,
+         bool shouldCloneInlinedRegion) {
+        // Check to see if the region is being cloned, or moved inline. In
+        // either case, move the new blocks after the 'insertBlock' to improve
+        // IR readability.
+        Region *insertRegion = inlineBlock->getParent();
+        if (shouldCloneInlinedRegion)
+          src->cloneInto(insertRegion, postInsertBlock->getIterator(), mapper);
+        else
+          insertRegion->getBlocks().splice(postInsertBlock->getIterator(),
+                                           src->getBlocks(), src->begin(),
+                                           src->end());
+      };
+  return doWork;
+}
 
 //===----------------------------------------------------------------------===//
 // InlinerInterfaceImpl
@@ -729,19 +751,22 @@ bool Inliner::Impl::shouldInline(ResolvedCall &resolvedCall) {
 
   // Don't allow inlining if the callee has multiple blocks (unstructured
   // control flow) but we cannot be sure that the caller region supports that.
-  bool calleeHasMultipleBlocks =
-      llvm::hasNItemsOrMore(*callableRegion, /*N=*/2);
-  // If both parent ops have the same type, it is safe to inline. Otherwise,
-  // decide based on whether the op has the SingleBlock trait or not.
-  // Note: This check does currently not account for SizedRegion/MaxSizedRegion.
-  auto callerRegionSupportsMultipleBlocks = [&]() {
-    return callableRegion->getParentOp()->getName() ==
-               resolvedCall.call->getParentOp()->getName() ||
-           !resolvedCall.call->getParentOp()
-                ->mightHaveTrait<OpTrait::SingleBlock>();
-  };
-  if (calleeHasMultipleBlocks && !callerRegionSupportsMultipleBlocks())
-    return false;
+  if (!inliner.canHandleMultipleBlocks()) {
+    bool calleeHasMultipleBlocks =
+        llvm::hasNItemsOrMore(*callableRegion, /*N=*/2);
+    // If both parent ops have the same type, it is safe to inline. Otherwise,
+    // decide based on whether the op has the SingleBlock trait or not.
+    // Note: This check does currently not account for
+    // SizedRegion/MaxSizedRegion.
+    auto callerRegionSupportsMultipleBlocks = [&]() {
+      return callableRegion->getParentOp()->getName() ==
+                 resolvedCall.call->getParentOp()->getName() ||
+             !resolvedCall.call->getParentOp()
+                  ->mightHaveTrait<OpTrait::SingleBlock>();
+    };
+    if (calleeHasMultipleBlocks && !callerRegionSupportsMultipleBlocks())
+      return false;
+  }
 
   if (!inliner.isProfitableToInline(resolvedCall))
     return false;

mlir/lib/Transforms/Utils/InliningUtils.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/InliningUtils.h"
+#include "mlir/Transforms/Inliner.h"
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IRMapping.h"
@@ -275,16 +276,10 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   if (call && callable)
     handleArgumentImpl(interface, builder, call, callable, mapper);
 
-  // Check to see if the region is being cloned, or moved inline. In either
-  // case, move the new blocks after the 'insertBlock' to improve IR
-  // readability.
+  // Clone the callee's source into the caller.
   Block *postInsertBlock = inlineBlock->splitBlock(inlinePoint);
-  if (shouldCloneInlinedRegion)
-    src->cloneInto(insertRegion, postInsertBlock->getIterator(), mapper);
-  else
-    insertRegion->getBlocks().splice(postInsertBlock->getIterator(),
-                                     src->getBlocks(), src->begin(),
-                                     src->end());
+  Inliner::doClone()(builder, src, inlineBlock, postInsertBlock, mapper,
+                     shouldCloneInlinedRegion);
 
   // Get the range of newly inserted blocks.
   auto newBlocks = llvm::make_range(std::next(inlineBlock->getIterator()),