[MemProf] Fix inline propagation of memprof metadata

[tejohnson]

Summary

It isn't correct to always remove memprof metadata MIBs from the
original allocation call after inlining.

Let's say we have the following partial call graph:

C     D
 \   /
  v v
   B   E
   |  /
   v v
    A

where A contains an allocation call. If both contexts including B have
the same allocation behavior, the context in the memprof metadata on the
allocation will be pruned, and we will have 2 MIBs with contexts:
A,B and A,E.

Previously, if we inlined A into B we propagate the matching MIBs onto
the inlined allocation call in B' (A,B in this case), and remove it from
the original out of line allocation in A. This is correct if we have a
single round of bottom up inlining.

However, in the compiler we can have multiple invocations of the inliner
pass (e.g. LTO). We may also inline non-bottom up with an alternative
inliner such as the ModuleInliner. In that case, we could end up first
inlining B into C, without having inlined A into B. The call graph then
looks like:

    D
    |
    v
C'  B   E
 \  |  /
  v v v
    A

If we subsequently (perhaps on a later invocation of bottom up inlining)
inline A into B, the previous handling would propagate the memprof MIB
context A,B up into the inlined allocation in B', and remove it from the
original allocation in A. The propagation into B' is fine, however, by
removing it from A's allocation, we no longer reflect the context coming
from C'.

To fix this, simply prevent the removal of MIB from the original
allocation callsites.

[tejohnson]

Fixed the rendering of the call graphs in phabricator.

[davidxl]

Is the MIB removal purely for the purpose of reducing memory usage? It seems safe to always keep them.

[tejohnson]

In D140764#4019771, @davidxl wrote:

Is the MIB removal purely for the purpose of reducing memory usage? It seems safe to always keep them.

Yep, it was purely a "space" optimization. I noticed this problem when creating test cases for the context cloning patches I'm prepping.

[snehasish]

LGTM.

Please add a note to the commit message about the noncold -> notcold change in the tests.

[tejohnson]

In D140764#4019852, @snehasish wrote:

LGTM.

Please add a note to the commit message about the noncold -> notcold change in the tests.

Added a note to the commit message.

llvm/lib/Transforms/Utils/InlineFunction.cpp

static void updateMemprofMetadata(CallBase *CI,
   if (!MemprofMDAttached)
     // If we used a function attribute remove the callsite metadata as well.
     removeCallsiteMetadata(CI);
 }
 
 // Update the metadata on the inlined copy ClonedCall of a call OrigCall in the
 // inlined callee body, based on the callsite metadata InlinedCallsiteMD from
 // the call that was inlined.
