[CSSPGO] Refactoring findIRAnchors

[wlei]

Summary

Address feedback in https://reviews.llvm.org/D158817. Since extractProbe can be used for both calliste and BB probe, we can leverage this to unify the callsite handling code.

[hoy]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2154–2158

nit: may be good to hoist this code into a lamda and share it with the line-number path.

[wlei]

addressing feedback

llvm/lib/Transforms/IPO/SampleProfile.cpp

2154–2158

Good point, done!

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2170–2172

Actually can this be unified with FindTopLevelFrame? For not inlined calls, DIL->getSubprogramLinkageName should be equivalent to FindCalleeName(dyn_cast<CallBase>(&I))?

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2170–2172

Good point, but to clarify:
DIL->getSubprogramLinkageName is the current/top-level Frame.
FindCalleeName is the current callee name, it's the second top Frame.

Maybe my variable name is misleading in the new patch. Here is actually not FindTopLevelFrame, For example: "main:1 @ foo", the term "top-level" frame is "main:1", but here what it returns is {1, foo} , i,e {the top-level frame's location, the second top's frame's func name}

uh..I need to fix it.

That said, maybe we still can merge FindTopLevelFrame and FindCalleeName , like check if it's inlined code inside of the lambda, let me think.

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2170–2172

OK.. I'm done one version, please take a look https://reviews.llvm.org/D159221

I ran into some bugs and spent some times debugging, the tricky parts are :

1. For pseudo-probe mode, a) for inline code, we need to use getCallSiteIdentifier to extract the callsite but b) for non-inlined code, we need to use the probe id.

2. The pseudo instruction is a CallBase and the getCalledFunction will return a non-empty callee name "llvm.pseudoprobe" which actually should be a empty StringRef.

I then made it work, but as it introduces more complexity, maybe it's still not what we want, any thoughts on this?

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2114

Maybe assert that DIL->getInlinedAt() is not null first. For not inlined DIL, this will not get what we want. Assert helps make it clear.

With the assert, the loop can also changed to do-while to make it explicit. I find do-while slightly more readable for this kind of cases.

do {
  PrevDIL = DIL;
  DIL = DIL->getInlinedAt();
} while (DIL->getInlinedAt())

2124

nit: this does not involve search/find, so maybe just GetCanonicalCalleeName.

2170–2172

I guess, FindTopLevelFrame can be called FindTopLevelInlinedCallsite? It would also help to explain in a comment what it does with an example, like the main:1 @ foo one you mentioned above.

For pseudo-probe mode, a) for inline code, we need to use getCallSiteIdentifier to extract the callsite but b) for non-inlined code, we need to use the probe id.

Thanks for clarification - I think that answered my original question, and the answer is no. Getting callsite (loc, callee_name pair) from debug info is only doable for inlined callsite, so for not inlined code, we need to get callee from the call instructions.

I looked at https://reviews.llvm.org/D159221, i think that with better naming, the version in this patch is actually more readable. wdyt?

[wlei]

addressing feedback

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2114

Sounds good, done.

2170–2172

thanks for the suggestion, renamed and updated the comments.

[wenlei]

lgtm, thanks.

llvm/lib/Transforms/IPO/SampleProfile.cpp

