[llvm] [memprof] Add CallStackRadixTreeBuilder (PR #93784)

via llvm-commits llvm-commits at lists.llvm.org
Thu May 30 01:31:30 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-pgo

Author: Kazu Hirata (kazutakahirata)

<details>
<summary>Changes</summary>

Call stacks are a huge portion of the MemProf profile, taking up 70+%
of the profile file size.

This patch implements a radix tree to compress call stacks, which are
known to have long common prefixes.  Specifically,
CallStackRadixTreeBuilder, introduced in this patch, takes call stacks
in the MemProf profile, sorts them in the dictionary order to maximize
the common prefix between adjacent call stacks, and then encodes a
radix tree into a single array that is ready for serialization.

The resulting radix array is essentially a concatenation of call stack
arrays, each encoded with its length followed by the payload, except
that these arrays contain "instructions" like "skip 7 elements
forward" to borrow common prefixes from other call stacks.

This patch does not integrate with the MemProf
serialization/deserialization infrastructure yet.  Once integrated,
the radix tree is expected to roughly halve the file size of the
MemProf profile.

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


3 Files Affected:

- (modified) llvm/include/llvm/ProfileData/MemProf.h (+100) 
- (modified) llvm/lib/ProfileData/MemProf.cpp (+155) 
- (modified) llvm/unittests/ProfileData/MemProfTest.cpp (+107) 


``````````diff
diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index d44a2d1e2fb11..16e6fa71020e8 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -849,6 +849,106 @@ struct IndexedMemProfData {
   llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> CallStackData;
 };
 
+// Construct a radix tree of call stacks.
+//
+// A set of call stacks might look like:
+//
+// CallStackId 1:  f1 -> f2 -> f3
+// CallStackId 2:  f1 -> f2 -> f4 -> f5
+// CallStackId 3:  f1 -> f2 -> f4 -> f6
+// CallStackId 4:  f7 -> f8 -> f9
+//
+// where each fn refers to a stack frame.
+//
+// Since we expect a lot of common prefixes, we can compress the call stacks
+// into a radix tree like:
+//
+// CallStackId 1:  f1 -> f2 -> f3
+//                       |
+// CallStackId 2:        +---> f4 -> f5
+//                             |
+// CallStackId 3:              +---> f6
+//
+// CallStackId 4:  f7 -> f8 -> f9
+//
+// Now, we are interested in retrieving call stacks for a given CallStackId, so
+// we just need a pointer from a given call stack to its parent.  For example,
+// CallStackId 2 would point to CallStackId 1 as a parent.
+//
+// We serialize the radix tree above into a single array along with the length
+// of each call stack and pointers to the parent call stacks.
+//
+// Index:              0  1  2  3  4  5  6  7  8  9 10 11 12 13 14
+// Array:             L3 f9 f8 f7 L4 f6 J3 L4 f5 f4 J3 L3 f3 f2 f1
+//                     ^           ^        ^           ^
+//                     |           |        |           |
+// CallStackId 4:  0 --+           |        |           |
+// CallStackId 3:  4 --------------+        |           |
+// CallStackId 2:  7 -----------------------+           |
+// CallStackId 1: 11 -----------------------------------+
+//
+// - LN indicates the length of a call stack, encoded as ordinary integer N.
+//
+// - JN indicates a pointer to the parent, encoded as -N.
+//
+// For example, if we are decoding CallStackId 2, we start a forward traversal
+// at Index 7, noting the call stack length of 4 and obtaining f5 and f4.  When
+// we see J3 at Index 10, we resume a forward traversal at Index 13 = 10 + 3,
+// picking up f2 and f1.  We are done after collecting 4 frames as indicated at
+// the beginning of the traversal.
+//
+// On-disk IndexedMemProfRecord will refer to call stacks by their indexes into
+// the radix tree array, so we do not explicitly encode mappings like:
+// "CallStackId 1 -> 11".
+class CallStackRadixTreeBuilder {
+  // The radix tree array.
+  std::vector<uint32_t> RadixArray;
+
+  // Mapping from CallStackIds to indexes into RadixArray.
+  llvm::DenseMap<CallStackId, uint32_t> CallStackPos;
+
+  // The indexes within RadixArray of the last call stack's frames encoded
+  // satisfying:
+  //
+  //   RadixArray[Indexes[I]] == (*Prev)[Prev->size() - I - 1]
+  //
+  // where Prev is one of the parameters to build.
+  std::vector<uint32_t> Indexes;
+
+  using CSIdPair = std::pair<CallStackId, llvm::SmallVector<FrameId> *>;
+
+  // Returns the sorted list of call stacks.
+  std::vector<CSIdPair>
+  sortCallStacks(llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
+                     &MemProfCallStackData);
+
+  // Encode a call stack into RadixArray.  Return the starting index within
+  // RadixArray.
+  uint32_t
+  encodeCallStack(const llvm::SmallVector<FrameId> *CallStack,
+                  const llvm::SmallVector<FrameId> *Prev,
+                  const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes);
+
+public:
+  CallStackRadixTreeBuilder() = default;
+
+  void build(llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
+                 &MemProfCallStackData,
+             const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes);
+
+  const std::vector<uint32_t> &getRadixArray() const { return RadixArray; }
+
+  const llvm::DenseMap<CallStackId, uint32_t> &getCallStackPos() const {
+    return CallStackPos;
+  }
+};
+
+llvm::DenseMap<CallStackId, uint32_t>
+writeCallStackRadixTree(raw_ostream &OS,
+                        llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
+                            &MemProfCallStackData,
+                        llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes);
+
 // Verify that each CallStackId is computed with hashCallStack.  This function
 // is intended to help transition from CallStack to CSId in
 // IndexedAllocationInfo.
diff --git a/llvm/lib/ProfileData/MemProf.cpp b/llvm/lib/ProfileData/MemProf.cpp
index 2f0e53736c82e..57ec4b26f0510 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -321,6 +321,161 @@ CallStackId hashCallStack(ArrayRef<FrameId> CS) {
   return CSId;
 }
 
+// Returns the sorted list of call stacks.
+std::vector<CallStackRadixTreeBuilder::CSIdPair>
+CallStackRadixTreeBuilder::sortCallStacks(
+    llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
+        &MemProfCallStackData) {
+  // Create a list of call stacks to be sorted.
+  std::vector<CSIdPair> CallStacks;
+  CallStacks.reserve(MemProfCallStackData.size());
+  for (auto &[CSId, CallStack] : MemProfCallStackData) {
+    CallStacks.emplace_back(CSId, &CallStack);
+  }
+
+  // Sort the list of call stacks in the dictionary order to maximize the length
+  // of the common prefix between two adjacent call stacks.
+  auto LessThan = [&](const CSIdPair &L, const CSIdPair &R) {
+    // Call stacks are stored from leaf to root.  Perform comparisons from the
+    // root.
+    return std::lexicographical_compare(
+        L.second->rbegin(), L.second->rend(), R.second->rbegin(),
+        R.second->rend(), [&](FrameId F1, FrameId F2) { return F1 < F2; });
+  };
+  llvm::sort(CallStacks, LessThan);
+
+  return CallStacks;
+}
+
+// Encode a call stack into RadixArray.  Return the starting index within
+// RadixArray.  For each call stack we encode, we emit two or three components
+// into RadixArray.  If a given call stack doesn't have a common prefix relative
+// to the previous one, we emit:
+//
+// - the frames in the given call stack in the reverse order
+//
+// - the length of the given call stack
+//
+// If a given call stack has a non-empty common prefix relative to the previous
+// one, we emit:
+//
+// - the relative location of the common prefix, encoded as a negative number.
+//
+// - a portion of the given call stack that's beyond the common prefix
+//
+// - the length of the given call stack, including the length of the common
+//   prefix.
+//
+// To quickly determine the location of the common prefix within RadixArray,
+// Indexes caches the indexes of the previous call stack's frames within
+// RadixArray.
+uint32_t CallStackRadixTreeBuilder::encodeCallStack(
+    const llvm::SmallVector<FrameId> *CallStack,
+    const llvm::SmallVector<FrameId> *Prev,
+    const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes) {
+  // Compute the length of the common prefix between Prev and CallStack.
+  uint32_t CommonLen = 0;
+  if (Prev) {
+    auto Pos = std::mismatch(Prev->rbegin(), Prev->rend(), CallStack->rbegin(),
+                             CallStack->rend());
+    CommonLen = std::distance(CallStack->rbegin(), Pos.second);
+  }
+
+  // Drop the portion beyond CommonLen.
+  assert(CommonLen <= Indexes.size());
+  Indexes.resize(CommonLen);
+
+  // Append a pointer to the parent.
+  if (CommonLen) {
+    uint32_t CurrentIndex = RadixArray.size();
+    uint32_t ParentIndex = Indexes.back();
+    // The offset to the parent must be negative because we are pointing to an
+    // element we've already added to RadixArray.
+    assert(ParentIndex < CurrentIndex);
+    RadixArray.push_back(ParentIndex - CurrentIndex);
+  }
+
+  // Copy the part of the call stack beyond the common prefix to RadixArray.
+  assert(CommonLen <= CallStack->size());
+  for (FrameId F : llvm::drop_begin(llvm::reverse(*CallStack), CommonLen)) {
+    // Remember the index of F in RadixArray.
+    Indexes.push_back(RadixArray.size());
+    RadixArray.push_back(MemProfFrameIndexes.find(F)->second);
+  }
+  assert(CallStack->size() == Indexes.size());
+
+  // End with the call stack length.
+  RadixArray.push_back(CallStack->size());
+
+  // Return the index within RadixArray where we can start reconstructing a
+  // given call stack from.
+  return RadixArray.size() - 1;
+}
+
+// Build a radix tree array.
+void CallStackRadixTreeBuilder::build(
+    llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
+        &MemProfCallStackData,
+    const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes) {
+  std::vector<CSIdPair> CallStacks = sortCallStacks(MemProfCallStackData);
+  assert(CallStacks.size() == MemProfCallStackData.size());
+
+  // Reserve some reasonable amount of storage.
+  RadixArray.clear();
+  RadixArray.reserve(CallStacks.size() * 8);
+
+  // Indexes will grow as long as the longest call stack.
+  Indexes.clear();
+  Indexes.reserve(512);
+
+  // Compute the radix array.  We encode one call stack at a time, computing the
+  // longest prefix that's shared with the previous call stack we encode.  For
+  // each call stack we encode, we remember a mapping from CallStackId to its
+  // position within RadixArray.
+  //
+  // As an optimization, we encode from the last call stack in CallStacks to
+  // reduce the number of times we follow pointers to the parents.  Consider the
+  // list of call stacks that has been sorted in the dictionary order:
+  //
+  // Call Stack 1: F1
+  // Call Stack 2: F1 -> F2
+  // Call Stack 3: F1 -> F2 -> F3
+  //
+  // If we traversed CallStacks in the forward order, we would end up with a
+  // radix tree like:
+  //
+  // Call Stack 1:  F1
+  //                |
+  // Call Stack 2:  +---> F2
+  //                      |
+  // Call Stack 3:        +---> F3
+  //
+  // Notice that each call stack jumps to the previous one.  However, if we
+  // traverse CallStacks in the reverse order, then Call Stack 3 has the
+  // complete call stack encoded without any pointers.  Call Stack 1 and 2 point
+  // to appropriate prefixes of Call Stack 3.
+  llvm::SmallVector<FrameId> *Prev = nullptr;
+  for (const auto &[CSId, CallStack] : llvm::reverse(CallStacks)) {
+    uint32_t Pos = encodeCallStack(CallStack, Prev, MemProfFrameIndexes);
+    CallStackPos.insert({CSId, Pos});
+    Prev = CallStack;
+  }
+  assert(CallStackPos.size() == MemProfCallStackData.size());
+
+  if (RadixArray.size() >= 2) {
+    // Reverse the radix array in place.  We do so mostly for intuitive
+    // deserialization where we would read the length field and then the call
+    // stack frames proper just like any other array deserialization, except
+    // that we have occasional jumps to take advantage of prefixes.
+    for (size_t I = 0, J = RadixArray.size() - 1; I < J; ++I, --J)
+      std::swap(RadixArray[I], RadixArray[J]);
+
+    // "Reverse" the indexes stored in CallStackPos.
+    for (auto &[K, V] : CallStackPos)
+      V = RadixArray.size() - 1 - V;
+  }
+}
+
 void verifyIndexedMemProfRecord(const IndexedMemProfRecord &Record) {
   for (const auto &AS : Record.AllocSites) {
     assert(AS.CSId == hashCallStack(AS.CallStack));
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index a913718d0fe06..1a160260b6024 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -662,4 +662,111 @@ TEST(MemProf, MissingFrameId) {
   ASSERT_TRUE(FrameIdConv.LastUnmappedId.has_value());
   EXPECT_EQ(*FrameIdConv.LastUnmappedId, 3U);
 }
+
+// Verify CallStackRadixTreeBuilder can handle empty inputs.
+TEST(MemProf, RadixTreeBuilderEmpty) {
+  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes;
+  llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
+  llvm::memprof::CallStackRadixTreeBuilder Builder;
+  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  ASSERT_THAT(Builder.getRadixArray(), testing::IsEmpty());
+  const auto &Mappings = Builder.getCallStackPos();
+  ASSERT_THAT(Mappings, testing::IsEmpty());
+}
+
+// Verify CallStackRadixTreeBuilder can handle one trivial call stack.
+TEST(MemProf, RadixTreeBuilderOne) {
+  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes = {
+      {11, 1}, {12, 2}, {13, 3}};
+  llvm::SmallVector<llvm::memprof::FrameId> CS1 = {13, 12, 11};
+  llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS1), CS1});
+  llvm::memprof::CallStackRadixTreeBuilder Builder;
+  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  EXPECT_THAT(Builder.getRadixArray(), testing::ElementsAreArray({
+                                           3U, // Size of CS1,
+                                           3U, // MemProfFrameIndexes[13]
+                                           2U, // MemProfFrameIndexes[12]
+                                           1U  // MemProfFrameIndexes[11]
+                                       }));
+  const auto &Mappings = Builder.getCallStackPos();
+  ASSERT_THAT(Mappings, SizeIs(1));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS1), 0U)));
+}
+
+// Verify CallStackRadixTreeBuilder can form a link between two call stacks.
+TEST(MemProf, RadixTreeBuilderTwo) {
+  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes = {
+      {11, 1}, {12, 2}, {13, 3}};
+  llvm::SmallVector<llvm::memprof::FrameId> CS1 = {12, 11};
+  llvm::SmallVector<llvm::memprof::FrameId> CS2 = {13, 12, 11};
+  llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS1), CS1});
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS2), CS2});
+  llvm::memprof::CallStackRadixTreeBuilder Builder;
+  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  EXPECT_THAT(Builder.getRadixArray(),
+              testing::ElementsAreArray({
+                  2U,                        // Size of CS1
+                  static_cast<uint32_t>(-3), // Jump 3 steps
+                  3U,                        // Size of CS2
+                  3U,                        // MemProfFrameIndexes[13]
+                  2U,                        // MemProfFrameIndexes[12]
+                  1U                         // MemProfFrameIndexes[11]
+              }));
+  const auto &Mappings = Builder.getCallStackPos();
+  ASSERT_THAT(Mappings, SizeIs(2));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS1), 0U)));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS2), 2U)));
+}
+
+// Verify CallStackRadixTreeBuilder can form a jump to a prefix that itself has
+// another jump to another prefix.
+TEST(MemProf, RadixTreeBuilderSuccessiveJumps) {
+  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes = {
+      {11, 1}, {12, 2}, {13, 3}, {14, 4}, {15, 5}, {16, 6}, {17, 7}, {18, 8},
+  };
+  llvm::SmallVector<llvm::memprof::FrameId> CS1 = {14, 13, 12, 11};
+  llvm::SmallVector<llvm::memprof::FrameId> CS2 = {15, 13, 12, 11};
+  llvm::SmallVector<llvm::memprof::FrameId> CS3 = {17, 16, 12, 11};
+  llvm::SmallVector<llvm::memprof::FrameId> CS4 = {18, 16, 12, 11};
+  llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS1), CS1});
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS2), CS2});
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS3), CS3});
+  MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS4), CS4});
+  llvm::memprof::CallStackRadixTreeBuilder Builder;
+  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  EXPECT_THAT(Builder.getRadixArray(),
+              testing::ElementsAreArray({
+                  4U,                        // Size of CS1
+                  4U,                        // MemProfFrameIndexes[14]
+                  static_cast<uint32_t>(-3), // Jump 3 steps
+                  4U,                        // Size of CS2
+                  5U,                        // MemProfFrameIndexes[15]
+                  3U,                        // MemProfFrameIndexes[13]
+                  static_cast<uint32_t>(-7), // Jump 7 steps
+                  4U,                        // Size of CS3
+                  7U,                        // MemProfFrameIndexes[17]
+                  static_cast<uint32_t>(-3), // Jump 3 steps
+                  4U,                        // Size of CS4
+                  8U,                        // MemProfFrameIndexes[18]
+                  6U,                        // MemProfFrameIndexes[16]
+                  2U,                        // MemProfFrameIndexes[12]
+                  1U                         // MemProfFrameIndexes[11]
+              }));
+  const auto &Mappings = Builder.getCallStackPos();
+  ASSERT_THAT(Mappings, SizeIs(4));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS1), 0U)));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS2), 3U)));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS3), 7U)));
+  EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
+                            llvm::memprof::hashCallStack(CS4), 10U)));
+}
 } // namespace

``````````

</details>


https://github.com/llvm/llvm-project/pull/93784


More information about the llvm-commits mailing list