-static void
-propagateMemProfHelper(const CallBase *OrigCall, CallBase *ClonedCall,
-                       MDNode *InlinedCallsiteMD,
-                       std::map<const CallBase *, std::vector<Metadata *>>
-                           &OrigCallToNewMemProfMDMap) {
+static void propagateMemProfHelper(const CallBase *OrigCall,
+                                   CallBase *ClonedCall,
+                                   MDNode *InlinedCallsiteMD) {
   MDNode *OrigCallsiteMD = ClonedCall->getMetadata(LLVMContext::MD_callsite);
   MDNode *ClonedCallsiteMD = nullptr;
   // Check if the call originally had callsite metadata, and update it for the
   // new call in the inlined body.
   if (OrigCallsiteMD) {
     // The cloned call's context is now the concatenation of the original call's
     // callsite metadata and the callsite metadata on the call where it was
     // inlined.
     ClonedCallsiteMD = MDNode::concatenate(OrigCallsiteMD, InlinedCallsiteMD);
     ClonedCall->setMetadata(LLVMContext::MD_callsite, ClonedCallsiteMD);
   }
 
   // Update any memprof metadata on the cloned call.
   MDNode *OrigMemProfMD = ClonedCall->getMetadata(LLVMContext::MD_memprof);
   if (!OrigMemProfMD)
     return;
   // We currently expect that allocations with memprof metadata also have
   // callsite metadata for the allocation's part of the context.
   assert(OrigCallsiteMD);
 
   // New call's MIB list.
   std::vector<Metadata *> NewMIBList;
-  // Updated MIB list for the original call in the out-of-line callee.
-  std::vector<Metadata *> UpdatedOrigMIBList;
 
   // For each MIB metadata, check if its call stack context starts with the
   // new clone's callsite metadata. If so, that MIB goes onto the cloned call in
   // the inlined body. If not, it stays on the out-of-line original call.
   for (auto &MIBOp : OrigMemProfMD->operands()) {
     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
     // Stack is first operand of MIB.
     MDNode *StackMD = getMIBStackNode(MIB);
     assert(StackMD);
     // See if the new cloned callsite context matches this profiled context.
     if (haveCommonPrefix(StackMD, ClonedCallsiteMD))
       // Add it to the cloned call's MIB list.
       NewMIBList.push_back(MIB);
-    else
-      // Keep it on the original call.
-      UpdatedOrigMIBList.push_back(MIB);
   }
   if (NewMIBList.empty()) {
     removeMemProfMetadata(ClonedCall);
     removeCallsiteMetadata(ClonedCall);
     return;
   }
-  if (NewMIBList.size() < OrigMemProfMD->getNumOperands()) {
-    assert(!UpdatedOrigMIBList.empty());
-    OrigCallToNewMemProfMDMap[OrigCall] = UpdatedOrigMIBList;
+  if (NewMIBList.size() < OrigMemProfMD->getNumOperands())
     updateMemprofMetadata(ClonedCall, NewMIBList);
-  } else
-    OrigCallToNewMemProfMDMap[OrigCall] = {};
 }
 
 // Update memprof related metadata (!memprof and !callsite) based on the
 // inlining of Callee into the callsite at CB. The updates include merging the
 // inlined callee's callsite metadata with that of the inlined call,
 // and moving the subset of any memprof contexts to the inlined callee
 // allocations if they match the new inlined call stack.
 // FIXME: Replace memprof metadata with function attribute if all MIB end up
 // having the same behavior. Do other context trimming/merging optimizations
 // too.
 static void
 propagateMemProfMetadata(Function *Callee, CallBase &CB,
                          bool ContainsMemProfMetadata,
                          const ValueMap<const Value *, WeakTrackingVH> &VMap) {
   MDNode *CallsiteMD = CB.getMetadata(LLVMContext::MD_callsite);
   // Only need to update if the inlined callsite had callsite metadata, or if
   // there was any memprof metadata inlined.
   if (!CallsiteMD && !ContainsMemProfMetadata)
     return;
 
   // Propagate metadata onto the cloned calls in the inlined callee.
-  // Can't update the original call using the VMap since it holds a const
-  // pointer, those will be updated in the subsequent loop.
-  std::map<const CallBase *, std::vector<Metadata *>> OrigCallToNewMemProfMDMap;
   for (const auto &Entry : VMap) {
     // See if this is a call that has been inlined and remapped, and not
     // simplified away in the process.
     auto *OrigCall = dyn_cast_or_null<CallBase>(Entry.first);
     auto *ClonedCall = dyn_cast_or_null<CallBase>(Entry.second);
     if (!OrigCall || !ClonedCall)
       continue;
     // If the inlined callsite did not have any callsite metadata, then it isn't
     // involved in any profiled call contexts, and we can remove any memprof
     // metadata on the cloned call.
     if (!CallsiteMD) {
       removeMemProfMetadata(ClonedCall);
       removeCallsiteMetadata(ClonedCall);
       continue;
     }
-    propagateMemProfHelper(OrigCall, ClonedCall, CallsiteMD,
-                           OrigCallToNewMemProfMDMap);
-  }
-
-  // Update memprof MD on calls within the original callee function to remove
-  // MIB with stacks that matched the inlined context (those moved to a new
-  // memprof MD on the inlined version of the call).
-  for (BasicBlock &BB : *Callee) {
-    for (Instruction &I : BB) {
-      CallBase *Call = dyn_cast<CallBase>(&I);
-      if (!Call || !OrigCallToNewMemProfMDMap.count(Call))
-        continue;
-      std::vector<Metadata *> &UpdatedMemProfMD =
-          OrigCallToNewMemProfMDMap[Call];
-      if (!UpdatedMemProfMD.empty())
-        updateMemprofMetadata(Call, UpdatedMemProfMD);
-      else {
-        removeMemProfMetadata(Call);
-        removeCallsiteMetadata(Call);
-      }
-    }
+    propagateMemProfHelper(OrigCall, ClonedCall, CallsiteMD);
   }
 }
 
 /// When inlining a call site that has !llvm.mem.parallel_loop_access,
 /// !llvm.access.group, !alias.scope or !noalias metadata, that metadata should
 /// be propagated to all memory-accessing cloned instructions.
 static void PropagateCallSiteMetadata(CallBase &CB, Function::iterator FStart,
                                       Function::iterator FEnd) {

llvm/test/Transforms/Inline/memprof_inline.ll

 source_filename = "memprof_inline.cc"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: mustprogress uwtable
 ; CHECK-LABEL: define dso_local noundef ptr @_Z3foov
 define dso_local noundef ptr @_Z3foov() #0 !dbg !39 {
 entry:
-  ; CHECK: call {{.*}} @_Znam
-  ; CHECK-NOT: !memprof
-  ; CHECK-NOT: !callsite
+  ;; We should keep the original memprof metadata intact.
+  ; CHECK: call {{.*}} @_Znam{{.*}} !memprof ![[ORIGMEMPROF:[0-9]+]]
   %call = call noalias noundef nonnull ptr @_Znam(i64 noundef 10) #6, !dbg !42, !memprof !43, !callsite !50
   ; CHECK-NEXT: ret
   ret ptr %call, !dbg !51
 }
 
 ; Function Attrs: nobuiltin allocsize(0)
 declare noundef nonnull ptr @_Znam(i64 noundef) #1
 

 !38 = !{!"clang version 15.0.0 (https://github.com/llvm/llvm-project.git e09c924f98ec157adeaa74819b0aec9a07a1b552)"}
 !39 = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov", scope: !1, file: !1, line: 4, type: !40, scopeLine: 4, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !41)
 !40 = !DISubroutineType(types: !41)
 !41 = !{}
 !42 = !DILocation(line: 5, column: 10, scope: !39)
 !43 = !{!44, !46, !48}
 !44 = !{!45, !"cold"}
 !45 = !{i64 -2458008693472584243, i64 7394638144382192936}
-!46 = !{!47, !"noncold"}
+!46 = !{!47, !"notcold"}
 !47 = !{i64 -2458008693472584243, i64 -8908997186479157179}
 !48 = !{!49, !"cold"}
 !49 = !{i64 -2458008693472584243, i64 -8079659623765193173}
+; CHECK: ![[ORIGMEMPROF]] = !{![[ORIGMIB1:[0-9]+]], ![[ORIGMIB2:[0-9]+]], ![[ORIGMIB3:[0-9]+]]}
+; CHECK: ![[ORIGMIB1]] = !{![[ORIGMIBSTACK1:[0-9]+]], !"cold"}
+; CHECK: ![[ORIGMIBSTACK1]] = !{i64 -2458008693472584243, i64 7394638144382192936}
+; CHECK: ![[ORIGMIB2]] = !{![[ORIGMIBSTACK2:[0-9]+]], !"notcold"}
+; CHECK: ![[ORIGMIBSTACK2]] = !{i64 -2458008693472584243, i64 -8908997186479157179}
+; CHECK: ![[ORIGMIB3]] = !{![[ORIGMIBSTACK3:[0-9]+]], !"cold"}
+; CHECK: ![[ORIGMIBSTACK3]] = !{i64 -2458008693472584243, i64 -8079659623765193173}
 !50 = !{i64 -2458008693472584243}
 !51 = !DILocation(line: 5, column: 3, scope: !39)
 !52 = distinct !DISubprogram(name: "foo2", linkageName: "_Z4foo2v", scope: !1, file: !1, line: 7, type: !40, scopeLine: 7, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !41)
 !53 = !DILocation(line: 8, column: 10, scope: !52)
 !54 = !{i64 -8079659623765193173}
 !55 = !DILocation(line: 8, column: 3, scope: !52)
 !56 = distinct !DISubprogram(name: "main", scope: !1, file: !1, line: 16, type: !40, scopeLine: 16, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !41)
 !57 = !DILocation(line: 17, column: 13, scope: !56)

llvm/test/Transforms/Inline/memprof_inline2.ll

 source_filename = "memprof_inline2.cc"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: mustprogress uwtable
 ; CHECK-LABEL: define dso_local noundef ptr @_Z3foov
 define dso_local noundef ptr @_Z3foov() #0 !dbg !39 {
 entry:
-  ;; We should still have memprof/callsite metadata for the non-inlined calls
-  ;; from main, but should have removed those from the inlined call in_Z4foo2v.
-  ;; CHECK: call {{.*}} @_Znam{{.*}} !memprof ![[ORIGMEMPROF:[0-9]+]]
+  ;; We should keep the original memprof metadata intact.
+  ; CHECK: call {{.*}} @_Znam{{.*}} !memprof ![[ORIGMEMPROF:[0-9]+]]
   %call = call noalias noundef nonnull ptr @_Znam(i64 noundef 10) #7, !dbg !42, !memprof !43, !callsite !52
   ret ptr %call, !dbg !53
 }
 
 ; Function Attrs: nobuiltin allocsize(0)
 declare noundef nonnull ptr @_Znam(i64 noundef) #1
 
 ;; Mark noinline so we don't inline into calls from bar and baz. We should end

 ;; callsite in foo, which isn't inlined into main due to the callsites in main
 ;; being annotated as noinline.
 ;; The second 2 are from the calls from foo2, which inlines its callsite to foo
 ;; but is not itself inlined into its callers. Therefore they get moved to a
 ;; new memprof metadata within foo2.
 !43 = !{!44, !46, !48, !50}
 !44 = !{!45, !"cold"}
 !45 = !{i64 -2458008693472584243, i64 7394638144382192936}
