From ffea26185baf30bbc7f959fb471746918ad29b46 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Thu, 13 Mar 2025 14:45:49 -0700
Subject: [PATCH 01/15] [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@nvidia.com>
Date: Mon, 17 Mar 2025 08:07:14 -0700
Subject: [PATCH 02/15] Update mlir/lib/Transforms/InlinerPass.cpp

Co-authored-by: jeanPerier <jean.perier.polytechnique@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; };
 

From ee862e69204423fbe35c84d4fe8a821ee113bdf1 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Thu, 20 Mar 2025 16:56:34 -0700
Subject: [PATCH 03/15] Add a test for Inliner callback functions.

---
 .../Transforms/test-inlining-callback.mlir    |  24 +++
 mlir/test/lib/Transforms/CMakeLists.txt       |   1 +
 .../lib/Transforms/TestInliningCallback.cpp   | 152 ++++++++++++++++++
 mlir/tools/mlir-opt/mlir-opt.cpp              |   2 +
 4 files changed, 179 insertions(+)
 create mode 100644 mlir/test/Transforms/test-inlining-callback.mlir
 create mode 100644 mlir/test/lib/Transforms/TestInliningCallback.cpp

diff --git a/mlir/test/Transforms/test-inlining-callback.mlir b/mlir/test/Transforms/test-inlining-callback.mlir
new file mode 100644
index 0000000000000..c012c31e7e490
--- /dev/null
+++ b/mlir/test/Transforms/test-inlining-callback.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s -test-inline-callback | FileCheck %s
+
+// Test inlining with multiple blocks and scf.execute_region transformation
+// CHECK-LABEL: func @test_inline_multiple_blocks
+func.func @test_inline_multiple_blocks(%arg0: i32) -> i32 {
+  // CHECK: %[[RES:.*]] = scf.execute_region -> i32
+  // CHECK-NEXT: %[[ADD1:.*]] = arith.addi %arg0, %arg0
+  // CHECK-NEXT: cf.br ^bb1(%[[ADD1]] : i32)
+  // CHECK: ^bb1(%[[ARG:.*]]: i32):
+  // CHECK-NEXT: %[[ADD2:.*]] = arith.addi %[[ARG]], %[[ARG]]
+  // CHECK-NEXT: scf.yield %[[ADD2]]
+  // CHECK: return %[[RES]]
+  %fn = "test.functional_region_op"() ({
+  ^bb0(%a : i32):
+    %b = arith.addi %a, %a : i32
+    cf.br ^bb1(%b: i32)
+  ^bb1(%c: i32):
+    %d = arith.addi %c, %c : i32
+    "test.return"(%d) : (i32) -> ()
+  }) : () -> ((i32) -> i32)
+
+  %0 = call_indirect %fn(%arg0) : (i32) -> i32
+  return %0 : i32
+}
diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt
index c053fd4b20473..76041cd6cd791 100644
--- a/mlir/test/lib/Transforms/CMakeLists.txt
+++ b/mlir/test/lib/Transforms/CMakeLists.txt
@@ -29,6 +29,7 @@ add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestControlFlowSink.cpp
   TestInlining.cpp
+  TestInliningCallback.cpp
   TestMakeIsolatedFromAbove.cpp
   TestTransformsOps.cpp
   ${MLIRTestTransformsPDLSrc}
diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
new file mode 100644
index 0000000000000..738e861ccdad1
--- /dev/null
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -0,0 +1,152 @@
+//===- TestInliningCallback.cpp - Pass to inline calls in the test dialect
+//--------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// This file implements a pass to test inlining callbacks including
+// canHandleMultipleBlocks and doClone.
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "TestOps.h"
+#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Inliner.h"
+#include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/StringSet.h"
+
+using namespace mlir;
+using namespace test;
+
+namespace {
+struct InlinerCallback
+    : public PassWrapper<InlinerCallback, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InlinerCallback)
+
+  StringRef getArgument() const final { return "test-inline-callback"; }
+  StringRef getDescription() const final {
+    return "Test inlining region calls with call back functions";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<scf::SCFDialect>();
+  }
+
+  static LogicalResult runPipelineHelper(Pass &pass, OpPassManager &pipeline,
+                                         Operation *op) {
+    return mlir::cast<InlinerCallback>(pass).runPipeline(pipeline, op);
+  }
+
+  // Customize the implementation of Inliner::doClone
+  // Wrap the callee into scf.execute_region operation
+  static void testDoClone(OpBuilder &builder, Region *src, Block *inlineBlock,
+                          Block *postInsertBlock, IRMapping &mapper,
+                          bool shouldCloneInlinedRegion) {
+    // Create a new scf.execute_region operation
+    mlir::Operation &call = inlineBlock->back();
+    builder.setInsertionPointAfter(&call);
+
+    auto executeRegionOp = builder.create<mlir::scf::ExecuteRegionOp>(
+        call.getLoc(), call.getResultTypes());
+    mlir::Region &region = executeRegionOp.getRegion();
+
+    // Move the inlined blocks into the region
+    src->cloneInto(&region, mapper);
+
+    // Split block before scf operation.
+    Block *continueBlock =
+        inlineBlock->splitBlock(executeRegionOp.getOperation());
+
+    // Replace all test.return with scf.yield
+    for (mlir::Block &block : region) {
+
+      for (mlir::Operation &op : llvm::make_early_inc_range(block)) {
+        if (test::TestReturnOp returnOp =
+                llvm::dyn_cast<test::TestReturnOp>(&op)) {
+          mlir::OpBuilder returnBuilder(returnOp);
+          returnBuilder.create<mlir::scf::YieldOp>(returnOp.getLoc(),
+                                                   returnOp.getOperands());
+          returnOp.erase();
+        }
+      }
+    }
+
+    // Add test.return after scf.execute_region
+    builder.setInsertionPointAfter(executeRegionOp);
+    builder.create<test::TestReturnOp>(executeRegionOp.getLoc(),
+                                       executeRegionOp.getResults());
+  }
+
+  void runOnOperation() override {
+    InlinerConfig config;
+    CallGraph &cg = getAnalysis<CallGraph>();
+
+    auto function = getOperation();
+
+    // By default, assume that any inlining is profitable.
+    auto profitabilityCb = [&](const mlir::Inliner::ResolvedCall &call) {
+      return true;
+    };
+
+    // This customized inliner can turn multiple blocks into a single block.
+    auto canHandleMultipleBlocksCb = [&]() { return true; };
+
+    // Get an instance of the inliner.
+    Inliner inliner(function, cg, *this, getAnalysisManager(),
+                    runPipelineHelper, config, profitabilityCb,
+                    canHandleMultipleBlocksCb);
+
+    // Customize the implementation of Inliner::doClone
+    inliner.setCloneCallback([](OpBuilder &builder, Region *src,
+                                Block *inlineBlock, Block *postInsertBlock,
+                                IRMapping &mapper,
+                                bool shouldCloneInlinedRegion) {
+      return testDoClone(builder, src, inlineBlock, postInsertBlock, mapper,
+                         shouldCloneInlinedRegion);
+    });
+
+    // Collect each of the direct function calls within the module.
+    SmallVector<func::CallIndirectOp, 16> callers;
+    function.walk(
+        [&](func::CallIndirectOp caller) { callers.push_back(caller); });
+
+    // Build the inliner interface.
+    InlinerInterface interface(&getContext());
+
+    // Try to inline each of the call operations.
+    for (auto caller : callers) {
+      auto callee = dyn_cast_or_null<FunctionalRegionOp>(
+          caller.getCallee().getDefiningOp());
+      if (!callee)
+        continue;
+
+      // Inline the functional region operation, but only clone the internal
+      // region if there is more than one use.
+      if (failed(inlineRegion(
+              interface, &callee.getBody(), caller, caller.getArgOperands(),
+              caller.getResults(), caller.getLoc(),
+              /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
+        continue;
+
+      // If the inlining was successful then erase the call and callee if
+      // possible.
+      caller.erase();
+      if (callee.use_empty())
+        callee.erase();
+    }
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerInlinerCallback() { PassRegistration<InlinerCallback>(); }
+} // namespace test
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index d06ff8070e7cf..ca4706e96787f 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -73,6 +73,7 @@ void registerCommutativityUtils();
 void registerConvertCallOpPass();
 void registerConvertFuncOpPass();
 void registerInliner();
+void registerInlinerCallback();
 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
@@ -215,6 +216,7 @@ void registerTestPasses() {
   mlir::test::registerConvertCallOpPass();
   mlir::test::registerConvertFuncOpPass();
   mlir::test::registerInliner();
+  mlir::test::registerInlinerCallback();
   mlir::test::registerMemRefBoundCheck();
   mlir::test::registerPatternsTestPass();
   mlir::test::registerSimpleParametricTilingPass();

From 86cf37f3326540f424486df4ac16ce9de10e7a7b Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Wed, 2 Apr 2025 17:54:19 -0700
Subject: [PATCH 04/15] Move to InlinerConfig

---
 mlir/include/mlir/Transforms/Inliner.h        | 40 ++++++++--
 mlir/include/mlir/Transforms/InliningUtils.h  | 16 ++--
 mlir/lib/Transforms/InlinerPass.cpp           |  6 +-
 mlir/lib/Transforms/Utils/Inliner.cpp         | 26 +------
 mlir/lib/Transforms/Utils/InliningUtils.cpp   | 77 ++++++++++---------
 mlir/test/lib/Transforms/TestInlining.cpp     | 15 ++--
 .../lib/Transforms/TestInliningCallback.cpp   | 29 ++++---
 7 files changed, 111 insertions(+), 98 deletions(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 86fa6ce6a0c77..84c96230f98f7 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -27,6 +27,10 @@ class InlinerConfig {
 public:
   using DefaultPipelineTy = std::function<void(OpPassManager &)>;
   using OpPipelinesTy = llvm::StringMap<OpPassManager>;
+  using CloneCallbackTy =
+      std::function<void(OpBuilder &builder, Region *src, Block *inlineBlock,
+                         Block *postInsertBlock, IRMapping &mapper,
+                         bool shouldCloneInlinedRegion)>;
 
   InlinerConfig() = default;
   InlinerConfig(DefaultPipelineTy defaultPipeline,
@@ -39,6 +43,9 @@ 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);
   }
@@ -46,6 +53,12 @@ class InlinerConfig {
     opPipelines = std::move(pipelines);
   }
   void setMaxInliningIterations(unsigned max) { maxInliningIterations = max; }
+  void setCloneCallback(CloneCallbackTy callback) {
+    cloneCallback = std::move(callback);
+  }
+  void setCanHandleMultipleBlocks(bool value) {
+    canHandleMultipleBlocks = value;
+  }
 
 private:
   /// An optional function that constructs an optimization pipeline for
@@ -60,6 +73,24 @@ 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());
+  };
+  /// Determining if multiple blocks can be handled
+  bool canHandleMultipleBlocks{false};
 };
 
 /// This is an implementation of the inliner
@@ -104,12 +135,10 @@ class Inliner {
 
   Inliner(Operation *op, CallGraph &cg, Pass &pass, AnalysisManager am,
           RunPipelineHelperTy runPipelineHelper, const InlinerConfig &config,
-          ProfitabilityCallbackTy isProfitableToInline,
-          canHandleMultipleBlocksCbTy canHandleMultipleBlocks)
+          ProfitabilityCallbackTy isProfitableToInline)
       : op(op), cg(cg), pass(pass), am(am),
         runPipelineHelper(std::move(runPipelineHelper)), config(config),
-        isProfitableToInline(std::move(isProfitableToInline)),
-        canHandleMultipleBlocks(std::move(canHandleMultipleBlocks)) {}
+        isProfitableToInline(std::move(isProfitableToInline)) {}
   Inliner(Inliner &) = delete;
   void operator=(const Inliner &) = delete;
 
@@ -141,9 +170,6 @@ 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;
diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index becfe9b047ef4..e14b2412f7853 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/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,13 +254,15 @@ 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,
+LogicalResult inlineRegion(InlinerInterface &interface,
+                           const InlinerConfig &config, 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,
+LogicalResult inlineRegion(InlinerInterface &interface,
+                           const InlinerConfig &config, Region *src,
                            Block *inlineBlock, Block::iterator inlinePoint,
                            IRMapping &mapper, ValueRange resultsToReplace,
                            TypeRange regionResultTypes,
@@ -269,12 +272,14 @@ LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
 /// 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,
+LogicalResult inlineRegion(InlinerInterface &interface,
+                           const InlinerConfig &config, Region *src,
                            Operation *inlinePoint, ValueRange inlinedOperands,
                            ValueRange resultsToReplace,
                            std::optional<Location> inlineLoc = std::nullopt,
                            bool shouldCloneInlinedRegion = true);
-LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
+LogicalResult inlineRegion(InlinerInterface &interface,
+                           const InlinerConfig &config, Region *src,
                            Block *inlineBlock, Block::iterator inlinePoint,
                            ValueRange inlinedOperands,
                            ValueRange resultsToReplace,
@@ -287,7 +292,8 @@ LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
 /// 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,
+LogicalResult inlineCall(InlinerInterface &interface,
+                         const InlinerConfig &config, CallOpInterface call,
                          CallableOpInterface callable, Region *src,
                          bool shouldCloneInlinedRegion = true);
 
diff --git a/mlir/lib/Transforms/InlinerPass.cpp b/mlir/lib/Transforms/InlinerPass.cpp
index f39de35139e5e..703e517d45374 100644
--- a/mlir/lib/Transforms/InlinerPass.cpp
+++ b/mlir/lib/Transforms/InlinerPass.cpp
@@ -142,13 +142,9 @@ void InlinerPass::runOnOperation() {
     return isProfitableToInline(call, inliningThreshold);
   };
 
-  // By default, prevent inlining a function 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, canHandleMultipleBlocksCb);
+                  config, profitabilityCb);
 
   // Run the inlining.
   if (failed(inliner.doInlining()))
diff --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index f34fd089a8d98..445d4bd5efc75 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -344,28 +344,6 @@ 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
@@ -673,7 +651,7 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
     bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
 
     LogicalResult inlineResult =
-        inlineCall(inlinerIface, call,
+        inlineCall(inlinerIface, inliner.config, call,
                    cast<CallableOpInterface>(targetRegion->getParentOp()),
                    targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
     if (failed(inlineResult)) {
@@ -751,7 +729,7 @@ 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.
-  if (!inliner.canHandleMultipleBlocks()) {
+  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,
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 4531dd17ee3b2..1e5d62cec0b3d 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -32,11 +32,11 @@ static void
 remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
                       Location callerLoc) {
   DenseMap<Location, LocationAttr> mappedLocations;
-  auto remapLoc = [&](Location loc) {
+    LocationAttr newLoc = stackLocations(loc, callerLoc);
     auto [it, inserted] = mappedLocations.try_emplace(loc);
     // Only query the attribute uniquer once per callsite attribute.
     if (inserted) {
-      auto newLoc = CallSiteLoc::get(loc, callerLoc);
+      LocationAttr newLoc = stackLocations(loc, callerLoc);
       it->getSecond() = newLoc;
     }
     return it->second;
@@ -246,10 +246,10 @@ 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, const InlinerConfig &config,
+                 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.
@@ -278,8 +278,8 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
 
   // Clone the callee's source into the caller.
   Block *postInsertBlock = inlineBlock->splitBlock(inlinePoint);
-  Inliner::doClone()(builder, src, inlineBlock, postInsertBlock, mapper,
-                     shouldCloneInlinedRegion);
+  config.getCloneCallback()(builder, src, inlineBlock, postInsertBlock, mapper,
+                            shouldCloneInlinedRegion);
 
   // Get the range of newly inserted blocks.
   auto newBlocks = llvm::make_range(std::next(inlineBlock->getIterator()),
@@ -304,7 +304,7 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   bool singleBlockFastPath = interface.allowSingleBlockOptimization(newBlocks);
 
   // Handle the case where only a single block was inlined.
-  if (singleBlockFastPath && std::next(newBlocks.begin()) == newBlocks.end()) {
+  if (singleBlockFastPath && llvm::hasSingleElement(newBlocks)) {
     // Run the result attribute handler on the terminator operands.
     Operation *firstBlockTerminator = firstNewBlock->getTerminator();
     builder.setInsertionPoint(firstBlockTerminator);
@@ -348,9 +348,10 @@ 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, const InlinerConfig &config,
+                 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())
@@ -372,51 +373,52 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   }
 
   // Call into the main region inliner function.
-  return inlineRegionImpl(interface, src, inlineBlock, inlinePoint, mapper,
-                          resultsToReplace, resultsToReplace.getTypes(),
+  return inlineRegionImpl(interface, config, src, inlineBlock, inlinePoint,
+                          mapper, resultsToReplace, resultsToReplace.getTypes(),
                           inlineLoc, shouldCloneInlinedRegion, call);
 }
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
+LogicalResult mlir::inlineRegion(InlinerInterface &interface,
+                                 const InlinerConfig &config, Region *src,
                                  Operation *inlinePoint, IRMapping &mapper,
                                  ValueRange resultsToReplace,
                                  TypeRange regionResultTypes,
                                  std::optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, src, inlinePoint->getBlock(),
+  return inlineRegion(interface, config, 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, const InlinerConfig &config, Region *src,
+    Block *inlineBlock, Block::iterator inlinePoint, IRMapping &mapper,
+    ValueRange resultsToReplace, TypeRange regionResultTypes,
+    std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(interface, config, src, inlineBlock, inlinePoint,
+                          mapper, resultsToReplace, regionResultTypes,
+                          inlineLoc, shouldCloneInlinedRegion);
 }
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
+LogicalResult mlir::inlineRegion(InlinerInterface &interface,
+                                 const InlinerConfig &config, Region *src,
                                  Operation *inlinePoint,
                                  ValueRange inlinedOperands,
                                  ValueRange resultsToReplace,
                                  std::optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, src, inlinePoint->getBlock(),
+  return inlineRegion(interface, config, 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,
+
+LogicalResult
+mlir::inlineRegion(InlinerInterface &interface, const InlinerConfig &config,
+                   Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+                   ValueRange inlinedOperands, ValueRange resultsToReplace,
+                   std::optional<Location> inlineLoc,
+                   bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(interface, config, src, inlineBlock, inlinePoint,
                           inlinedOperands, resultsToReplace, inlineLoc,
                           shouldCloneInlinedRegion);
 }
@@ -450,6 +452,7 @@ static Value materializeConversion(const DialectInlinerInterface *interface,
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
 LogicalResult mlir::inlineCall(InlinerInterface &interface,
+                               const InlinerConfig &config,
                                CallOpInterface call,
                                CallableOpInterface callable, Region *src,
                                bool shouldCloneInlinedRegion) {
@@ -526,7 +529,7 @@ LogicalResult mlir::inlineCall(InlinerInterface &interface,
     return cleanupState();
 
   // Attempt to inline the call.
-  if (failed(inlineRegionImpl(interface, src, call->getBlock(),
+  if (failed(inlineRegionImpl(interface, config, src, call->getBlock(),
                               ++call->getIterator(), mapper, callResults,
                               callableResultTypes, call.getLoc(),
                               shouldCloneInlinedRegion, call)))
diff --git a/mlir/test/lib/Transforms/TestInlining.cpp b/mlir/test/lib/Transforms/TestInlining.cpp
index 223cc78dd1e21..6c9ae845a16c2 100644
--- a/mlir/test/lib/Transforms/TestInlining.cpp
+++ b/mlir/test/lib/Transforms/TestInlining.cpp
@@ -18,6 +18,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Inliner.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSet.h"
 
@@ -25,8 +26,9 @@ using namespace mlir;
 using namespace test;
 
 namespace {
-struct Inliner : public PassWrapper<Inliner, OperationPass<func::FuncOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Inliner)
+struct InlinerTest
+    : public PassWrapper<InlinerTest, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InlinerTest)
 
   StringRef getArgument() const final { return "test-inline"; }
   StringRef getDescription() const final {
@@ -34,6 +36,9 @@ struct Inliner : public PassWrapper<Inliner, OperationPass<func::FuncOp>> {
   }
 
   void runOnOperation() override {
+    InlinerConfig config;
+    config.setCanHandleMultipleBlocks(false);
+
     auto function = getOperation();
 
     // Collect each of the direct function calls within the module.
@@ -54,8 +59,8 @@ struct Inliner : public PassWrapper<Inliner, OperationPass<func::FuncOp>> {
       // Inline the functional region operation, but only clone the internal
       // region if there is more than one use.
       if (failed(inlineRegion(
-              interface, &callee.getBody(), caller, caller.getArgOperands(),
-              caller.getResults(), caller.getLoc(),
+              interface, config, &callee.getBody(), caller,
+              caller.getArgOperands(), caller.getResults(), caller.getLoc(),
               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
         continue;
 
@@ -71,6 +76,6 @@ struct Inliner : public PassWrapper<Inliner, OperationPass<func::FuncOp>> {
 
 namespace mlir {
 namespace test {
-void registerInliner() { PassRegistration<Inliner>(); }
+void registerInliner() { PassRegistration<InlinerTest>(); }
 } // namespace test
 } // namespace mlir
diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
index 738e861ccdad1..9e06d4dba9180 100644
--- a/mlir/test/lib/Transforms/TestInliningCallback.cpp
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -95,23 +95,22 @@ struct InlinerCallback
       return true;
     };
 
-    // This customized inliner can turn multiple blocks into a single block.
-    auto canHandleMultipleBlocksCb = [&]() { return true; };
-
-    // Get an instance of the inliner.
-    Inliner inliner(function, cg, *this, getAnalysisManager(),
-                    runPipelineHelper, config, profitabilityCb,
-                    canHandleMultipleBlocksCb);
-
-    // Customize the implementation of Inliner::doClone
-    inliner.setCloneCallback([](OpBuilder &builder, Region *src,
-                                Block *inlineBlock, Block *postInsertBlock,
-                                IRMapping &mapper,
-                                bool shouldCloneInlinedRegion) {
+    // Set the clone callback in the config
+    config.setCloneCallback([](OpBuilder &builder, Region *src,
+                               Block *inlineBlock, Block *postInsertBlock,
+                               IRMapping &mapper,
+                               bool shouldCloneInlinedRegion) {
       return testDoClone(builder, src, inlineBlock, postInsertBlock, mapper,
                          shouldCloneInlinedRegion);
     });
 
+    // Set canHandleMultipleBlocks to true in the config
+    config.setCanHandleMultipleBlocks(true);
+
+    // Get an instance of the inliner.
+    Inliner inliner(function, cg, *this, getAnalysisManager(),
+                    runPipelineHelper, config, profitabilityCb);
+
     // Collect each of the direct function calls within the module.
     SmallVector<func::CallIndirectOp, 16> callers;
     function.walk(
@@ -130,8 +129,8 @@ struct InlinerCallback
       // Inline the functional region operation, but only clone the internal
       // region if there is more than one use.
       if (failed(inlineRegion(
-              interface, &callee.getBody(), caller, caller.getArgOperands(),
-              caller.getResults(), caller.getLoc(),
+              interface, config, &callee.getBody(), caller,
+              caller.getArgOperands(), caller.getResults(), caller.getLoc(),
               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
         continue;
 

From e8a70cc1949fecc6ed66909da6cee633f032cae5 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Wed, 2 Apr 2025 18:05:16 -0700
Subject: [PATCH 05/15] fix format.

---
 mlir/lib/Transforms/Utils/InliningUtils.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 1e5d62cec0b3d..73d8c25540bea 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -32,11 +32,11 @@ static void
 remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
                       Location callerLoc) {
   DenseMap<Location, LocationAttr> mappedLocations;
-    LocationAttr newLoc = stackLocations(loc, callerLoc);
+  auto remapLoc = [&](Location loc) {
     auto [it, inserted] = mappedLocations.try_emplace(loc);
     // Only query the attribute uniquer once per callsite attribute.
     if (inserted) {
-      LocationAttr newLoc = stackLocations(loc, callerLoc);
+      auto newLoc = CallSiteLoc::get(loc, callerLoc);
       it->getSecond() = newLoc;
     }
     return it->second;

From 4b34cfdfe3d059ac8f5a2be5e2b906d1cb1fc16c Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Wed, 2 Apr 2025 19:43:18 -0700
Subject: [PATCH 06/15] Remove unused codes

---
 mlir/include/mlir/Transforms/Inliner.h    | 27 +++++------------------
 mlir/test/lib/Transforms/TestInlining.cpp |  1 -
 2 files changed, 5 insertions(+), 23 deletions(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 84c96230f98f7..f90a71ecd0089 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -89,7 +89,11 @@ class InlinerConfig {
                                        src->getBlocks(), src->begin(),
                                        src->end());
   };
-  /// Determining if multiple blocks can be handled
+  /// Determining 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 shuold perform the extra
+  /// transformation. see the example in
+  /// mlir/test/lib/Transforms/TestInliningCallback.cpp
   bool canHandleMultipleBlocks{false};
 };
 
@@ -121,18 +125,6 @@ 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)
@@ -145,14 +137,6 @@ class Inliner {
   /// 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;
@@ -174,7 +158,6 @@ class Inliner {
   /// Forward declaration of the class providing the actual implementation.
   class Impl;
 };
-
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_INLINER_H
diff --git a/mlir/test/lib/Transforms/TestInlining.cpp b/mlir/test/lib/Transforms/TestInlining.cpp
index 6c9ae845a16c2..8e6f41b8d87b3 100644
--- a/mlir/test/lib/Transforms/TestInlining.cpp
+++ b/mlir/test/lib/Transforms/TestInlining.cpp
@@ -37,7 +37,6 @@ struct InlinerTest
 
   void runOnOperation() override {
     InlinerConfig config;
-    config.setCanHandleMultipleBlocks(false);
 
     auto function = getOperation();
 

From d5849f9df86712a4c2eeccad6d06c98692a4f6c3 Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:38:35 -0700
Subject: [PATCH 07/15] Update mlir/include/mlir/Transforms/Inliner.h

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/include/mlir/Transforms/Inliner.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index f90a71ecd0089..5db542f98ba67 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -89,7 +89,7 @@ class InlinerConfig {
                                        src->getBlocks(), src->begin(),
                                        src->end());
   };
-  /// Determining if the inliner can inline a function containing multiple
+  /// 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 shuold perform the extra
   /// transformation. see the example in

From 0e6fd3fea88e582dc489b5e55aa622a320d1180b Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:39:03 -0700
Subject: [PATCH 08/15] Update mlir/include/mlir/Transforms/Inliner.h

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/include/mlir/Transforms/Inliner.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 5db542f98ba67..058bc4f7cf470 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -91,7 +91,7 @@ class InlinerConfig {
   };
   /// 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 shuold perform the extra
+  /// 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};

From 94235daa2c6c100116002923a4681f22171cf272 Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:39:45 -0700
Subject: [PATCH 09/15] Update mlir/include/mlir/Transforms/Inliner.h

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/include/mlir/Transforms/Inliner.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 058bc4f7cf470..0a8c276076d0b 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -56,7 +56,7 @@ class InlinerConfig {
   void setCloneCallback(CloneCallbackTy callback) {
     cloneCallback = std::move(callback);
   }
-  void setCanHandleMultipleBlocks(bool value) {
+  void setCanHandleMultipleBlocks(bool value = true) {
     canHandleMultipleBlocks = value;
   }
 

From 6e7e6e180c3f349e3be8c4c4741a5f6e012a4e86 Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:40:30 -0700
Subject: [PATCH 10/15] Update
 mlir/test/lib/Transforms/TestInliningCallback.cpp

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/test/lib/Transforms/TestInliningCallback.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
index 9e06d4dba9180..fb66d56b0021f 100644
--- a/mlir/test/lib/Transforms/TestInliningCallback.cpp
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -105,7 +105,7 @@ struct InlinerCallback
     });
 
     // Set canHandleMultipleBlocks to true in the config
-    config.setCanHandleMultipleBlocks(true);
+    config.setCanHandleMultipleBlocks();
 
     // Get an instance of the inliner.
     Inliner inliner(function, cg, *this, getAnalysisManager(),

From e9b7963ee2e7502d23436b281810a75eb1eef985 Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:41:21 -0700
Subject: [PATCH 11/15] Update
 mlir/test/lib/Transforms/TestInliningCallback.cpp

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/test/lib/Transforms/TestInliningCallback.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
index fb66d56b0021f..68f4b1476c77b 100644
--- a/mlir/test/lib/Transforms/TestInliningCallback.cpp
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -88,7 +88,7 @@ struct InlinerCallback
     InlinerConfig config;
     CallGraph &cg = getAnalysis<CallGraph>();
 
-    auto function = getOperation();
+    func::FuncOp function = getOperation();
 
     // By default, assume that any inlining is profitable.
     auto profitabilityCb = [&](const mlir::Inliner::ResolvedCall &call) {

From 59b08390aeca577b078d1dbc728fb1530992a400 Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 13:41:54 -0700
Subject: [PATCH 12/15] Update
 mlir/test/lib/Transforms/TestInliningCallback.cpp

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/test/lib/Transforms/TestInliningCallback.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
index 68f4b1476c77b..486cfeb62999e 100644
--- a/mlir/test/lib/Transforms/TestInliningCallback.cpp
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -112,7 +112,7 @@ struct InlinerCallback
                     runPipelineHelper, config, profitabilityCb);
 
     // Collect each of the direct function calls within the module.
-    SmallVector<func::CallIndirectOp, 16> callers;
+    SmallVector<func::CallIndirectOp> callers;
     function.walk(
         [&](func::CallIndirectOp caller) { callers.push_back(caller); });
 

From de67999e9f6c6c74ed7b643a840d3c4dfb77e1f9 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Thu, 3 Apr 2025 18:52:39 -0700
Subject: [PATCH 13/15] Pass cloneCallback function referece instead of whole
 InlinerConfig class

---
 mlir/include/mlir/Transforms/InliningUtils.h  | 29 ++++----
 mlir/lib/Transforms/Utils/Inliner.cpp         |  2 +-
 mlir/lib/Transforms/Utils/InliningUtils.cpp   | 68 ++++++++++---------
 mlir/test/lib/Transforms/TestInlining.cpp     |  2 +-
 .../lib/Transforms/TestInliningCallback.cpp   |  2 +-
 5 files changed, 54 insertions(+), 49 deletions(-)

diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index e14b2412f7853..6d8aecb1ff64e 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -255,16 +255,17 @@ class InlinerInterface
 /// information. 'shouldCloneInlinedRegion' corresponds to whether the source
 /// region should be cloned into the 'inlinePoint' or spliced directly.
 LogicalResult inlineRegion(InlinerInterface &interface,
-                           const InlinerConfig &config, Region *src,
-                           Operation *inlinePoint, IRMapping &mapper,
-                           ValueRange resultsToReplace,
+                           InlinerConfig::CloneCallbackTy cloneCallback,
+                           Region *src, Operation *inlinePoint,
+                           IRMapping &mapper, ValueRange resultsToReplace,
                            TypeRange regionResultTypes,
                            std::optional<Location> inlineLoc = std::nullopt,
                            bool shouldCloneInlinedRegion = true);
 LogicalResult inlineRegion(InlinerInterface &interface,
-                           const InlinerConfig &config, Region *src,
-                           Block *inlineBlock, Block::iterator inlinePoint,
-                           IRMapping &mapper, ValueRange resultsToReplace,
+                           InlinerConfig::CloneCallbackTy cloneCallback,
+                           Region *src, Block *inlineBlock,
+                           Block::iterator inlinePoint, IRMapping &mapper,
+                           ValueRange resultsToReplace,
                            TypeRange regionResultTypes,
                            std::optional<Location> inlineLoc = std::nullopt,
                            bool shouldCloneInlinedRegion = true);
@@ -273,14 +274,16 @@ LogicalResult inlineRegion(InlinerInterface &interface,
 /// providing the set of operands ('inlinedOperands') that should be used
 /// in-favor of the region arguments when inlining.
 LogicalResult inlineRegion(InlinerInterface &interface,
-                           const InlinerConfig &config, Region *src,
-                           Operation *inlinePoint, ValueRange inlinedOperands,
+                           InlinerConfig::CloneCallbackTy cloneCallback,
+                           Region *src, Operation *inlinePoint,
+                           ValueRange inlinedOperands,
                            ValueRange resultsToReplace,
                            std::optional<Location> inlineLoc = std::nullopt,
                            bool shouldCloneInlinedRegion = true);
 LogicalResult inlineRegion(InlinerInterface &interface,
-                           const InlinerConfig &config, Region *src,
-                           Block *inlineBlock, Block::iterator inlinePoint,
+                           InlinerConfig::CloneCallbackTy cloneCallback,
+                           Region *src, Block *inlineBlock,
+                           Block::iterator inlinePoint,
                            ValueRange inlinedOperands,
                            ValueRange resultsToReplace,
                            std::optional<Location> inlineLoc = std::nullopt,
@@ -293,9 +296,9 @@ LogicalResult inlineRegion(InlinerInterface &interface,
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
 LogicalResult inlineCall(InlinerInterface &interface,
-                         const InlinerConfig &config, CallOpInterface call,
-                         CallableOpInterface callable, Region *src,
-                         bool shouldCloneInlinedRegion = true);
+                         InlinerConfig::CloneCallbackTy cloneCallback,
+                         CallOpInterface call, CallableOpInterface callable,
+                         Region *src, bool shouldCloneInlinedRegion = true);
 
 } // namespace mlir
 
diff --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 445d4bd5efc75..2c1daab67a0d3 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -651,7 +651,7 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
     bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
 
     LogicalResult inlineResult =
-        inlineCall(inlinerIface, inliner.config, call,
+        inlineCall(inlinerIface, inliner.config.getCloneCallback(), call,
                    cast<CallableOpInterface>(targetRegion->getParentOp()),
                    targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
     if (failed(inlineResult)) {
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 73d8c25540bea..550ec6b6215eb 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -246,8 +246,9 @@ static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
 }
 
 static LogicalResult
-inlineRegionImpl(InlinerInterface &interface, const InlinerConfig &config,
-                 Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+inlineRegionImpl(InlinerInterface &interface,
+                 InlinerConfig::CloneCallbackTy cloneCallback, Region *src,
+                 Block *inlineBlock, Block::iterator inlinePoint,
                  IRMapping &mapper, ValueRange resultsToReplace,
                  TypeRange regionResultTypes, std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, CallOpInterface call = {}) {
@@ -278,8 +279,8 @@ inlineRegionImpl(InlinerInterface &interface, const InlinerConfig &config,
 
   // Clone the callee's source into the caller.
   Block *postInsertBlock = inlineBlock->splitBlock(inlinePoint);
-  config.getCloneCallback()(builder, src, inlineBlock, postInsertBlock, mapper,
-                            shouldCloneInlinedRegion);
+  cloneCallback(builder, src, inlineBlock, postInsertBlock, mapper,
+                shouldCloneInlinedRegion);
 
   // Get the range of newly inserted blocks.
   auto newBlocks = llvm::make_range(std::next(inlineBlock->getIterator()),
@@ -348,8 +349,9 @@ inlineRegionImpl(InlinerInterface &interface, const InlinerConfig &config,
 }
 
 static LogicalResult
-inlineRegionImpl(InlinerInterface &interface, const InlinerConfig &config,
-                 Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+inlineRegionImpl(InlinerInterface &interface,
+                 InlinerConfig::CloneCallbackTy cloneCallback, Region *src,
+                 Block *inlineBlock, Block::iterator inlinePoint,
                  ValueRange inlinedOperands, ValueRange resultsToReplace,
                  std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, CallOpInterface call = {}) {
@@ -373,54 +375,54 @@ inlineRegionImpl(InlinerInterface &interface, const InlinerConfig &config,
   }
 
   // Call into the main region inliner function.
-  return inlineRegionImpl(interface, config, 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,
-                                 const InlinerConfig &config, Region *src,
-                                 Operation *inlinePoint, IRMapping &mapper,
-                                 ValueRange resultsToReplace,
+                                 InlinerConfig::CloneCallbackTy cloneCallback,
+                                 Region *src, Operation *inlinePoint,
+                                 IRMapping &mapper, ValueRange resultsToReplace,
                                  TypeRange regionResultTypes,
                                  std::optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, config, src, inlinePoint->getBlock(),
+  return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),
                       ++inlinePoint->getIterator(), mapper, resultsToReplace,
                       regionResultTypes, inlineLoc, shouldCloneInlinedRegion);
 }
 
 LogicalResult mlir::inlineRegion(
-    InlinerInterface &interface, const InlinerConfig &config, Region *src,
-    Block *inlineBlock, Block::iterator inlinePoint, IRMapping &mapper,
-    ValueRange resultsToReplace, TypeRange regionResultTypes,
+    InlinerInterface &interface, InlinerConfig::CloneCallbackTy cloneCallback,
+    Region *src, Block *inlineBlock, Block::iterator inlinePoint,
+    IRMapping &mapper, ValueRange resultsToReplace, TypeRange regionResultTypes,
     std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {
-  return inlineRegionImpl(interface, config, src, inlineBlock, inlinePoint,
-                          mapper, resultsToReplace, regionResultTypes,
-                          inlineLoc, shouldCloneInlinedRegion);
+  return inlineRegionImpl(
+      interface, cloneCallback, src, inlineBlock, inlinePoint, mapper,
+      resultsToReplace, regionResultTypes, inlineLoc, shouldCloneInlinedRegion);
 }
 
 LogicalResult mlir::inlineRegion(InlinerInterface &interface,
-                                 const InlinerConfig &config, Region *src,
-                                 Operation *inlinePoint,
+                                 InlinerConfig::CloneCallbackTy cloneCallback,
+                                 Region *src, Operation *inlinePoint,
                                  ValueRange inlinedOperands,
                                  ValueRange resultsToReplace,
                                  std::optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
-  return inlineRegion(interface, config, src, inlinePoint->getBlock(),
+  return inlineRegion(interface, cloneCallback, src, inlinePoint->getBlock(),
                       ++inlinePoint->getIterator(), inlinedOperands,
                       resultsToReplace, inlineLoc, shouldCloneInlinedRegion);
 }
 
-LogicalResult
-mlir::inlineRegion(InlinerInterface &interface, const InlinerConfig &config,
-                   Region *src, Block *inlineBlock, Block::iterator inlinePoint,
-                   ValueRange inlinedOperands, ValueRange resultsToReplace,
-                   std::optional<Location> inlineLoc,
-                   bool shouldCloneInlinedRegion) {
-  return inlineRegionImpl(interface, config, src, inlineBlock, inlinePoint,
-                          inlinedOperands, resultsToReplace, inlineLoc,
-                          shouldCloneInlinedRegion);
+LogicalResult mlir::inlineRegion(
+    InlinerInterface &interface, InlinerConfig::CloneCallbackTy 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,
@@ -452,7 +454,7 @@ static Value materializeConversion(const DialectInlinerInterface *interface,
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
 LogicalResult mlir::inlineCall(InlinerInterface &interface,
-                               const InlinerConfig &config,
+                               InlinerConfig::CloneCallbackTy cloneCallback,
                                CallOpInterface call,
                                CallableOpInterface callable, Region *src,
                                bool shouldCloneInlinedRegion) {
@@ -529,7 +531,7 @@ LogicalResult mlir::inlineCall(InlinerInterface &interface,
     return cleanupState();
 
   // Attempt to inline the call.
-  if (failed(inlineRegionImpl(interface, config, src, call->getBlock(),
+  if (failed(inlineRegionImpl(interface, cloneCallback, src, call->getBlock(),
                               ++call->getIterator(), mapper, callResults,
                               callableResultTypes, call.getLoc(),
                               shouldCloneInlinedRegion, call)))
diff --git a/mlir/test/lib/Transforms/TestInlining.cpp b/mlir/test/lib/Transforms/TestInlining.cpp
index 8e6f41b8d87b3..ae904a92a5d68 100644
--- a/mlir/test/lib/Transforms/TestInlining.cpp
+++ b/mlir/test/lib/Transforms/TestInlining.cpp
@@ -58,7 +58,7 @@ struct InlinerTest
       // Inline the functional region operation, but only clone the internal
       // region if there is more than one use.
       if (failed(inlineRegion(
-              interface, config, &callee.getBody(), caller,
+              interface, config.getCloneCallback(), &callee.getBody(), caller,
               caller.getArgOperands(), caller.getResults(), caller.getLoc(),
               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
         continue;
diff --git a/mlir/test/lib/Transforms/TestInliningCallback.cpp b/mlir/test/lib/Transforms/TestInliningCallback.cpp
index 486cfeb62999e..012d62b7b1b42 100644
--- a/mlir/test/lib/Transforms/TestInliningCallback.cpp
+++ b/mlir/test/lib/Transforms/TestInliningCallback.cpp
@@ -129,7 +129,7 @@ struct InlinerCallback
       // Inline the functional region operation, but only clone the internal
       // region if there is more than one use.
       if (failed(inlineRegion(
-              interface, config, &callee.getBody(), caller,
+              interface, config.getCloneCallback(), &callee.getBody(), caller,
               caller.getArgOperands(), caller.getResults(), caller.getLoc(),
               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
         continue;

From 395196e8e0a21205cf2ad81ba1be668998900f0d Mon Sep 17 00:00:00 2001
From: junfengd-nv <junfengd@nvidia.com>
Date: Fri, 4 Apr 2025 08:45:11 -0700
Subject: [PATCH 14/15] Update mlir/include/mlir/Transforms/Inliner.h

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
---
 mlir/include/mlir/Transforms/Inliner.h | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 0a8c276076d0b..6bfc5706fecb8 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -27,10 +27,12 @@ class InlinerConfig {
 public:
   using DefaultPipelineTy = std::function<void(OpPassManager &)>;
   using OpPipelinesTy = llvm::StringMap<OpPassManager>;
-  using CloneCallbackTy =
-      std::function<void(OpBuilder &builder, Region *src, Block *inlineBlock,
+  using CloneCallbackSigTy =
+      void(*)(OpBuilder &builder, Region *src, Block *inlineBlock,
                          Block *postInsertBlock, IRMapping &mapper,
-                         bool shouldCloneInlinedRegion)>;
+                         bool shouldCloneInlinedRegion);
+  using CloneCallbackTy =
+      std::function<CloneCallbackSigTy>;
 
   InlinerConfig() = default;
   InlinerConfig(DefaultPipelineTy defaultPipeline,

From 8975cee359de7b0d5587660283c5145a68156676 Mon Sep 17 00:00:00 2001
From: Junfeng Dong <junfengd@nvidia.com>
Date: Fri, 4 Apr 2025 14:26:24 -0700
Subject: [PATCH 15/15] Correct the parameter to use function_ref.

---
 mlir/include/mlir/Transforms/Inliner.h       | 11 ++--
 mlir/include/mlir/Transforms/InliningUtils.h | 68 ++++++++++----------
 mlir/lib/Transforms/Utils/InliningUtils.cpp  | 60 ++++++++---------
 3 files changed, 69 insertions(+), 70 deletions(-)

diff --git a/mlir/include/mlir/Transforms/Inliner.h b/mlir/include/mlir/Transforms/Inliner.h
index 6bfc5706fecb8..506b4455af646 100644
--- a/mlir/include/mlir/Transforms/Inliner.h
+++ b/mlir/include/mlir/Transforms/Inliner.h
@@ -27,12 +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>;
+  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,
diff --git a/mlir/include/mlir/Transforms/InliningUtils.h b/mlir/include/mlir/Transforms/InliningUtils.h
index 6d8aecb1ff64e..552030983d724 100644
--- a/mlir/include/mlir/Transforms/InliningUtils.h
+++ b/mlir/include/mlir/Transforms/InliningUtils.h
@@ -254,40 +254,39 @@ 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,
-                           InlinerConfig::CloneCallbackTy cloneCallback,
-                           Region *src, Operation *inlinePoint,
-                           IRMapping &mapper, ValueRange resultsToReplace,
-                           TypeRange regionResultTypes,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
-LogicalResult inlineRegion(InlinerInterface &interface,
-                           InlinerConfig::CloneCallbackTy cloneCallback,
-                           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,
-                           InlinerConfig::CloneCallbackTy cloneCallback,
-                           Region *src, Operation *inlinePoint,
-                           ValueRange inlinedOperands,
-                           ValueRange resultsToReplace,
-                           std::optional<Location> inlineLoc = std::nullopt,
-                           bool shouldCloneInlinedRegion = true);
-LogicalResult inlineRegion(InlinerInterface &interface,
-                           InlinerConfig::CloneCallbackTy cloneCallback,
-                           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
@@ -295,10 +294,11 @@ LogicalResult inlineRegion(InlinerInterface &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 inlineCall(InlinerInterface &interface,
-                         InlinerConfig::CloneCallbackTy cloneCallback,
-                         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
 
diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 550ec6b6215eb..bb10f3fcb0184 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -247,8 +247,8 @@ static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
 
 static LogicalResult
 inlineRegionImpl(InlinerInterface &interface,
-                 InlinerConfig::CloneCallbackTy cloneCallback, Region *src,
-                 Block *inlineBlock, Block::iterator inlinePoint,
+                 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 = {}) {
@@ -350,8 +350,8 @@ inlineRegionImpl(InlinerInterface &interface,
 
 static LogicalResult
 inlineRegionImpl(InlinerInterface &interface,
-                 InlinerConfig::CloneCallbackTy cloneCallback, Region *src,
-                 Block *inlineBlock, Block::iterator inlinePoint,
+                 function_ref<InlinerConfig::CloneCallbackSigTy> cloneCallback,
+                 Region *src, Block *inlineBlock, Block::iterator inlinePoint,
                  ValueRange inlinedOperands, ValueRange resultsToReplace,
                  std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, CallOpInterface call = {}) {
@@ -381,45 +381,45 @@ inlineRegionImpl(InlinerInterface &interface,
                           shouldCloneInlinedRegion, call);
 }
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface,
-                                 InlinerConfig::CloneCallbackTy cloneCallback,
-                                 Region *src, Operation *inlinePoint,
-                                 IRMapping &mapper, ValueRange resultsToReplace,
-                                 TypeRange regionResultTypes,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
+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, InlinerConfig::CloneCallbackTy cloneCallback,
-    Region *src, Block *inlineBlock, Block::iterator inlinePoint,
-    IRMapping &mapper, ValueRange resultsToReplace, TypeRange regionResultTypes,
+    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,
-                                 InlinerConfig::CloneCallbackTy cloneCallback,
-                                 Region *src, Operation *inlinePoint,
-                                 ValueRange inlinedOperands,
-                                 ValueRange resultsToReplace,
-                                 std::optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
+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, InlinerConfig::CloneCallbackTy cloneCallback,
-    Region *src, Block *inlineBlock, Block::iterator inlinePoint,
-    ValueRange inlinedOperands, ValueRange resultsToReplace,
-    std::optional<Location> inlineLoc, bool shouldCloneInlinedRegion) {
+    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);
@@ -453,11 +453,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,
-                               InlinerConfig::CloneCallbackTy cloneCallback,
-                               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();