[llvm-profgen] Decoding pseudo probe for profiled function only.

Complete pseudo probes decoding can result in large memory usage. In practice only a small porting of the decoded probes are used in profile generation. I'm changing the full decoding mode to be decoding for profiled functions only, though we still do a full scan of the .pseudoprobe section due to a missing table-of-content but we don't have to build the in-memory data structure for functions not sampled.

To build the in-memory data structure for profiled functions only, I'm rewriting the previous non-recursive probe decoding logic to be recursive. This is easy to read and maintain.

I also have to change the previous representation of unsymbolized context from probe-based stack to address-based stack since the profiled functions are unknown yet by the time of virtual unwinding. The address-based stack will be converted to probe-based stack after virtual unwinding and on-demand probe decoding.

I'm seeing 20GB memory is saved for one of our internal large service.

Reviewed By: wenlei

Differential Revision: https://reviews.llvm.org/D121643

Author: htyu

llvm/include/llvm/MC/MCPseudoProbe.h

@@ -55,6 +55,7 @@
 #include <tuple>
 #include <type_traits>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace llvm {
@@ -353,6 +354,15 @@ class MCPseudoProbeDecoder {
   // Decode pseudo_probe section to build address to probes map.
   bool buildAddress2ProbeMap(const uint8_t *Start, std::size_t Size);
 
+  // Decode pseudo_probe section to build address to probes map for specifed
+  // functions only.
+  bool buildAddress2ProbeMap(const uint8_t *Start, std::size_t Size,
+                             std::unordered_set<uint64_t> &GuildFilter);
+
+  bool buildAddress2ProbeMap(MCDecodedPseudoProbeInlineTree *Cur,
+                             uint64_t &LastAddr,
+                             std::unordered_set<uint64_t> &GuildFilter);
+
   // Print pseudo_probe_desc section info
   void printGUID2FuncDescMap(raw_ostream &OS);
 

llvm/lib/MC/MCPseudoProbe.cpp

@@ -358,8 +358,9 @@ bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start,
   return true;
 }
 
