[Mlir-commits] [mlir] [mlir][inliner] Add doClone and canHandleMultipleBlocks callbacks to Inliner (PR #131226)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Mar 17 08:07:22 PDT 2025
https://github.com/junfengd-nv updated https://github.com/llvm/llvm-project/pull/131226
>From ffea26185baf30bbc7f959fb471746918ad29b46 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd at nvidia.com>
Date: Thu, 13 Mar 2025 14:45:49 -0700
Subject: [PATCH 1/2] [mlir][inliner] Add doClone and canHandleMultipleBlocks
callbacks to the Inliner
These callback functions enhance the flexibility to customize inliner behavior.
. 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.
---
mlir/include/mlir/Transforms/Inliner.h | 30 +++++++++++-
mlir/lib/Transforms/InlinerPass.cpp | 6 ++-
mlir/lib/Transforms/Utils/Inliner.cpp | 51 +++++++++++++++------
mlir/lib/Transforms/Utils/InliningUtils.cpp | 13 ++----
4 files changed, 75 insertions(+), 25 deletions(-)
diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index ec77319d6ac88..86fa6ce6a0c77 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/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
diff --git a/mlir/lib/Transforms/InlinerPass.cpp b/mlir/lib/Transforms/InlinerPass.cpp
index 703e517d45374..fc831a5ebf4fb 100644
--- a/mlir/lib/Transforms/InlinerPass.cpp
+++ b/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()))
diff --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 756f5e379e7dd..f34fd089a8d98 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/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;
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0cae63c58ca7b..4531dd17ee3b2 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/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()),
>From 065e2109c44b5b36f703b40bc7dd2b357100628a Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd at nvidia.com>
Date: Mon, 17 Mar 2025 08:07:14 -0700
Subject: [PATCH 2/2] Update mlir/lib/Transforms/InlinerPass.cpp
Co-authored-by: jeanPerier <jean.perier.polytechnique at gmail.com>
---
mlir/lib/Transforms/InlinerPass.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/InlinerPass.cpp b/mlir/lib/Transforms/InlinerPass.cpp
index fc831a5ebf4fb..f39de35139e5e 100644
--- a/mlir/lib/Transforms/InlinerPass.cpp
+++ b/mlir/lib/Transforms/InlinerPass.cpp
@@ -142,7 +142,7 @@ void InlinerPass::runOnOperation() {
return isProfitableToInline(call, inliningThreshold);
};
- // By default, prevent inlining a functon containing multiple blocks into a
+ // By default, prevent inlining a function containing multiple blocks into a
// region that requires a single block.
auto canHandleMultipleBlocksCb = [=]() { return false; };
More information about the Mlir-commits
mailing list