Skip to content

Commit 6b59b33

Browse files
authored
[MLIR] Handle call site locations when inlining (llvm#132247)
When inlining a `callee` with a call site debug location, the inlining infrastructure was trivially combining the `callee` and the `caller` locations, forming a "tree" of call stacks. Because of this, the remarks were printing an incomplete inlining stack. This commit handles this case and appends the `caller` location at the end of the `callee`'s stack, extending the chain.
1 parent 42d49a1 commit 6b59b33

File tree

2 files changed

+61
-7
lines changed

2 files changed

+61
-7
lines changed

mlir/lib/Transforms/Utils/InliningUtils.cpp

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,29 @@
2525

2626
using namespace mlir;
2727

28+
/// Combine `callee` location with `caller` location to create a stack that
29+
/// represents the call chain.
30+
/// If `callee` location is a `CallSiteLoc`, indicating an existing stack of
31+
/// locations, the `caller` location is appended to the end of it, extending
32+
/// the chain.
33+
/// Otherwise, a single `CallSiteLoc` is created, representing a direct call
34+
/// from `caller` to `callee`.
35+
static LocationAttr stackLocations(Location callee, Location caller) {
36+
Location lastCallee = callee;
37+
SmallVector<CallSiteLoc> calleeInliningStack;
38+
while (auto nextCallSite = dyn_cast<CallSiteLoc>(lastCallee)) {
39+
calleeInliningStack.push_back(nextCallSite);
40+
lastCallee = nextCallSite.getCaller();
41+
}
42+
43+
CallSiteLoc firstCallSite = CallSiteLoc::get(lastCallee, caller);
44+
for (CallSiteLoc currentCallSite : reverse(calleeInliningStack))
45+
firstCallSite =
46+
CallSiteLoc::get(currentCallSite.getCallee(), firstCallSite);
47+
48+
return firstCallSite;
49+
}
50+
2851
/// Remap all locations reachable from the inlined blocks with CallSiteLoc
2952
/// locations with the provided caller location.
3053
static void
@@ -35,7 +58,7 @@ remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
3558
auto [it, inserted] = mappedLocations.try_emplace(loc);
3659
// Only query the attribute uniquer once per callsite attribute.
3760
if (inserted) {
38-
auto newLoc = CallSiteLoc::get(loc, callerLoc);
61+
LocationAttr newLoc = stackLocations(loc, callerLoc);
3962
it->getSecond() = newLoc;
4063
}
4164
return it->second;

mlir/test/Transforms/inlining.mlir

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,20 +74,51 @@ func.func @inline_with_multi_return() -> i32 {
7474
}
7575

7676
// Check that location information is updated for inlined instructions.
77-
func.func @func_with_locations(%c : i32) -> i32 {
77+
78+
#inline_stack1 = loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at "mysource3.cc":16:2)))
79+
#inline_stack2 = loc(callsite("mysource4.cc":55:4 at callsite("mysource5.cc":25:8 at "mysource6.cc":32:4)))
80+
81+
// INLINE-LOC-LABEL: func @func_with_file_locations
82+
func.func @func_with_file_locations(%c : i32) -> i32 {
7883
%b = arith.addi %c, %c : i32 loc("mysource.cc":10:8)
7984
return %b : i32 loc("mysource.cc":11:2)
8085
}
8186

82-
// INLINE-LOC-LABEL: func @inline_with_locations
83-
func.func @inline_with_locations(%arg0 : i32) -> i32 {
87+
// INLINE-LOC-LABEL: func @func_with_callsite_locations
88+
func.func @func_with_callsite_locations(%c : i32) -> i32 {
89+
%b = arith.addi %c, %c : i32 loc(#inline_stack1)
90+
return %b : i32 loc(#inline_stack1)
91+
}
92+
93+
// INLINE-LOC-LABEL: func @inline_func_with_file_locations
94+
func.func @inline_func_with_file_locations(%arg0 : i32) -> i32 {
8495
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32 loc(callsite("mysource.cc":10:8 at "mysource.cc":55:14))
85-
// INLINE-LOC-NEXT: return
96+
%0 = call @func_with_file_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
8697

87-
%0 = call @func_with_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
88-
return %0 : i32
98+
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
99+
// INLINE-LOC-SAME: loc(callsite("mysource.cc":10:8 at callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6
100+
// INLINE-LOC-SAME: at "mysource3.cc":16:2))))
101+
%1 = call @func_with_file_locations(%0) : (i32) -> i32 loc(#inline_stack1)
102+
103+
// INLINE-LOC-NEXT: return
104+
return %1 : i32
89105
}
90106

107+
// INLINE-LOC-LABEL: func @inline_func_with_callsite_locations
108+
func.func @inline_func_with_callsite_locations(%arg0 : i32) -> i32 {
109+
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
110+
// INLINE-LOC-SAME: loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at callsite("mysource3.cc":16:2
111+
// INLINE-LOC-SAME: at "mysource.cc":10:8))))
112+
%0 = call @func_with_callsite_locations(%arg0) : (i32) -> i32 loc("mysource.cc":10:8)
113+
114+
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
115+
// INLINE-LOC-SAME: loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at callsite("mysource3.cc":16:2
116+
// INLINE-LOC-SAME: at callsite("mysource4.cc":55:4 at callsite("mysource5.cc":25:8 at "mysource6.cc":32:4))))))
117+
%1 = call @func_with_callsite_locations(%0) : (i32) -> i32 loc(#inline_stack2)
118+
119+
// INLINE-LOC-NEXT: return
120+
return %1 : i32
121+
}
91122

92123
// Check that external function declarations are not inlined.
93124
func.func private @func_external()

0 commit comments

Comments
 (0)