-!46 = !{!47, !"noncold"}
+!46 = !{!47, !"notcold"}
 !47 = !{i64 -2458008693472584243, i64 -8908997186479157179}
-!48 = !{!49, !"noncold"}
+!48 = !{!49, !"notcold"}
 !49 = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -4805294506621015872}
 !50 = !{!51, !"cold"}
 !51 = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -972865200055133905}
-; CHECK: ![[ORIGMEMPROF]] = !{![[ORIGMIB1:[0-9]+]], ![[ORIGMIB2:[0-9]+]]}
+; CHECK: ![[ORIGMEMPROF]] = !{![[ORIGMIB1:[0-9]+]], ![[ORIGMIB2:[0-9]+]], ![[ORIGMIB3:[0-9]+]], ![[ORIGMIB4:[0-9]+]]}
 ; CHECK: ![[ORIGMIB1]] = !{![[ORIGMIBSTACK1:[0-9]+]], !"cold"}
 ; CHECK: ![[ORIGMIBSTACK1]] = !{i64 -2458008693472584243, i64 7394638144382192936}
 ; CHECK: ![[ORIGMIB2]] = !{![[ORIGMIBSTACK2:[0-9]+]], !"notcold"}
 ; CHECK: ![[ORIGMIBSTACK2]] = !{i64 -2458008693472584243, i64 -8908997186479157179}