MatchingManager =
         std::make_unique<SampleProfileMatcher>(M, *Reader, ProbeManager.get());
   }
 
   return true;
 }
 
 void SampleProfileMatcher::findIRAnchors(
     const Function &F, std::map<LineLocation, StringRef> &IRAnchors) {
-  // Extract profile matching anchors in the IR.
-  for (auto &BB : F) {
-    for (auto &I : BB) {
-      // TODO: To support AutoFDO, we need to parse all the non-call
-      // instructions to extract the line-number based locations. For
-      // pseudo-probe mode, since each block is instrumented with one probe
-      // inst, parsing probe inst is enough.
-      if (FunctionSamples::ProfileIsProbeBased && isa<PseudoProbeInst>(&I)) {
-        std::optional<PseudoProbe> Probe = extractProbe(I);
-        assert(Probe &&
-               "Probe should not be null for pseudo-probe instruction");
-        // Flatten inlined IR for the matching. Recover the original callsite
-        // and call target by analyzing the inline frames from the debug info.
-        if (DILocation *DIL = I.getDebugLoc()) {
-          if (DIL->getInlinedAt()) {
-            // Find the top-level inline frame.
-            const DILocation *PrevDIL = DIL;
-            for (; DIL->getInlinedAt(); DIL = DIL->getInlinedAt())
-              PrevDIL = DIL;
+  // For inlined code, recover the original callsite and callee by finding the
+  // top-level inline frame. e.g. For frame stack "main:1 @ foo:2 @ bar:3", the
+  // top-level frame is "main:1", the callsite is "1" and the callee is "foo".
+  auto FindTopLevelInlinedCallsite = [](const DILocation *DIL) {
+    assert((DIL && DIL->getInlinedAt()) && "No inlined callsite");
+    const DILocation *PrevDIL = nullptr;
+    do {
+      DIL = DIL->getInlinedAt();
+    } while (DIL->getInlinedAt());
 
-            LineLocation Callsite = FunctionSamples::getCallSiteIdentifier(DIL);
-            StringRef CalleeName = PrevDIL->getSubprogramLinkageName();
-            IRAnchors.emplace(Callsite, CalleeName);
-          } else
-            IRAnchors.emplace(LineLocation(Probe->Id, 0), StringRef());
-        }
-      }
-
-      if (!isa<CallBase>(&I) || isa<IntrinsicInst>(&I))
-        continue;
-
-      const auto *CB = dyn_cast<CallBase>(&I);
-      if (auto &DLoc = I.getDebugLoc()) {
-        // Skip the inlined callsite. For pseudo probe mode, extracting inline
-        // info from the probe inst is enough.
-        if (DLoc.getInlinedAt())
-          continue;
+    return std::make_pair(Callsite, CalleeName);
+  };
 
-        LineLocation IRCallsite = FunctionSamples::getCallSiteIdentifier(DLoc);
-        StringRef CalleeName = UnknownIndirectCallee;
-        if (Function *Callee = CB->getCalledFunction())
-          CalleeName = FunctionSamples::getCanonicalFnName(Callee->getName());
+  auto GetCanonicalCalleeName = [](const CallBase *CB) {
+    return CalleeName;
+  };
+
+  // Extract profile matching anchors in the IR.
+  for (auto &BB : F) {
+    for (auto &I : BB) {
+      DILocation *DIL = I.getDebugLoc();
+      if (!DIL)
+        continue;
+
+      if (FunctionSamples::ProfileIsProbeBased) {
+        if (auto Probe = extractProbe(I)) {
+          // Flatten inlined IR for the matching.
+          if (DIL->getInlinedAt()) {
+            IRAnchors.emplace(FindTopLevelInlinedCallsite(DIL));
+          } else {
+            // Use empty StringRef for basic block probe.
+            StringRef CalleeName;
+            if (const auto *CB = dyn_cast<CallBase>(&I)) {
+              // Skip the probe inst whose callee name is "llvm.pseudoprobe".
+              if (!isa<IntrinsicInst>(&I))
+                CalleeName = GetCanonicalCalleeName(CB);
+            }
+            IRAnchors.emplace(LineLocation(Probe->Id, 0), CalleeName);
+          }
+        }
+      } else {
+        // TODO: For line-number based profile(AutoFDO), currently only support
+        // find callsite anchors. In future, we need to parse all the non-call
+        // instructions to extract the line locations for profile matching.
+        if (!isa<CallBase>(&I) || isa<IntrinsicInst>(&I))
+          continue;
 
-        // Force to overwrite the callee name in case any non-call location was
-        // written before.
-        auto R = IRAnchors.emplace(IRCallsite, CalleeName);
-        R.first->second = CalleeName;
-        assert((!FunctionSamples::ProfileIsProbeBased || R.second ||
-                R.first->second == CalleeName) &&
-               "Overwrite non-call or different callee name location for "
-               "pseudo probe callsite");
+        if (DIL->getInlinedAt()) {
+          IRAnchors.emplace(FindTopLevelInlinedCallsite(DIL));
+        } else {
+          LineLocation Callsite = FunctionSamples::getCallSiteIdentifier(DIL);
+          StringRef CalleeName = GetCanonicalCalleeName(dyn_cast<CallBase>(&I));
+          IRAnchors.emplace(Callsite, CalleeName);
+        }
       }
     }
   }
 }
 
 void SampleProfileMatcher::countMismatchedSamples(const FunctionSamples &FS) {
   const auto *FuncDesc = ProbeManager->getDesc(FS.getName());
   // Skip the function that is external or renamed.