-bool MCPseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
-                                                 std::size_t Size) {
+bool MCPseudoProbeDecoder::buildAddress2ProbeMap(
+    MCDecodedPseudoProbeInlineTree *Cur, uint64_t &LastAddr,
+    std::unordered_set<uint64_t> &GuildFilter) {
   // The pseudo_probe section encodes an inline forest and each tree has a
   // format like:
   //  FUNCTION BODY (one for each uninlined function present in the text
@@ -390,101 +391,110 @@ bool MCPseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
   //           FUNCTION BODY
   //             A FUNCTION BODY entry describing the inlined function.
 
-  Data = Start;
-  End = Data + Size;
-
-  MCDecodedPseudoProbeInlineTree *Root = &DummyInlineRoot;
-  MCDecodedPseudoProbeInlineTree *Cur = &DummyInlineRoot;
-  uint64_t LastAddr = 0;
   uint32_t Index = 0;
-  // A DFS-based decoding
-  while (Data < End) {
-    if (Root == Cur) {
-      // Use a sequential id for top level inliner.
-      Index = Root->getChildren().size();
-    } else {
-      // Read inline site for inlinees
-      auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
-      if (!ErrorOrIndex)
-        return false;
-      Index = std::move(*ErrorOrIndex);
-    }
+  if (Cur == &DummyInlineRoot) {
+    // Use a sequential id for top level inliner.
+    Index = Cur->getChildren().size();
+  } else {
+    // Read inline site for inlinees
+    auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
+    if (!ErrorOrIndex)
+      return false;
+    Index = std::move(*ErrorOrIndex);
+  }
+
+  // Read guid
+  auto ErrorOrCurGuid = readUnencodedNumber<uint64_t>();
+  if (!ErrorOrCurGuid)
+    return false;
+  uint64_t Guid = std::move(*ErrorOrCurGuid);
+
+  // Decide if top-level node should be disgarded.
+  if (Cur == &DummyInlineRoot && !GuildFilter.empty() &&
+      !GuildFilter.count(Guid))
+    Cur = nullptr;
+
+  // If the incoming node is null, all its children nodes should be disgarded.
+  if (Cur) {
     // Switch/add to a new tree node(inlinee)
     Cur = Cur->getOrAddNode(std::make_tuple(Cur->Guid, Index));
-    // Read guid
-    auto ErrorOrCurGuid = readUnencodedNumber<uint64_t>();
-    if (!ErrorOrCurGuid)
-      return false;
-    Cur->Guid = std::move(*ErrorOrCurGuid);
-    // Read number of probes in the current node.
-    auto ErrorOrNodeCount = readUnsignedNumber<uint32_t>();
-    if (!ErrorOrNodeCount)
+    Cur->Guid = Guid;
+  }
+
+  // Read number of probes in the current node.
+  auto ErrorOrNodeCount = readUnsignedNumber<uint32_t>();
+  if (!ErrorOrNodeCount)
+    return false;
+  uint32_t NodeCount = std::move(*ErrorOrNodeCount);
+  // Read number of direct inlinees
+  auto ErrorOrCurChildrenToProcess = readUnsignedNumber<uint32_t>();
+  if (!ErrorOrCurChildrenToProcess)
+    return false;
+  // Read all probes in this node
+  for (std::size_t I = 0; I < NodeCount; I++) {
+    // Read index
+    auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
+    if (!ErrorOrIndex)
       return false;
-    uint32_t NodeCount = std::move(*ErrorOrNodeCount);
-    // Read number of direct inlinees
-    auto ErrorOrCurChildrenToProcess = readUnsignedNumber<uint32_t>();
-    if (!ErrorOrCurChildrenToProcess)
+    uint32_t Index = std::move(*ErrorOrIndex);
+    // Read type | flag.
+    auto ErrorOrValue = readUnencodedNumber<uint8_t>();
+    if (!ErrorOrValue)
       return false;
-    Cur->ChildrenToProcess = std::move(*ErrorOrCurChildrenToProcess);
-    // Read all probes in this node
-    for (std::size_t I = 0; I < NodeCount; I++) {
-      // Read index
-      auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
-      if (!ErrorOrIndex)
+    uint8_t Value = std::move(*ErrorOrValue);
+    uint8_t Kind = Value & 0xf;
+    uint8_t Attr = (Value & 0x70) >> 4;
+    // Read address
+    uint64_t Addr = 0;
+    if (Value & 0x80) {
+      auto ErrorOrOffset = readSignedNumber<int64_t>();
+      if (!ErrorOrOffset)
         return false;
-      uint32_t Index = std::move(*ErrorOrIndex);
-      // Read type | flag.
-      auto ErrorOrValue = readUnencodedNumber<uint8_t>();
-      if (!ErrorOrValue)
+      int64_t Offset = std::move(*ErrorOrOffset);
+      Addr = LastAddr + Offset;
+    } else {
+      auto ErrorOrAddr = readUnencodedNumber<int64_t>();
+      if (!ErrorOrAddr)
         return false;
-      uint8_t Value = std::move(*ErrorOrValue);
-      uint8_t Kind = Value & 0xf;
-      uint8_t Attr = (Value & 0x70) >> 4;
-      // Read address
-      uint64_t Addr = 0;
-      if (Value & 0x80) {
-        auto ErrorOrOffset = readSignedNumber<int64_t>();
-        if (!ErrorOrOffset)
-          return false;
-        int64_t Offset = std::move(*ErrorOrOffset);
-        Addr = LastAddr + Offset;
-      } else {
-        auto ErrorOrAddr = readUnencodedNumber<int64_t>();
-        if (!ErrorOrAddr)
-          return false;
-        Addr = std::move(*ErrorOrAddr);
-      }
+      Addr = std::move(*ErrorOrAddr);
+    }
+
+    if (Cur) {
       // Populate Address2ProbesMap
       auto &Probes = Address2ProbesMap[Addr];
       Probes.emplace_back(Addr, Cur->Guid, Index, PseudoProbeType(Kind), Attr,
                           Cur);
       Cur->addProbes(&Probes.back());
-      LastAddr = Addr;
     }
+    LastAddr = Addr;
+  }
 
-    // Look for the parent for the next node by subtracting the current
-    // node count from tree counts along the parent chain. The first node
-    // in the chain that has a non-zero tree count is the target.
-    while (Cur != Root) {
-      if (Cur->ChildrenToProcess == 0) {
-        Cur = static_cast<MCDecodedPseudoProbeInlineTree *>(Cur->Parent);
-        if (Cur != Root) {
-          assert(Cur->ChildrenToProcess > 0 &&
-                 "Should have some unprocessed nodes");
-          Cur->ChildrenToProcess -= 1;
-        }
-      } else {
-        break;
-      }
-    }
+  uint32_t ChildrenToProcess = std::move(*ErrorOrCurChildrenToProcess);
+  for (uint32_t I = 0; I < ChildrenToProcess; I++) {
+    buildAddress2ProbeMap(Cur, LastAddr, GuildFilter);
   }
 
+  return true;
+}
+
+bool MCPseudoProbeDecoder::buildAddress2ProbeMap(
+    const uint8_t *Start, std::size_t Size,
+    std::unordered_set<uint64_t> &GuildFilter) {
+  Data = Start;
+  End = Data + Size;
+  uint64_t LastAddr = 0;
+  while (Data < End)
+    buildAddress2ProbeMap(&DummyInlineRoot, LastAddr, GuildFilter);
   assert(Data == End && "Have unprocessed data in pseudo_probe section");
-  assert(Cur == Root &&
-         " Cur should point to root when the forest is fully built up");
   return true;
 }
 
+bool MCPseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
+                                                 std::size_t Size) {
+  std::unordered_set<uint64_t> GuildFilter;
+  return buildAddress2ProbeMap(Start, Size, GuildFilter);
+}
+
 void MCPseudoProbeDecoder::printGUID2FuncDescMap(raw_ostream &OS) {
   OS << "Pseudo Probe Desc:\n";
   // Make the output deterministic

llvm/test/tools/llvm-profgen/noinline-cs-pseudoprobe.test

@@ -24,14 +24,14 @@
 ; CHECK-NEXT: 4: 15
 ; CHECK-NEXT: !CFGChecksum: 72617220756
 
-; CHECK-UNWINDER:      [main:2]
+; CHECK-UNWINDER:      [0x7f4]
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   79e-7bf:15
 ; CHECK-UNWINDER-NEXT:   7c4-7cf:15
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7bf->760:15
 ; CHECK-UNWINDER-NEXT:   7cf->79e:16
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:8]
+; CHECK-UNWINDER-NEXT: [0x7f4 @ 0x7bf]
 ; CHECK-UNWINDER-NEXT:   1
 ; CHECK-UNWINDER-NEXT:   760-77f:15
 ; CHECK-UNWINDER-NEXT:   1

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

@@ -123,7 +123,7 @@
 ; CHECK:  6: 1 fa:1
 ; CHECK:  !CFGChecksum: 563022570642068
 
-; CHECK-UNWINDER:      [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5]
+; CHECK-UNWINDER:      [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab]
 ; CHECK-UNWINDER-NEXT:   3
 ; CHECK-UNWINDER-NEXT:   7a0-7a7:1
 ; CHECK-UNWINDER-NEXT:   7a0-7ab:3
@@ -132,33 +132,33 @@
 ; CHECK-UNWINDER-NEXT:   7a7->7b2:1
 ; CHECK-UNWINDER-NEXT:   7ab->7a0:4
 ; CHECK-UNWINDER-NEXT:   7b5->7c0:1
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6]
+; CHECK-UNWINDER-NEXT: [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab @ 0x7b5]
 ; CHECK-UNWINDER-NEXT:   1
 ; CHECK-UNWINDER-NEXT:   7c0-7d4:1
 ; CHECK-UNWINDER-NEXT:   1
 ; CHECK-UNWINDER-NEXT:   7d4->7c0:1
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8]
+; CHECK-UNWINDER-NEXT: [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab @ 0x7b5 @ 0x7d4]
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7c0-7cd:1
 ; CHECK-UNWINDER-NEXT:   7db-7e0:1
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7cd->7db:1
 ; CHECK-UNWINDER-NEXT:   7e0->7a0:1
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7]
+; CHECK-UNWINDER-NEXT: [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab @ 0x7b5 @ 0x7d4 @ 0x7e0]
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7a0-7a7:1
 ; CHECK-UNWINDER-NEXT:   7b2-7b5:1
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7a7->7b2:1
 ; CHECK-UNWINDER-NEXT:   7b5->7c0:1
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6]
+; CHECK-UNWINDER-NEXT: [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab @ 0x7b5 @ 0x7d4 @ 0x7e0 @ 0x7b5]
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7c0-7cd:2
 ; CHECK-UNWINDER-NEXT:   7db-7e0:1
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7cd->7db:2
 ; CHECK-UNWINDER-NEXT:   7e0->7a0:1