-; CHECK: ![[NEWMEMPROF]] = !{![[NEWMIB1:[0-9]+]], ![[NEWMIB2:[0-9]+]]}
-; CHECK: ![[NEWMIB1]] = !{![[NEWMIBSTACK1:[0-9]+]], !"notcold"}
-; CHECK: ![[NEWMIBSTACK1]] = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -4805294506621015872}
-; CHECK: ![[NEWMIB2]] = !{![[NEWMIBSTACK2:[0-9]+]], !"cold"}
-; CHECK: ![[NEWMIBSTACK2]] = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -972865200055133905}
+; CHECK: ![[ORIGMIB3]] = !{![[ORIGMIBSTACK3:[0-9]+]], !"notcold"}
+; CHECK: ![[ORIGMIBSTACK3]] = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -4805294506621015872}
+; CHECK: ![[ORIGMIB4]] = !{![[ORIGMIBSTACK4:[0-9]+]], !"cold"}
+; CHECK: ![[ORIGMIBSTACK4]] = !{i64 -2458008693472584243, i64 -8079659623765193173, i64 -972865200055133905}
+; CHECK: ![[NEWMEMPROF]] = !{![[ORIGMIB3:[0-9]+]], ![[ORIGMIB4:[0-9]+]]}
 ; CHECK: ![[NEWCALLSITE]] = !{i64 -2458008693472584243, i64 -8079659623765193173}
 !52 = !{i64 -2458008693472584243}
 !53 = !DILocation(line: 5, column: 3, scope: !39)
 !54 = distinct !DISubprogram(name: "foo2", linkageName: "_Z4foo2v", scope: !1, file: !1, line: 7, type: !40, scopeLine: 7, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !41)
 !55 = !DILocation(line: 8, column: 10, scope: !54)
 !56 = !{i64 -8079659623765193173}
 !57 = !DILocation(line: 8, column: 3, scope: !54)
 !58 = distinct !DISubprogram(name: "bar", linkageName: "_Z3barv", scope: !1, file: !1, line: 10, type: !40, scopeLine: 10, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !41)