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

Kazu Hirata via llvm-commits llvm-commits at lists.llvm.org
Thu Jun 6 15:52:03 PDT 2024


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

>From ebd1d7047c55e238f6bc48837c14b1c84cdbe9cf Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Wed, 29 May 2024 14:19:37 -0700
Subject: [PATCH 1/5] [memprof] Add CallStackRadixTreeBuilder

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.
---
 llvm/include/llvm/ProfileData/MemProf.h    | 100 +++++++++++++
 llvm/lib/ProfileData/MemProf.cpp           | 155 +++++++++++++++++++++
 llvm/unittests/ProfileData/MemProfTest.cpp | 107 ++++++++++++++
 3 files changed, 362 insertions(+)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 2a7c853cade66..4015d2ed47dad 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -905,6 +905,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 1d9860e0ea7e8..ec6f2f2ef48b0 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -410,6 +410,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

>From ca0ea2eb09dd5d3a286e8078ac98253fe071987f Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Tue, 4 Jun 2024 01:22:07 -0700
Subject: [PATCH 2/5] Expand comments.

Call MemProfCallStackData.takeVector() to take advantage of the move
semantics.
---
 llvm/include/llvm/ProfileData/MemProf.h    | 40 ++++++++------
 llvm/lib/ProfileData/MemProf.cpp           | 63 ++++++++++------------
 llvm/unittests/ProfileData/MemProfTest.cpp |  8 +--
 3 files changed, 54 insertions(+), 57 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 4015d2ed47dad..4a07f14f935fb 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -947,6 +947,10 @@ struct IndexedMemProfData {
 //
 // - JN indicates a pointer to the parent, encoded as -N.
 //
+// The radix tree allows us to reconstruct call stacks in the leaf-to-root
+// order as we scan the array from left ro right while following pointers to
+// parents along the way
+//
 // 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,
@@ -963,20 +967,27 @@ class CallStackRadixTreeBuilder {
   // 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:
+  // In build, we partition a given call stack into two parts -- the prefix
+  // that's common with the previously encoded call stack and the frames beyond
+  // the common prefix -- the unique portion.  Then we want to find out where
+  // the common prefix is stored in RadixArray so that we can link the unique
+  // portion to the common prefix.  Indexes, declared below, helps with our
+  // needs.  Intuitively, Indexes tells us where each of the previously encoded
+  // call stack is stored in RadixArray.  More formally, Indexes satisfies:
+  //
+  //   RadixArray[Indexes[I]] == Prev[I]
   //
-  //   RadixArray[Indexes[I]] == (*Prev)[Prev->size() - I - 1]
+  // for every I, where Prev is the the call stack in the root-to-leaf order
+  // previously encoded by build.  (Note that Prev, as passed to
+  // encodeCallStack, is in the leaf-to-root order.)
   //
-  // where Prev is one of the parameters to build.
+  // For example, if the call stack being encoded shares 5 frames at the root of
+  // the call stack with the previously encoded call stack,
+  // RadixArray[Indexes[0]] is the root frame of the common prefix.
+  // RadixArray[Indexes[5 - 1]] is the last frame of the common prefix.
   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);
+  using CSIdPair = std::pair<CallStackId, llvm::SmallVector<FrameId>>;
 
   // Encode a call stack into RadixArray.  Return the starting index within
   // RadixArray.
@@ -988,8 +999,9 @@ class CallStackRadixTreeBuilder {
 public:
   CallStackRadixTreeBuilder() = default;
 
+  // Build a radix tree array.
   void build(llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
-                 &MemProfCallStackData,
+                 &&MemProfCallStackData,
              const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes);
 
   const std::vector<uint32_t> &getRadixArray() const { return RadixArray; }
@@ -999,12 +1011,6 @@ class CallStackRadixTreeBuilder {
   }
 };
 
-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 ec6f2f2ef48b0..327f4216ef5ca 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -410,38 +410,12 @@ 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 frames in the given call stack in the root-to-leaf order
 //
 // - the length of the given call stack
 //
@@ -455,6 +429,14 @@ CallStackRadixTreeBuilder::sortCallStacks(
 // - the length of the given call stack, including the length of the common
 //   prefix.
 //
+// The resulting RadixArray requires a somewhat unintuitive backward traversal
+// to reconstruct a call stack -- read the call stack length and scan backward
+// while collecting frames in the leaf to root order.  build, the caller of this
+// function, reverses RadixArray in place so that we can reconstruct a call
+// stack as if we were deserializing an array in a typical way -- the call stack
+// length followed by the frames in the leaf-to-root order except that we need
+// to handle pointers to parents along the way.
+//
 // To quickly determine the location of the common prefix within RadixArray,
 // Indexes caches the indexes of the previous call stack's frames within
 // RadixArray.
@@ -462,7 +444,7 @@ 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.
+  // Compute the length of the common root prefix between Prev and CallStack.
   uint32_t CommonLen = 0;
   if (Prev) {
     auto Pos = std::mismatch(Prev->rbegin(), Prev->rend(), CallStack->rbegin(),
@@ -501,13 +483,23 @@ uint32_t CallStackRadixTreeBuilder::encodeCallStack(
   return RadixArray.size() - 1;
 }
 
-// Build a radix tree array.
 void CallStackRadixTreeBuilder::build(
     llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
-        &MemProfCallStackData,
+        &&MemProfCallStackData,
     const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes) {
-  std::vector<CSIdPair> CallStacks = sortCallStacks(MemProfCallStackData);
-  assert(CallStacks.size() == MemProfCallStackData.size());
+  // Take the vector portion of MemProfCallStackData.  The vector is exactly
+  // what we need to sort.  Also, we no longer need its lookup capability.
+  llvm::SmallVector<CSIdPair, 0> CallStacks = MemProfCallStackData.takeVector();
+
+  // Sort the list of call stacks in the dictionary order to maximize the length
+  // of the common prefix between two adjacent call stacks.
+  llvm::sort(CallStacks, [&](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; });
+  });
 
   // Reserve some reasonable amount of storage.
   RadixArray.clear();
@@ -543,13 +535,12 @@ void CallStackRadixTreeBuilder::build(
   // 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;
+  const llvm::SmallVector<FrameId> *Prev = nullptr;
   for (const auto &[CSId, CallStack] : llvm::reverse(CallStacks)) {
-    uint32_t Pos = encodeCallStack(CallStack, Prev, MemProfFrameIndexes);
+    uint32_t Pos = encodeCallStack(&CallStack, Prev, MemProfFrameIndexes);
     CallStackPos.insert({CSId, Pos});
-    Prev = CallStack;
+    Prev = &CallStack;
   }
-  assert(CallStackPos.size() == MemProfCallStackData.size());
 
   if (RadixArray.size() >= 2) {
     // Reverse the radix array in place.  We do so mostly for intuitive
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 1a160260b6024..408fb548de9bb 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -668,7 +668,7 @@ TEST(MemProf, RadixTreeBuilderEmpty) {
   llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes;
   llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
   llvm::memprof::CallStackRadixTreeBuilder Builder;
-  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
   ASSERT_THAT(Builder.getRadixArray(), testing::IsEmpty());
   const auto &Mappings = Builder.getCallStackPos();
   ASSERT_THAT(Mappings, testing::IsEmpty());
@@ -682,7 +682,7 @@ TEST(MemProf, RadixTreeBuilderOne) {
   llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
   MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS1), CS1});
   llvm::memprof::CallStackRadixTreeBuilder Builder;
-  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
   EXPECT_THAT(Builder.getRadixArray(), testing::ElementsAreArray({
                                            3U, // Size of CS1,
                                            3U, // MemProfFrameIndexes[13]
@@ -705,7 +705,7 @@ TEST(MemProf, RadixTreeBuilderTwo) {
   MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS1), CS1});
   MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS2), CS2});
   llvm::memprof::CallStackRadixTreeBuilder Builder;
-  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
   EXPECT_THAT(Builder.getRadixArray(),
               testing::ElementsAreArray({
                   2U,                        // Size of CS1
@@ -739,7 +739,7 @@ TEST(MemProf, RadixTreeBuilderSuccessiveJumps) {
   MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS3), CS3});
   MemProfCallStackData.insert({llvm::memprof::hashCallStack(CS4), CS4});
   llvm::memprof::CallStackRadixTreeBuilder Builder;
-  Builder.build(MemProfCallStackData, MemProfFrameIndexes);
+  Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
   EXPECT_THAT(Builder.getRadixArray(),
               testing::ElementsAreArray({
                   4U,                        // Size of CS1

>From 33a2e6ccee50d5925282195b95f5730148c73826 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Tue, 4 Jun 2024 12:26:37 -0700
Subject: [PATCH 3/5] Use LinearFrameId and LinearCallStackId.

---
 llvm/include/llvm/ProfileData/MemProf.h | 10 +++++-----
 llvm/lib/ProfileData/MemProf.cpp        |  9 +++++----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 29e26cefd13f8..1ae57656deda2 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -1006,10 +1006,10 @@ class CallStackRadixTreeBuilder {
 
   // 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);
+  LinearCallStackId encodeCallStack(
+      const llvm::SmallVector<FrameId> *CallStack,
+      const llvm::SmallVector<FrameId> *Prev,
+      const llvm::DenseMap<FrameId, LinearFrameId> &MemProfFrameIndexes);
 
 public:
   CallStackRadixTreeBuilder() = default;
@@ -1017,7 +1017,7 @@ class CallStackRadixTreeBuilder {
   // Build a radix tree array.
   void build(llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
                  &&MemProfCallStackData,
-             const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes);
+             const llvm::DenseMap<FrameId, LinearFrameId> &MemProfFrameIndexes);
 
   const std::vector<uint32_t> &getRadixArray() const { return RadixArray; }
 
diff --git a/llvm/lib/ProfileData/MemProf.cpp b/llvm/lib/ProfileData/MemProf.cpp
index 327f4216ef5ca..6f4725d2590fc 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -440,10 +440,10 @@ CallStackId hashCallStack(ArrayRef<FrameId> CS) {
 // 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(
+LinearCallStackId CallStackRadixTreeBuilder::encodeCallStack(
     const llvm::SmallVector<FrameId> *CallStack,
     const llvm::SmallVector<FrameId> *Prev,
-    const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes) {
+    const llvm::DenseMap<FrameId, LinearFrameId> &MemProfFrameIndexes) {
   // Compute the length of the common root prefix between Prev and CallStack.
   uint32_t CommonLen = 0;
   if (Prev) {
@@ -486,7 +486,7 @@ uint32_t CallStackRadixTreeBuilder::encodeCallStack(
 void CallStackRadixTreeBuilder::build(
     llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>>
         &&MemProfCallStackData,
-    const llvm::DenseMap<FrameId, uint32_t> &MemProfFrameIndexes) {
+    const llvm::DenseMap<FrameId, LinearFrameId> &MemProfFrameIndexes) {
   // Take the vector portion of MemProfCallStackData.  The vector is exactly
   // what we need to sort.  Also, we no longer need its lookup capability.
   llvm::SmallVector<CSIdPair, 0> CallStacks = MemProfCallStackData.takeVector();
@@ -537,7 +537,8 @@ void CallStackRadixTreeBuilder::build(
   // to appropriate prefixes of Call Stack 3.
   const llvm::SmallVector<FrameId> *Prev = nullptr;
   for (const auto &[CSId, CallStack] : llvm::reverse(CallStacks)) {
-    uint32_t Pos = encodeCallStack(&CallStack, Prev, MemProfFrameIndexes);
+    LinearCallStackId Pos =
+        encodeCallStack(&CallStack, Prev, MemProfFrameIndexes);
     CallStackPos.insert({CSId, Pos});
     Prev = &CallStack;
   }

>From 1d2ddf6ca044ffd5188ad0dc79c11661c1d89920 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Tue, 4 Jun 2024 12:33:13 -0700
Subject: [PATCH 4/5] Use LinearFrameId and LinearCallStackId in more places.

---
 llvm/include/llvm/ProfileData/MemProf.h    | 11 ++++++-----
 llvm/unittests/ProfileData/MemProfTest.cpp |  8 ++++----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 1ae57656deda2..532a709c5690f 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -977,10 +977,10 @@ struct IndexedMemProfData {
 // "CallStackId 1 -> 11".
 class CallStackRadixTreeBuilder {
   // The radix tree array.
-  std::vector<uint32_t> RadixArray;
+  std::vector<LinearFrameId> RadixArray;
 
   // Mapping from CallStackIds to indexes into RadixArray.
-  llvm::DenseMap<CallStackId, uint32_t> CallStackPos;
+  llvm::DenseMap<CallStackId, LinearCallStackId> CallStackPos;
 
   // In build, we partition a given call stack into two parts -- the prefix
   // that's common with the previously encoded call stack and the frames beyond
@@ -1000,7 +1000,7 @@ class CallStackRadixTreeBuilder {
   // the call stack with the previously encoded call stack,
   // RadixArray[Indexes[0]] is the root frame of the common prefix.
   // RadixArray[Indexes[5 - 1]] is the last frame of the common prefix.
-  std::vector<uint32_t> Indexes;
+  std::vector<LinearCallStackId> Indexes;
 
   using CSIdPair = std::pair<CallStackId, llvm::SmallVector<FrameId>>;
 
@@ -1019,9 +1019,10 @@ class CallStackRadixTreeBuilder {
                  &&MemProfCallStackData,
              const llvm::DenseMap<FrameId, LinearFrameId> &MemProfFrameIndexes);
 
-  const std::vector<uint32_t> &getRadixArray() const { return RadixArray; }
+  const std::vector<LinearFrameId> &getRadixArray() const { return RadixArray; }
 
-  const llvm::DenseMap<CallStackId, uint32_t> &getCallStackPos() const {
+  const llvm::DenseMap<CallStackId, LinearCallStackId> &
+  getCallStackPos() const {
     return CallStackPos;
   }
 };
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 408fb548de9bb..e120bc3e35947 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -665,7 +665,7 @@ TEST(MemProf, MissingFrameId) {
 
 // Verify CallStackRadixTreeBuilder can handle empty inputs.
 TEST(MemProf, RadixTreeBuilderEmpty) {
-  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes;
+  llvm::DenseMap<FrameId, llvm::memprof::LinearFrameId> MemProfFrameIndexes;
   llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
   llvm::memprof::CallStackRadixTreeBuilder Builder;
   Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
@@ -676,7 +676,7 @@ TEST(MemProf, RadixTreeBuilderEmpty) {
 
 // Verify CallStackRadixTreeBuilder can handle one trivial call stack.
 TEST(MemProf, RadixTreeBuilderOne) {
-  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes = {
+  llvm::DenseMap<FrameId, llvm::memprof::LinearFrameId> MemProfFrameIndexes = {
       {11, 1}, {12, 2}, {13, 3}};
   llvm::SmallVector<llvm::memprof::FrameId> CS1 = {13, 12, 11};
   llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> MemProfCallStackData;
@@ -697,7 +697,7 @@ TEST(MemProf, RadixTreeBuilderOne) {
 
 // Verify CallStackRadixTreeBuilder can form a link between two call stacks.
 TEST(MemProf, RadixTreeBuilderTwo) {
-  llvm::DenseMap<FrameId, uint32_t> MemProfFrameIndexes = {
+  llvm::DenseMap<FrameId, llvm::memprof::LinearFrameId> MemProfFrameIndexes = {
       {11, 1}, {12, 2}, {13, 3}};
   llvm::SmallVector<llvm::memprof::FrameId> CS1 = {12, 11};
   llvm::SmallVector<llvm::memprof::FrameId> CS2 = {13, 12, 11};
@@ -726,7 +726,7 @@ TEST(MemProf, RadixTreeBuilderTwo) {
 // 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 = {
+  llvm::DenseMap<FrameId, llvm::memprof::LinearFrameId> 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};

>From 5a9c8ab47ac6d96522ef6579580b0289267a47f0 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Thu, 6 Jun 2024 15:51:19 -0700
Subject: [PATCH 5/5] Teach CallStackRadixTreeBuilder::build to return early.

---
 llvm/lib/ProfileData/MemProf.cpp | 36 +++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/ProfileData/MemProf.cpp b/llvm/lib/ProfileData/MemProf.cpp
index 6f4725d2590fc..c0b9faafacaf4 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/llvm/lib/ProfileData/MemProf.cpp
@@ -491,6 +491,13 @@ void CallStackRadixTreeBuilder::build(
   // what we need to sort.  Also, we no longer need its lookup capability.
   llvm::SmallVector<CSIdPair, 0> CallStacks = MemProfCallStackData.takeVector();
 
+  // Return early if we have no work to do.
+  if (CallStacks.empty()) {
+    RadixArray.clear();
+    CallStackPos.clear();
+    return;
+  }
+
   // Sort the list of call stacks in the dictionary order to maximize the length
   // of the common prefix between two adjacent call stacks.
   llvm::sort(CallStacks, [&](const CSIdPair &L, const CSIdPair &R) {
@@ -509,6 +516,10 @@ void CallStackRadixTreeBuilder::build(
   Indexes.clear();
   Indexes.reserve(512);
 
+  // CallStackPos will grow to exactly CallStacks.size() entries.
+  CallStackPos.clear();
+  CallStackPos.reserve(CallStacks.size());
+
   // 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
@@ -543,18 +554,19 @@ void CallStackRadixTreeBuilder::build(
     Prev = &CallStack;
   }
 
-  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;
-  }
+  // "RadixArray.size() - 1" below is problematic if RadixArray is empty.
+  assert(!RadixArray.empty());
+
+  // 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) {



More information about the llvm-commits mailing list