[MLIR] Handle call site locations when inlining

[VadimCurca]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Vadim Curcă (VadimCurca)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/132247.diff

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+16-1)
• (modified) mlir/test/Transforms/inlining.mlir (+37-6)

diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0cae63c58ca7b..7ddcdb7592e9a 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -25,6 +25,21 @@
 
 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 new `CallSiteLoc` is created, representing a direct call from
+/// `caller` to `callee`.
+static LocationAttr stackLocations(Location callee, Location caller) {
+  if (auto calleeCallSite = dyn_cast<CallSiteLoc>(callee)) {
+    return CallSiteLoc::get(calleeCallSite.getCallee(),
+                            stackLocations(calleeCallSite.getCaller(), caller));
+  }
+  return CallSiteLoc::get(callee, caller);
+}
+
 /// Remap all locations reachable from the inlined blocks with CallSiteLoc
 /// locations with the provided caller location.
 static void
@@ -35,7 +50,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);
+      auto newLoc = stackLocations(loc, callerLoc);
       it->getSecond() = newLoc;
     }
     return it->second;
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 65ffaad1fa859..1ed08878430b5 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -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-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))
-  // INLINE-LOC-NEXT: return
+  %0 = call @func_with_file_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
 
-  %0 = call @func_with_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
-  return %0 : i32
+  // 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()

[definelicht]

LGTM

[gysit]

Thanks for this improvement!

LGTM