[llvm] [memprof] Add another constructor to MemProfReader (PR #88952)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata created https://github.com/llvm/llvm-project/pull/88952

This patch enables users of MemProfReader to directly supply mappings
from CallStackId to actual call stacks.

Once the users of the current constructor without CSIdMap switch to
the new constructor, we'll have fewer users of:

- IndexedAllocationInfo::CallStack
- IndexedMemProfRecord::CallSites

bringing us one step closer to the removal of these fields in favor
of:

- IndexedAllocationInfo::CSId
- IndexedMemProfRecord::CallSiteIds


>From 2f91dcedf36af39a20fbf873f35961e9882a135a Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Tue, 16 Apr 2024 10:54:17 -0700
Subject: [PATCH] [memprof] Add another constructor to MemProfReader

This patch enables users of MemProfReader to directly supply mappings
from CallStackId to actual call stacks.

Once the users of the current constructor without CSIdMap switch to
the new constructor, we'll have fewer users of:

- IndexedAllocationInfo::CallStack
- IndexedMemProfRecord::CallSites

bringing us one step closer to the removal of these fields in favor
of:

- IndexedAllocationInfo::CSId
- IndexedMemProfRecord::CallSiteIds
---
 llvm/include/llvm/ProfileData/MemProfReader.h |  9 ++++
 llvm/unittests/ProfileData/MemProfTest.cpp    | 41 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index 1f84fefad03e39..7fa8af184dc93b 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -98,6 +98,15 @@ class MemProfReader {
       llvm::DenseMap<FrameId, Frame> FrameIdMap,
       llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> ProfData);
 
+  // Initialize the MemProfReader with the frame mappings, call stack mappings,
+  // and profile contents.
+  MemProfReader(
+      llvm::DenseMap<FrameId, Frame> FrameIdMap,
+      llvm::DenseMap<CallStackId, llvm::SmallVector<FrameId>> CSIdMap,
+      llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> ProfData)
+      : IdToFrame(std::move(FrameIdMap)), CSIdToCallStack(std::move(CSIdMap)),
+        FunctionProfileData(std::move(ProfData)) {}
+
 protected:
   // A helper method to extract the frame from the IdToFrame map.
   const Frame &idToFrame(const FrameId Id) const {
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index ab9227e9df881b..f596919ed039a8 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -436,6 +436,47 @@ TEST(MemProf, BaseMemProfReader) {
               FrameContains("bar", 10U, 2U, false));
 }
 
+TEST(MemProf, BaseMemProfReaderWithCSIdMap) {
+  llvm::DenseMap<FrameId, Frame> FrameIdMap;
+  Frame F1(/*Hash=*/IndexedMemProfRecord::getGUID("foo"), /*LineOffset=*/20,
+           /*Column=*/5, /*IsInlineFrame=*/true);
+  Frame F2(/*Hash=*/IndexedMemProfRecord::getGUID("bar"), /*LineOffset=*/10,
+           /*Column=*/2, /*IsInlineFrame=*/false);
+  FrameIdMap.insert({F1.hash(), F1});
+  FrameIdMap.insert({F2.hash(), F2});
+
+  llvm::DenseMap<CallStackId, llvm::SmallVector<FrameId>> CSIdMap;
+  llvm::SmallVector<FrameId> CallStack = {F1.hash(), F2.hash()};
+  CallStackId CSId = llvm::memprof::hashCallStack(CallStack);
+  CSIdMap.insert({CSId, CallStack});
+
+  llvm::MapVector<llvm::GlobalValue::GUID, IndexedMemProfRecord> ProfData;
+  IndexedMemProfRecord FakeRecord;
+  MemInfoBlock Block;
+  Block.AllocCount = 1U, Block.TotalAccessDensity = 4,
+  Block.TotalLifetime = 200001;
+  FakeRecord.AllocSites.emplace_back(
+      /*CS=*/llvm::SmallVector<FrameId>(),
+      /*CSId=*/llvm::memprof::hashCallStack(CallStack),
+      /*MB=*/Block);
+  ProfData.insert({F1.hash(), FakeRecord});
+
+  MemProfReader Reader(FrameIdMap, CSIdMap, ProfData);
+
+  llvm::SmallVector<MemProfRecord, 1> Records;
+  for (const auto &KeyRecordPair : Reader) {
+    Records.push_back(KeyRecordPair.second);
+  }
+
+  ASSERT_THAT(Records, SizeIs(1));
+  ASSERT_THAT(Records[0].AllocSites, SizeIs(1));
+  ASSERT_THAT(Records[0].AllocSites[0].CallStack, SizeIs(2));
+  EXPECT_THAT(Records[0].AllocSites[0].CallStack[0],
+              FrameContains("foo", 20U, 5U, true));
+  EXPECT_THAT(Records[0].AllocSites[0].CallStack[1],
+              FrameContains("bar", 10U, 2U, false));
+}
+
 TEST(MemProf, IndexedMemProfRecordToMemProfRecord) {
   // Verify that MemProfRecord can be constructed from IndexedMemProfRecord with
   // CallStackIds only.