[Mlir-commits] [mlir] [mlir][inliner] Bound recursive expansion across SCC iterations (PR #211377)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 22 13:57:34 PDT 2026
https://github.com/khaki3 created https://github.com/llvm/llvm-project/pull/211377
The MLIR inliner tracks inline history only within one invocation of `inlineCallsInSCC`. This correctly detects recursion while processing the current call worklist.
However, `inlineSCC` repeatedly runs optimization and inlining. On every iteration, the inline history was recreated. A recursive call left from the previous iteration was therefore treated as a new root call and expanded again.
For example:
```mlir
func.func @caller(%arg: i32) -> i32 {
%0 = call @a(%arg) : (i32) -> i32
return %0 : i32
}
func.func @a(%arg: i32) -> i32 {
%0 = arith.addi %arg, %c1 : i32
%1 = call @b(%0) : (i32) -> i32
return %1 : i32
}
func.func @b(%arg: i32) -> i32 {
%0 = arith.addi %arg, %c2 : i32
%1 = call @a(%0) : (i32) -> i32
return %1 : i32
}
```
### Before
The first inlining iteration produced one bounded expansion:
```text
caller -> a -> b -> a
```
The local history stopped the final `a`, but the next SCC iteration forgot that history and expanded it again:
```text
caller -> a -> b -> a -> b -> a -> b -> a ...
```
With four inliner iterations, the caller contained eight cloned additions before the remaining recursive call.
Conceptually:
```mlir
%0 = arith.addi %arg, %c1
%1 = arith.addi %0, %c2
%2 = arith.addi %1, %c1
%3 = arith.addi %2, %c2
%4 = arith.addi %3, %c1
%5 = arith.addi %4, %c2
%6 = arith.addi %5, %c1
%7 = arith.addi %6, %c2
%8 = call @a(%7)
```
This causes recursive IR growth on every SCC fixed-point iteration.
## Algorithm change
The existing per-worklist inline history remains unchanged. It still detects whether a target already appears in the current expansion lineage.
When that history detects a recursive expansion, the inliner now records the corresponding call-graph edge:
```text
(sourceNode, targetNode)
```
Two edge sets are used:
```cpp
BlockedEdges blockedEdges;
BlockedEdges newlyBlockedEdges;
```
The algorithm is:
```text
for each SCC iteration:
optimize the SCC
for each call in the current worklist:
if target occurs in the call's local inline history:
add (source, target) to newlyBlockedEdges
do not inline this recursive call
else if (source, target) is in blockedEdges:
do not inline it
else:
apply the normal legality and profitability checks
inline when permitted
after the complete worklist is processed:
merge newlyBlockedEdges into blockedEdges
```
Newly detected edges are not blocked immediately. Activation is delayed until the next SCC iteration so that independent calls on the same edge already present in the current worklist retain their normal inlining opportunities.
The edge state is scoped to `inlineSCC`. Other call-graph edges continue through the normal optimization and inlining fixed-point process, including opportunities exposed by canonicalization.
Call-graph edges are used instead of call-operation pointers because the configurable optimization pipeline may replace call operations between iterations. Exact call-site provenance cannot safely be retained across arbitrary transformations, so a proven recursive graph edge is conservatively blocked for the remaining iterations of that SCC.
### After
The first expansion remains unchanged, but the recursive edge is blocked in later SCC iterations:
```mlir
%0 = arith.addi %arg, %c1
%1 = arith.addi %0, %c2
%2 = call @a(%1)
```
With canonicalization, the additions may be folded:
```mlir
%c3 = arith.constant 3 : i32
%0 = arith.addi %arg, %c3
%1 = call @a(%0)
```
The recursive expansion is bounded while unrelated inlining and optimization continue.
>From 33c8581519e2db6e148bff87e19e66c6bbfb3747 Mon Sep 17 00:00:00 2001
From: Kazuaki Matsumura <kmatsumura at nvidia.com>
Date: Wed, 22 Jul 2026 13:50:09 -0700
Subject: [PATCH] [MLIR][Inliner] Bound recursive expansion across iterations
---
mlir/lib/Transforms/Utils/Inliner.cpp | 36 +++++---
mlir/test/Transforms/inlining-recursive.mlir | 94 ++++++++++++++++++++
2 files changed, 120 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/Inliner.cpp b/mlir/lib/Transforms/Utils/Inliner.cpp
index 40950312d566f..293789cf5873a 100644
--- a/mlir/lib/Transforms/Utils/Inliner.cpp
+++ b/mlir/lib/Transforms/Utils/Inliner.cpp
@@ -19,6 +19,7 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/DebugStringHelper.h"
#include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SCCIterator.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/DebugLog.h"
@@ -354,7 +355,7 @@ static std::string getNodeName(CallOpInterface op) {
return "_unnamed_callee_";
}
-/// Return true if the specified `inlineHistoryID` indicates an inline history
+/// Return true if the specified `inlineHistoryID` indicates an inline history
/// that already includes `node`.
static bool inlineHistoryIncludes(
CallGraphNode *node, std::optional<size_t> inlineHistoryID,
@@ -419,6 +420,9 @@ struct InlinerInterfaceImpl : public InlinerInterface {
namespace mlir {
+using CallGraphEdge = std::pair<CallGraphNode *, CallGraphNode *>;
+using BlockedEdges = DenseSet<CallGraphEdge>;
+
class Inliner::Impl {
public:
Impl(Inliner &inliner) : inliner(inliner) {}
@@ -453,7 +457,8 @@ class Inliner::Impl {
/// Attempt to inline calls within the given scc. This function returns
/// success if any calls were inlined, failure otherwise.
LogicalResult inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
- CGUseList &useList, CallGraphSCC ¤tSCC);
+ CGUseList &useList, CallGraphSCC ¤tSCC,
+ BlockedEdges &blockedEdges);
/// Returns true if the given call should be inlined.
bool shouldInline(ResolvedCall &resolvedCall);
@@ -471,10 +476,14 @@ LogicalResult Inliner::Impl::inlineSCC(InlinerInterfaceImpl &inlinerIface,
// hit the maximum iteration count. Simplifying early helps to refine the cost
// model, and in future iterations may devirtualize new calls.
unsigned iterationCount = 0;
+ // Optimization may replace calls, preventing exact provenance tracking.
+ // Conservatively retain proven recursive graph edges across iterations.
+ BlockedEdges blockedEdges;
do {
if (failed(optimizeSCC(inlinerIface.cg, useList, currentSCC, context)))
return failure();
- if (failed(inlineCallsInSCC(inlinerIface, useList, currentSCC)))
+ if (failed(
+ inlineCallsInSCC(inlinerIface, useList, currentSCC, blockedEdges)))
break;
} while (++iterationCount < inliner.config.getMaxInliningIterations());
return success();
@@ -580,7 +589,8 @@ Inliner::Impl::optimizeCallable(CallGraphNode *node,
/// success if any calls were inlined, failure otherwise.
LogicalResult
Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
- CGUseList &useList, CallGraphSCC ¤tSCC) {
+ CGUseList &useList, CallGraphSCC ¤tSCC,
+ BlockedEdges &blockedEdges) {
CallGraph &cg = inlinerIface.cg;
auto &calls = inlinerIface.calls;
@@ -604,12 +614,12 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
}
}
- // When inlining a callee produces new call sites, we want to keep track of
- // the fact that they were inlined from the callee. This allows us to avoid
- // infinite inlining.
+ // When inlining a callee produces new call sites, remember that they came
+ // from the callee to avoid recursively inlining through a cycle.
using InlineHistoryT = std::optional<size_t>;
SmallVector<std::pair<CallGraphNode *, InlineHistoryT>, 8> inlineHistory;
std::vector<InlineHistoryT> callHistory(calls.size(), InlineHistoryT{});
+ BlockedEdges newlyBlockedEdges;
LLVM_DEBUG({
LDBG() << "* Inliner: Initial calls in SCC are: {";
@@ -629,7 +639,11 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
InlineHistoryT inlineHistoryID = callHistory[i];
bool inHistory =
inlineHistoryIncludes(it.targetNode, inlineHistoryID, inlineHistory);
- bool doInline = !inHistory && shouldInline(it);
+ auto edge = std::make_pair(it.sourceNode, it.targetNode);
+ if (inHistory)
+ newlyBlockedEdges.insert(edge);
+ bool doInline =
+ !inHistory && !blockedEdges.contains(edge) && shouldInline(it);
CallOpInterface call = it.call;
LLVM_DEBUG({
if (doInline)
@@ -657,8 +671,7 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
}
inlinedAnyCalls = true;
- // Create a inline history entry for this inlined call, so that we remember
- // that new callsites came about due to inlining Callee.
+ // Record that the new callsites came from inlining the callee.
InlineHistoryT newInlineHistoryID{inlineHistory.size()};
inlineHistory.push_back(std::make_pair(it.targetNode, inlineHistoryID));
@@ -695,6 +708,9 @@ Inliner::Impl::inlineCallsInSCC(InlinerInterfaceImpl &inlinerIface,
currentSCC.remove(node);
inlinerIface.markForDeletion(node);
}
+ // Delay blocking until the next iteration so independent calls on the same
+ // edge in the current worklist retain their inlining opportunities.
+ blockedEdges.insert(newlyBlockedEdges.begin(), newlyBlockedEdges.end());
calls.clear();
return success(inlinedAnyCalls);
}
diff --git a/mlir/test/Transforms/inlining-recursive.mlir b/mlir/test/Transforms/inlining-recursive.mlir
index f953935475e1a..4ea72884a4d59 100644
--- a/mlir/test/Transforms/inlining-recursive.mlir
+++ b/mlir/test/Transforms/inlining-recursive.mlir
@@ -1,5 +1,6 @@
// RUN: mlir-opt %s -inline='default-pipeline=' | FileCheck %s
// RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=' | FileCheck %s
+// RUN: mlir-opt %s -inline | FileCheck %s --check-prefix=DEFAULT
// CHECK-LABEL: func.func @foo0
func.func @foo0(%arg0 : i32) -> i32 {
@@ -23,3 +24,96 @@ func.func @foo1(%arg0 : i32) -> i32 {
%2 = call @foo0(%1) : (i32) -> i32
return %2 : i32
}
+
+// Verify that recursive expansion is bounded across inliner iterations.
+// CHECK-LABEL: func.func @caller
+// CHECK-NEXT: %{{.*}} = arith.constant 1 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = arith.constant 2 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = call @cycle_a
+// CHECK-NEXT: %{{.*}} = constant @leaf
+// CHECK-NEXT: %{{.*}} = call_indirect
+// CHECK-NEXT: %{{.*}} = constant @cycle_a
+// CHECK-NEXT: %{{.*}} = call_indirect
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @caller(%arg0 : i32) -> i32 {
+ %0 = call @cycle_a(%arg0) : (i32) -> i32
+ %1 = call @wrapper(%0) : (i32) -> i32
+ %2 = call @same_edge_wrapper(%1) : (i32) -> i32
+ return %2 : i32
+}
+
+// DEFAULT-LABEL: func.func @caller
+// DEFAULT-NEXT: %{{.*}} = arith.constant 6 : i32
+// DEFAULT-NEXT: %{{.*}} = arith.constant 4 : i32
+// DEFAULT-NEXT: %{{.*}} = arith.constant 3 : i32
+// DEFAULT-NEXT: %{{.*}} = arith.addi
+// DEFAULT-NEXT: %{{.*}} = call @cycle_a
+// DEFAULT-NEXT: %{{.*}} = arith.muli
+// DEFAULT-NEXT: %{{.*}} = arith.addi
+// DEFAULT-NEXT: %{{.*}} = call @cycle_a
+// DEFAULT-NEXT: return
+// DEFAULT-NEXT: }
+
+// CHECK-LABEL: func.func @cycle_a
+// DEFAULT-LABEL: func.func @cycle_a
+func.func @cycle_a(%arg0 : i32) -> i32 {
+ %c1 = arith.constant 1 : i32
+ %0 = arith.addi %arg0, %c1 : i32
+ %1 = call @cycle_b(%0) : (i32) -> i32
+ return %1 : i32
+}
+
+func.func @cycle_b(%arg0 : i32) -> i32 {
+ %c2 = arith.constant 2 : i32
+ %0 = arith.addi %arg0, %c2 : i32
+ %1 = call @cycle_a(%0) : (i32) -> i32
+ return %1 : i32
+}
+
+func.func @wrapper(%arg0 : i32) -> i32 {
+ %fn = constant @leaf : (i32) -> i32
+ %0 = call_indirect %fn(%arg0) : (i32) -> i32
+ return %0 : i32
+}
+
+func.func @same_edge_wrapper(%arg0 : i32) -> i32 {
+ %fn = constant @cycle_a : (i32) -> i32
+ %0 = call_indirect %fn(%arg0) : (i32) -> i32
+ return %0 : i32
+}
+
+func.func @leaf(%arg0 : i32) -> i32 {
+ %c4 = arith.constant 4 : i32
+ %0 = arith.muli %arg0, %c4 : i32
+ return %0 : i32
+}
+
+// CHECK-LABEL: func.func @two_calls
+// CHECK-NEXT: %{{.*}} = arith.constant 1 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = arith.constant 2 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = call @cycle_a
+// CHECK-NEXT: %{{.*}} = arith.constant 1 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = arith.constant 2 : i32
+// CHECK-NEXT: %{{.*}} = arith.addi
+// CHECK-NEXT: %{{.*}} = call @cycle_a
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+// DEFAULT-LABEL: func.func @two_calls
+// DEFAULT-NEXT: %{{.*}} = arith.constant 3 : i32
+// DEFAULT-NEXT: %{{.*}} = arith.addi
+// DEFAULT-NEXT: %{{.*}} = call @cycle_a
+// DEFAULT-NEXT: %{{.*}} = arith.addi
+// DEFAULT-NEXT: %{{.*}} = call @cycle_a
+// DEFAULT-NEXT: return
+// DEFAULT-NEXT: }
+func.func @two_calls(%arg0 : i32) -> i32 {
+ %0 = call @cycle_a(%arg0) : (i32) -> i32
+ %1 = call @cycle_a(%0) : (i32) -> i32
+ return %1 : i32
+}
More information about the Mlir-commits
mailing list