[memprof] Dump call site matching information

[kazutakahirata]

MemProfiler.cpp annotates the IR with the memory profile so that we
can later duplicate context. Dumping the call site matching
information here allows us to analyze how well we manage to annotate
the IR. Specifically, this patch dumps:

• the full stack ID (to identify the profile call stack)
• the index within the profile call stack where we start matching
• the size of InlinedCallStack

This way, we get to see what part of profile call stack we are
matching, not just one frame somewhere in the profile call stack.

Now, obtaining the full stack ID requires a little bit of refactoring.
This patch modifies the value type of LocHashToCallSites so that it
contains the full stack as well as the starting index of a match.
Essentially, this patch partially reverts:

commit 7c294eb
Author: Kazu Hirata [email protected]
Date: Sat Dec 14 00:03:27 2024 -0800

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Kazu Hirata (kazutakahirata)

Changes

MemProfiler.cpp annotates the IR with the memory profile so that we
can later duplicate context. Dumping the call site matching
information here allows us to analyze how well we manage to annotate
the IR. Specifically, this patch dumps:

• the full stack ID (to identify the profile call stack)
• the index within the profile call stack where we start matching
• the size of InlinedCallStack

This way, we get to see what part of profile call stack we are
matching, not just one frame somewhere in the profile call stack.

Now, obtaining the full stack ID requires a little bit of refactoring.
This patch modifies the value type of LocHashToCallSites so that it
contains the full stack as well as the starting index of a match.
Essentially, this patch partially reverts:

commit 7c294eb
Author: Kazu Hirata <[email protected]>
Date: Sat Dec 14 00:03:27 2024 -0800

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Instrumentation/MemProfiler.cpp (+15-6)

diff --git a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
index 91c48338d032089..66d7466a92c3b39 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
@@ -1034,13 +1034,15 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
   std::map<uint64_t, std::set<const AllocationInfo *>> LocHashToAllocInfo;
   // A hash function for std::unordered_set<ArrayRef<Frame>> to work.
   struct CallStackHash {
-    size_t operator()(ArrayRef<Frame> CS) const {
-      return computeFullStackId(CS);
+    size_t operator()(const std::pair<ArrayRef<Frame>, unsigned> &CS) const {
+      auto &[CallStack, Idx] = CS;
+      return computeFullStackId(ArrayRef<Frame>(CallStack).drop_front(Idx));
     }
   };
   // For the callsites we need to record slices of the frame array (see comments
   // below where the map entries are added).
-  std::map<uint64_t, std::unordered_set<ArrayRef<Frame>, CallStackHash>>
+  std::map<uint64_t, std::unordered_set<std::pair<ArrayRef<Frame>, unsigned>,
+                                        CallStackHash>>
       LocHashToCallSites;
   for (auto &AI : MemProfRec->AllocSites) {
     NumOfMemProfAllocContextProfiles++;
@@ -1058,7 +1060,7 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
     unsigned Idx = 0;
     for (auto &StackFrame : CS) {
       uint64_t StackId = computeStackId(StackFrame);
-      LocHashToCallSites[StackId].insert(ArrayRef<Frame>(CS).drop_front(Idx++));
+      LocHashToCallSites[StackId].emplace(CS, Idx++);
       ProfileHasColumns |= StackFrame.Column;
       // Once we find this function, we can stop recording.
       if (StackFrame.Function == FuncGUID)
@@ -1201,15 +1203,22 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
       // instruction's leaf location in the callsites map and not the allocation
       // map.
       assert(CallSitesIter != LocHashToCallSites.end());
-      for (auto CallStackIdx : CallSitesIter->second) {
+      for (auto &[ProfileCallStack, Idx] : CallSitesIter->second) {
         // If we found and thus matched all frames on the call, create and
         // attach call stack metadata.
-        if (stackFrameIncludesInlinedCallStack(CallStackIdx,
+        if (stackFrameIncludesInlinedCallStack(ProfileCallStack.drop_front(Idx),
                                                InlinedCallStack)) {
           NumOfMemProfMatchedCallSites++;
           addCallsiteMetadata(I, InlinedCallStack, Ctx);
           // Only need to find one with a matching call stack and add a single
           // callsite metadata.
+
+          // Dump call site matching information upon request.
+          if (ClPrintMemProfMatchInfo) {
+            uint64_t FullStackId = computeFullStackId(ProfileCallStack);
+            errs() << "MemProf callsite " << FullStackId << " " << Idx << " "
+                   << InlinedCallStack.size() << "\n";
+          }
           break;
         }
       }

[teresajohnson]

You should probably clarify in the description that this will just be the full profiled context of frames on the interior callsite (i.e. including all possible inlined frames). Unlike in the allocation case where we have the whole allocation context.

Also, can you add a test?

[teresajohnson]

Maybe include a bit more info in the message about the latter 2 numbers being printed.