[MLIR] Handle call site locations when inlining (#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.
This commit is contained in:
Vadim Curcă
2025-03-21 14:06:53 +01:00
committed by GitHub
parent 42d49a1a5b
commit 6b59b33358
2 changed files with 63 additions and 9 deletions

View File

@@ -25,6 +25,29 @@
using namespace mlir;
/// Combine `callee` location with `caller` location to create a stack that
/// represents the call chain.
/// If `callee` location is a `CallSiteLoc`, indicating an existing stack of
/// locations, the `caller` location is appended to the end of it, extending
/// the chain.
/// Otherwise, a single `CallSiteLoc` is created, representing a direct call
/// from `caller` to `callee`.
static LocationAttr stackLocations(Location callee, Location caller) {
Location lastCallee = callee;
SmallVector<CallSiteLoc> calleeInliningStack;
while (auto nextCallSite = dyn_cast<CallSiteLoc>(lastCallee)) {
calleeInliningStack.push_back(nextCallSite);
lastCallee = nextCallSite.getCaller();
}
CallSiteLoc firstCallSite = CallSiteLoc::get(lastCallee, caller);
for (CallSiteLoc currentCallSite : reverse(calleeInliningStack))
firstCallSite =
CallSiteLoc::get(currentCallSite.getCallee(), firstCallSite);
return firstCallSite;
}
/// Remap all locations reachable from the inlined blocks with CallSiteLoc
/// locations with the provided caller location.
static void
@@ -35,7 +58,7 @@ remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
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;

View File

@@ -74,20 +74,51 @@ func.func @inline_with_multi_return() -> i32 {
}
// Check that location information is updated for inlined instructions.
func.func @func_with_locations(%c : i32) -> i32 {
#inline_stack1 = loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at "mysource3.cc":16:2)))
#inline_stack2 = loc(callsite("mysource4.cc":55:4 at callsite("mysource5.cc":25:8 at "mysource6.cc":32:4)))
// INLINE-LOC-LABEL: func @func_with_file_locations
func.func @func_with_file_locations(%c : i32) -> i32 {
%b = arith.addi %c, %c : i32 loc("mysource.cc":10:8)
return %b : i32 loc("mysource.cc":11:2)
}
// INLINE-LOC-LABEL: func @inline_with_locations
func.func @inline_with_locations(%arg0 : i32) -> i32 {
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32 loc(callsite("mysource.cc":10:8 at "mysource.cc":55:14))
// INLINE-LOC-NEXT: return
%0 = call @func_with_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
return %0 : i32
// INLINE-LOC-LABEL: func @func_with_callsite_locations
func.func @func_with_callsite_locations(%c : i32) -> i32 {
%b = arith.addi %c, %c : i32 loc(#inline_stack1)
return %b : i32 loc(#inline_stack1)
}
// INLINE-LOC-LABEL: func @inline_func_with_file_locations
func.func @inline_func_with_file_locations(%arg0 : i32) -> i32 {
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32 loc(callsite("mysource.cc":10:8 at "mysource.cc":55:14))
%0 = call @func_with_file_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
// INLINE-LOC-SAME: loc(callsite("mysource.cc":10:8 at callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6
// INLINE-LOC-SAME: at "mysource3.cc":16:2))))
%1 = call @func_with_file_locations(%0) : (i32) -> i32 loc(#inline_stack1)
// INLINE-LOC-NEXT: return
return %1 : i32
}
// INLINE-LOC-LABEL: func @inline_func_with_callsite_locations
func.func @inline_func_with_callsite_locations(%arg0 : i32) -> i32 {
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
// INLINE-LOC-SAME: loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at callsite("mysource3.cc":16:2
// INLINE-LOC-SAME: at "mysource.cc":10:8))))
%0 = call @func_with_callsite_locations(%arg0) : (i32) -> i32 loc("mysource.cc":10:8)
// INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32
// INLINE-LOC-SAME: loc(callsite("mysource1.cc":10:8 at callsite("mysource2.cc":13:6 at callsite("mysource3.cc":16:2
// INLINE-LOC-SAME: at callsite("mysource4.cc":55:4 at callsite("mysource5.cc":25:8 at "mysource6.cc":32:4))))))
%1 = call @func_with_callsite_locations(%0) : (i32) -> i32 loc(#inline_stack2)
// INLINE-LOC-NEXT: return
return %1 : i32
}
// Check that external function declarations are not inlined.
func.func private @func_external()