-; CHECK-UNWINDER-NEXT: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7]
+; CHECK-UNWINDER-NEXT: [0x842 @ 0x7d4 @ 0x7e0 @ 0x7ab @ 0x7b5 @ 0x7d4 @ 0x7e0 @ 0x7b5 @ 0x7e0]
 ; CHECK-UNWINDER-NEXT:   2
 ; CHECK-UNWINDER-NEXT:   7a0-7a7:1
 ; CHECK-UNWINDER-NEXT:   7b2-7b5:1

llvm/tools/llvm-profgen/PerfReader.cpp

@@ -179,17 +179,12 @@ std::shared_ptr<StringBasedCtxKey> FrameStack::getContextKey() {
   return KeyStr;
 }
 
-std::shared_ptr<ProbeBasedCtxKey> ProbeStack::getContextKey() {
-  std::shared_ptr<ProbeBasedCtxKey> ProbeBasedKey =
-      std::make_shared<ProbeBasedCtxKey>();
-  for (auto CallProbe : Stack) {
-    ProbeBasedKey->Probes.emplace_back(CallProbe);
-  }
-  CSProfileGenerator::compressRecursionContext<const MCDecodedPseudoProbe *>(
-      ProbeBasedKey->Probes);
-  CSProfileGenerator::trimContext<const MCDecodedPseudoProbe *>(
-      ProbeBasedKey->Probes);
-  return ProbeBasedKey;
+std::shared_ptr<AddrBasedCtxKey> AddressStack::getContextKey() {
+  std::shared_ptr<AddrBasedCtxKey> KeyStr = std::make_shared<AddrBasedCtxKey>();
+  KeyStr->Context = Stack;
+  CSProfileGenerator::compressRecursionContext<uint64_t>(KeyStr->Context);
+  CSProfileGenerator::trimContext<uint64_t>(KeyStr->Context);
+  return KeyStr;
 }
 
 template <typename T>
@@ -252,8 +247,8 @@ void VirtualUnwinder::collectSamplesFromFrameTrie(
 void VirtualUnwinder::collectSamplesFromFrameTrie(
     UnwindState::ProfiledFrame *Cur) {
   if (Binary->usePseudoProbes()) {
-    ProbeStack Stack(Binary);
-    collectSamplesFromFrameTrie<ProbeStack>(Cur, Stack);
+    AddressStack Stack(Binary);
+    collectSamplesFromFrameTrie<AddressStack>(Cur, Stack);
   } else {
     FrameStack Stack(Binary);
     collectSamplesFromFrameTrie<FrameStack>(Cur, Stack);
@@ -461,14 +456,17 @@ static std::string getContextKeyStr(ContextKey *K,
                                     const ProfiledBinary *Binary) {
   if (const auto *CtxKey = dyn_cast<StringBasedCtxKey>(K)) {
     return SampleContext::getContextString(CtxKey->Context);
-  } else if (const auto *CtxKey = dyn_cast<ProbeBasedCtxKey>(K)) {
-    SampleContextFrameVector ContextStack;
-    for (const auto *Probe : CtxKey->Probes) {
-      Binary->getInlineContextForProbe(Probe, ContextStack, true);
+  } else if (const auto *CtxKey = dyn_cast<AddrBasedCtxKey>(K)) {
+    std::ostringstream OContextStr;
+    for (uint32_t I = 0; I < CtxKey->Context.size(); I++) {
+      if (OContextStr.str().size())
+        OContextStr << " @ ";
+      OContextStr << "0x"
+                  << to_hexString(
+                         Binary->virtualAddrToOffset(CtxKey->Context[I]),
+                         false);
     }
-    // Probe context key at this point does not have leaf probe, so do not
-    // include the leaf inline location.
-    return SampleContext::getContextString(ContextStack, true);
+    return OContextStr.str();
   } else {
     llvm_unreachable("unexpected key type");
   }