[mlir][inliner] Add doClone and canHandleMultipleBlocks callbacks to Inliner Config (llvm#131226)

Current inliner disables inlining when the caller is in a region with
single block trait, while the callee function contains multiple blocks.
the SingleBlock trait is used in operations such as do/while loop, for
example fir.do_loop, fir.iterate_while and fir.if. Typically, calls within
loops are good candidates for inlining. However, functions with multiple
blocks are also common. for example, any function with "if () then
return" will result in multiple blocks in MLIR.

This change gives the flexibility of a customized inliner to handle such
cases.
doClone: clones instructions and other information from the callee
function into the caller function. .
canHandleMultipleBlocks: checks if functions with multiple blocks can be
inlined into a region with the SingleBlock trait.

The default behavior of the inliner remains unchanged.

---------

Co-authored-by: jeanPerier <jean.perier.polytechnique@gmail.com>
Co-authored-by: Mehdi Amini <joker.eph@gmail.com>

Author: 3 people

mlir/include/mlir/Transforms/Inliner.h

@@ -27,6 +27,11 @@ class InlinerConfig {
 public:
   using DefaultPipelineTy = std::function<void(OpPassManager &)>;
   using OpPipelinesTy = llvm::StringMap<OpPassManager>;
+  using CloneCallbackSigTy = void(OpBuilder &builder, Region *src,
+                                  Block *inlineBlock, Block *postInsertBlock,
+                                  IRMapping &mapper,
+                                  bool shouldCloneInlinedRegion);
+  using CloneCallbackTy = std::function<CloneCallbackSigTy>;
 
   InlinerConfig() = default;
   InlinerConfig(DefaultPipelineTy defaultPipeline,
@@ -39,13 +44,22 @@ class InlinerConfig {
   }
   const OpPipelinesTy &getOpPipelines() const { return opPipelines; }
   unsigned getMaxInliningIterations() const { return maxInliningIterations; }
+  const CloneCallbackTy &getCloneCallback() const { return cloneCallback; }
+  bool getCanHandleMultipleBlocks() const { return canHandleMultipleBlocks; }
+
   void setDefaultPipeline(DefaultPipelineTy pipeline) {
     defaultPipeline = std::move(pipeline);
   }
   void setOpPipelines(OpPipelinesTy pipelines) {
     opPipelines = std::move(pipelines);
   }
   void setMaxInliningIterations(unsigned max) { maxInliningIterations = max; }
+  void setCloneCallback(CloneCallbackTy callback) {
+    cloneCallback = std::move(callback);
+  }
+  void setCanHandleMultipleBlocks(bool value = true) {
+    canHandleMultipleBlocks = value;
+  }
 
 private:
   /// An optional function that constructs an optimization pipeline for
@@ -60,6 +74,28 @@ class InlinerConfig {
   /// For SCC-based inlining algorithms, specifies maximum number of iterations
   /// when inlining within an SCC.
   unsigned maxInliningIterations{0};
+  /// Callback for cloning operations during inlining
+  CloneCallbackTy cloneCallback = [](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());
+  };
+  /// Determine 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 it is true, cloneCallback should perform the extra
+  /// transformation. see the example in
+  /// mlir/test/lib/Transforms/TestInliningCallback.cpp
+  bool canHandleMultipleBlocks{false};
 };
 
 /// This is an implementation of the inliner

mlir/include/mlir/Transforms/InliningUtils.h

@@ -18,6 +18,7 @@
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Region.h"
 #include "mlir/IR/ValueRange.h"
+#include "mlir/Transforms/Inliner.h"
 #include <optional>
 
 namespace mlir {
@@ -253,43 +254,51 @@ class InlinerInterface
 /// provided, will be used to update the inlined operations' location
 /// information. 'shouldCloneInlinedRegion' corresponds to whether the source
 /// region should be cloned into the 'inlinePoint' or spliced directly.
-LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
-                           Operation *inlinePoint, IRMapping &mapper,
-                           ValueRange resultsToReplace,
-                           TypeRange regionResultTypes,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
-LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
-                           Block *inlineBlock, Block::iterator inlinePoint,
-                           IRMapping &mapper, ValueRange resultsToReplace,
-                           TypeRange regionResultTypes,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
+LogicalResult
+inlineRegion(InlinerInterface &interface,
+             function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+             Region *src, Operation *inlinePoint, IRMapping &mapper,
+             ValueRange resultsToReplace, TypeRange regionResultTypes,
+             std::optional<Location> inlineLoc = std::nullopt,
+             bool shouldCloneInlinedRegion = true);
+LogicalResult
+inlineRegion(InlinerInterface &interface,
+             function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+             Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+             IRMapping &mapper, ValueRange resultsToReplace,
+             TypeRange regionResultTypes,
+             std::optional<Location> inlineLoc = std::nullopt,
+             bool shouldCloneInlinedRegion = true);
 
 /// This function is an overload of the above 'inlineRegion' that allows for
 /// providing the set of operands ('inlinedOperands') that should be used
 /// in-favor of the region arguments when inlining.
-LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
-                           Operation *inlinePoint, ValueRange inlinedOperands,
-                           ValueRange resultsToReplace,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
-LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
-                           Block *inlineBlock, Block::iterator inlinePoint,
-                           ValueRange inlinedOperands,
-                           ValueRange resultsToReplace,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
+LogicalResult
+inlineRegion(InlinerInterface &interface,
+             function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+             Region *src, Operation *inlinePoint, ValueRange inlinedOperands,
+             ValueRange resultsToReplace,
+             std::optional<Location> inlineLoc = std::nullopt,
+             bool shouldCloneInlinedRegion = true);
+LogicalResult
+inlineRegion(InlinerInterface &interface,
+             function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+             Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+             ValueRange inlinedOperands, ValueRange resultsToReplace,
+             std::optional<Location> inlineLoc = std::nullopt,
+             bool shouldCloneInlinedRegion = true);
 
 /// This function inlines a given region, 'src', of a callable operation,
 /// 'callable', into the location defined by the given call operation. This
 /// function returns failure if inlining is not possible, success otherwise. On
 /// failure, no changes are made to the module. 'shouldCloneInlinedRegion'
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
-LogicalResult inlineCall(InlinerInterface &interface, CallOpInterface call,
-                         CallableOpInterface callable, Region *src,
-                         bool shouldCloneInlinedRegion = true);
+LogicalResult
+inlineCall(InlinerInterface &interface,
+           function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+           CallOpInterface call, CallableOpInterface callable, Region *src,
+           bool shouldCloneInlinedRegion = true);
 
 } // namespace mlir
 

mlir/lib/Transforms/Utils/Inliner.cpp

@@ -652,7 +652,7 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
     bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
 
     LogicalResult inlineResult =
-        inlineCall(inlinerIface, call,
+        inlineCall(inlinerIface, inliner.config.getCloneCallback(), call,
                    cast<CallableOpInterface>(targetRegion->getParentOp()),
                    targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
     if (failed(inlineResult)) {
@@ -730,19 +730,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.config.getCanHandleMultipleBlocks()) {
+    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"
@@ -266,10 +267,11 @@ static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
 }
 
 static LogicalResult
-inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
-                 Block::iterator inlinePoint, IRMapping &mapper,
-                 ValueRange resultsToReplace, TypeRange regionResultTypes,
-                 std::optional<Location> inlineLoc,
+inlineRegionImpl(InlinerInterface &interface,
+                 function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+                 Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+                 IRMapping &mapper, ValueRange resultsToReplace,
+                 TypeRange regionResultTypes, std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, CallOpInterface call = {}) {
   assert(resultsToReplace.size() == regionResultTypes.size());
   // We expect the region to have at least one block.
@@ -296,16 +298,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());
+  cloneCallback(builder, src, inlineBlock, postInsertBlock, mapper,
+                shouldCloneInlinedRegion);
 
   // Get the range of newly inserted blocks.
   auto newBlocks = llvm::make_range(std::next(inlineBlock->getIterator()),
@@ -374,9 +370,11 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
 }
 
 static LogicalResult
-inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
-                 Block::iterator inlinePoint, ValueRange inlinedOperands,
-                 ValueRange resultsToReplace, std::optional<Location> inlineLoc,
+inlineRegionImpl(InlinerInterface &interface,
+                 function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+                 Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+                 ValueRange inlinedOperands, ValueRange resultsToReplace,
+                 std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, CallOpInterface call = {}) {
   // We expect the region to have at least one block.
   if (src->empty())
@@ -398,53 +396,54 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   }
 
   // Call into the main region inliner function.
-  return inlineRegionImpl(interface, src, inlineBlock, inlinePoint, mapper,
-                          resultsToReplace, resultsToReplace.getTypes(),
-                          inlineLoc, shouldCloneInlinedRegion, call);
+  return inlineRegionImpl(interface, cloneCallback, src, inlineBlock,
+                          inlinePoint, mapper, resultsToReplace,
+                          resultsToReplace.getTypes(), inlineLoc,
+                          shouldCloneInlinedRegion, call);
 }
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Operation *inlinePoint, IRMapping &mapper,
-                                 ValueRange resultsToReplace,
-                                 TypeRange regionResultTypes,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, src, inlinePoint->getBlock(),
+LogicalResult mlir::inlineRegion(
+    InlinerInterface &interface,
+    function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback, Region *src,
+    Operation *inlinePoint, IRMapping &mapper, ValueRange resultsToReplace,
+    TypeRange regionResultTypes, std::optional<Location> inlineLoc,
+    bool shouldCloneInlinedRegion) {
+  return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),
                       ++inlinePoint->getIterator(), mapper, resultsToReplace,
                       regionResultTypes, inlineLoc, shouldCloneInlinedRegion);
 }
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Block *inlineBlock,
-                                 Block::iterator inlinePoint, IRMapping &mapper,
-                                 ValueRange resultsToReplace,
-                                 TypeRange regionResultTypes,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
-  return inlineRegionImpl(interface, src, inlineBlock, inlinePoint, mapper,
-                          resultsToReplace, regionResultTypes, inlineLoc,
-                          shouldCloneInlinedRegion);
+
+LogicalResult mlir::inlineRegion(
+    InlinerInterface &interface,
+    function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback, Region *src,
+    Block *inlineBlock, Block::iterator inlinePoint, IRMapping &mapper,
+    ValueRange resultsToReplace, TypeRange regionResultTypes,
+    std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(
+      interface, cloneCallback, src, inlineBlock, inlinePoint, mapper,
+      resultsToReplace, regionResultTypes, inlineLoc, shouldCloneInlinedRegion);
 }
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Operation *inlinePoint,
-                                 ValueRange inlinedOperands,
-                                 ValueRange resultsToReplace,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, src, inlinePoint->getBlock(),
+LogicalResult mlir::inlineRegion(
+    InlinerInterface &interface,
+    function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback, Region *src,
+    Operation *inlinePoint, ValueRange inlinedOperands,
+    ValueRange resultsToReplace, std::optional<Location> inlineLoc,
+    bool shouldCloneInlinedRegion) {
+  return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),
                       ++inlinePoint->getIterator(), inlinedOperands,
                       resultsToReplace, inlineLoc, shouldCloneInlinedRegion);
 }
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Block *inlineBlock,
-                                 Block::iterator inlinePoint,
-                                 ValueRange inlinedOperands,
-                                 ValueRange resultsToReplace,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
-  return inlineRegionImpl(interface, src, inlineBlock, inlinePoint,
-                          inlinedOperands, resultsToReplace, inlineLoc,
-                          shouldCloneInlinedRegion);
+
+LogicalResult mlir::inlineRegion(
+    InlinerInterface &interface,
+    function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback, Region *src,
+    Block *inlineBlock, Block::iterator inlinePoint, ValueRange inlinedOperands,
+    ValueRange resultsToReplace, std::optional<Location> inlineLoc,
+    bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(interface, cloneCallback, src, inlineBlock,
+                          inlinePoint, inlinedOperands, resultsToReplace,
+                          inlineLoc, shouldCloneInlinedRegion);
 }
 
 /// Utility function used to generate a cast operation from the given interface,
@@ -475,10 +474,11 @@ static Value materializeConversion(const DialectInlinerInterface *interface,
 /// failure, no changes are made to the module. 'shouldCloneInlinedRegion'
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
-LogicalResult mlir::inlineCall(InlinerInterface &interface,
-                               CallOpInterface call,
-                               CallableOpInterface callable, Region *src,
-                               bool shouldCloneInlinedRegion) {
+LogicalResult
+mlir::inlineCall(InlinerInterface &interface,
+                 function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+                 CallOpInterface call, CallableOpInterface callable,
+                 Region *src, bool shouldCloneInlinedRegion) {
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
@@ -552,7 +552,7 @@ LogicalResult mlir::inlineCall(InlinerInterface &interface,
     return cleanupState();
 
   // Attempt to inline the call.
-  if (failed(inlineRegionImpl(interface, src, call->getBlock(),
+  if (failed(inlineRegionImpl(interface, cloneCallback, src, call->getBlock(),
                               ++call->getIterator(), mapper, callResults,
                               callableResultTypes, call.getLoc(),
                               shouldCloneInlinedRegion, call)))