[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.

Author: VadimCurca

mlir/lib/Transforms/Utils/InliningUtils.cpp

@@ -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;

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()