[llvm] c348e26 - [memprof] Use CallStackRadixTreeBuilder in the V3 format (#94708)

[via llvm-commits]

Author: Kazu Hirata
Date: 2024-06-07T07:19:36-07:00
New Revision: c348e265bd1284f770e66639633199fefd8015ec

URL: https://github.com/llvm/llvm-project/commit/c348e265bd1284f770e66639633199fefd8015ec
DIFF: https://github.com/llvm/llvm-project/commit/c348e265bd1284f770e66639633199fefd8015ec.diff

LOG: [memprof] Use CallStackRadixTreeBuilder in the V3 format (#94708)

This patch integrates CallStackRadixTreeBuilder into the V3 format,
reducing the profile size to about 27% of the V2 profile size.

- Serialization: writeMemProfCallStackArray just needs to write out
  the radix tree array prepared by CallStackRadixTreeBuilder.
  Mappings from CallStackIds to LinearCallStackIds are moved by new
  function CallStackRadixTreeBuilder::takeCallStackPos.

- Deserialization: Deserializing a call stack is the same as
  deserializing an array encoded in the obvious manner -- the length
  followed by the payload, except that we need to follow a pointer to
  the parent to take advantage of common prefixes once in a while.
  This patch teaches LinearCallStackIdConverter to how to handle those
  pointers.

Added: 
    

Modified: 
    llvm/include/llvm/ProfileData/MemProf.h
    llvm/lib/ProfileData/InstrProfWriter.cpp
    llvm/unittests/ProfileData/MemProfTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 6b0fa6cd6541c..0e6245acb77e8 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -900,9 +900,18 @@ struct LinearCallStackIdConverter {
     Frames.reserve(NumFrames);
     for (; NumFrames; --NumFrames) {
       LinearFrameId Elem =
-          support::endian::readNext<LinearFrameId, llvm::endianness::little>(
-              Ptr);
+          support::endian::read<LinearFrameId, llvm::endianness::little>(Ptr);
+      // Follow a pointer to the parent, if any.  See comments below on
+      // CallStackRadixTreeBuilder for the description of the radix tree format.
+      if (static_cast<std::make_signed_t<LinearFrameId>>(Elem) < 0) {
+        Ptr += (-Elem) * sizeof(LinearFrameId);
+        Elem =
+            support::endian::read<LinearFrameId, llvm::endianness::little>(Ptr);
+      }
+      // We shouldn't encounter another pointer.
+      assert(static_cast<std::make_signed_t<LinearFrameId>>(Elem) >= 0);
       Frames.push_back(FrameIdToFrame(Elem));
+      Ptr += sizeof(LinearFrameId);
     }
 
     return Frames;
@@ -1024,9 +1033,8 @@ class CallStackRadixTreeBuilder {
 
   const std::vector<LinearFrameId> &getRadixArray() const { return RadixArray; }
 
-  const llvm::DenseMap<CallStackId, LinearCallStackId> &
-  getCallStackPos() const {
-    return CallStackPos;
+  llvm::DenseMap<CallStackId, LinearCallStackId> takeCallStackPos() {
+    return std::move(CallStackPos);
   }
 };
 

diff  --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index e58e6b8acfc81..a73f72a534f16 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -547,19 +547,11 @@ writeMemProfCallStackArray(
   llvm::DenseMap<memprof::CallStackId, memprof::LinearCallStackId>
       MemProfCallStackIndexes;
 
-  MemProfCallStackIndexes.reserve(MemProfCallStackData.size());
-  uint64_t CallStackBase = OS.tell();
-  for (const auto &[CSId, CallStack] : MemProfCallStackData) {
-    memprof::LinearCallStackId CallStackIndex =
-        (OS.tell() - CallStackBase) / sizeof(memprof::LinearCallStackId);
-    MemProfCallStackIndexes.insert({CSId, CallStackIndex});
-    const llvm::SmallVector<memprof::FrameId> CS = CallStack;
-    OS.write32(CS.size());
-    for (const auto F : CS) {
-      assert(MemProfFrameIndexes.contains(F));
-      OS.write32(MemProfFrameIndexes[F]);
-    }
-  }
+  memprof::CallStackRadixTreeBuilder Builder;
+  Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
+  for (auto I : Builder.getRadixArray())
+    OS.write32(I);
+  MemProfCallStackIndexes = Builder.takeCallStackPos();
 
   // Release the memory of this vector as it is no longer needed.
   MemProfCallStackData.clear();
@@ -695,8 +687,8 @@ static Error writeMemProfV2(ProfOStream &OS,
 // uint64_t Schema entry 1
 // ....
 // uint64_t Schema entry N - 1
-// OnDiskChainedHashTable MemProfFrameData
-// OnDiskChainedHashTable MemProfCallStackData
+// Frames serialized one after another
+// Call stacks encoded as a radix tree
 // OnDiskChainedHashTable MemProfRecordData
 static Error writeMemProfV3(ProfOStream &OS,
                             memprof::IndexedMemProfData &MemProfData,

diff  --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index e120bc3e35947..26421200e1a11 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -670,7 +670,7 @@ TEST(MemProf, RadixTreeBuilderEmpty) {
   llvm::memprof::CallStackRadixTreeBuilder Builder;
   Builder.build(std::move(MemProfCallStackData), MemProfFrameIndexes);
   ASSERT_THAT(Builder.getRadixArray(), testing::IsEmpty());
-  const auto &Mappings = Builder.getCallStackPos();
+  const auto Mappings = Builder.takeCallStackPos();
   ASSERT_THAT(Mappings, testing::IsEmpty());
 }
 
@@ -689,7 +689,7 @@ TEST(MemProf, RadixTreeBuilderOne) {
                                            2U, // MemProfFrameIndexes[12]
                                            1U  // MemProfFrameIndexes[11]
                                        }));
-  const auto &Mappings = Builder.getCallStackPos();
+  const auto Mappings = Builder.takeCallStackPos();
   ASSERT_THAT(Mappings, SizeIs(1));
   EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
                             llvm::memprof::hashCallStack(CS1), 0U)));
@@ -715,7 +715,7 @@ TEST(MemProf, RadixTreeBuilderTwo) {
                   2U,                        // MemProfFrameIndexes[12]
                   1U                         // MemProfFrameIndexes[11]
               }));
-  const auto &Mappings = Builder.getCallStackPos();
+  const auto Mappings = Builder.takeCallStackPos();
   ASSERT_THAT(Mappings, SizeIs(2));
   EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
                             llvm::memprof::hashCallStack(CS1), 0U)));
@@ -758,7 +758,7 @@ TEST(MemProf, RadixTreeBuilderSuccessiveJumps) {
                   2U,                        // MemProfFrameIndexes[12]
                   1U                         // MemProfFrameIndexes[11]
               }));
-  const auto &Mappings = Builder.getCallStackPos();
+  const auto Mappings = Builder.takeCallStackPos();
   ASSERT_THAT(Mappings, SizeIs(4));
   EXPECT_THAT(Mappings, testing::Contains(testing::Pair(
                             llvm::memprof::hashCallStack(CS1), 0U)));