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

[Kazu Hirata via llvm-commits]

================
@@ -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,
----------------
kazutakahirata wrote:

Sorry, this function declaration doesn't have its corresponding body.  I've deleted this in the latest iteration.

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