xref: /llvm-project/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp (revision b8d2f7177c39af7be371ba7f46cb00b9c63ef8f5)
1 //===- MemoryProfileInfoTest.cpp - Memory Profile Info Unit Tests-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/Analysis/MemoryProfileInfo.h"
10 #include "llvm/AsmParser/Parser.h"
11 #include "llvm/IR/Instructions.h"
12 #include "llvm/IR/LLVMContext.h"
13 #include "llvm/IR/Module.h"
14 #include "llvm/IR/ModuleSummaryIndex.h"
15 #include "llvm/Support/CommandLine.h"
16 #include "llvm/Support/SourceMgr.h"
17 #include "gtest/gtest.h"
18 #include <cstring>
19 #include <sys/types.h>
20 
21 using namespace llvm;
22 using namespace llvm::memprof;
23 
24 extern cl::opt<float> MemProfLifetimeAccessDensityColdThreshold;
25 extern cl::opt<unsigned> MemProfAveLifetimeColdThreshold;
26 extern cl::opt<unsigned> MemProfMinAveLifetimeAccessDensityHotThreshold;
27 
28 namespace {
29 
30 class MemoryProfileInfoTest : public testing::Test {
31 protected:
32   std::unique_ptr<Module> makeLLVMModule(LLVMContext &C, const char *IR) {
33     SMDiagnostic Err;
34     std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
35     if (!Mod)
36       Err.print("MemoryProfileInfoTest", errs());
37     return Mod;
38   }
39 
40   std::unique_ptr<ModuleSummaryIndex> makeLLVMIndex(const char *Summary) {
41     SMDiagnostic Err;
42     std::unique_ptr<ModuleSummaryIndex> Index =
43         parseSummaryIndexAssemblyString(Summary, Err);
44     if (!Index)
45       Err.print("MemoryProfileInfoTest", errs());
46     return Index;
47   }
48 
49   // This looks for a call that has the given value name, which
50   // is the name of the value being assigned the call return value.
51   CallBase *findCall(Function &F, const char *Name = nullptr) {
52     for (auto &BB : F)
53       for (auto &I : BB)
54         if (auto *CB = dyn_cast<CallBase>(&I))
55           if (!Name || CB->getName() == Name)
56             return CB;
57     return nullptr;
58   }
59 };
60 
61 // Test getAllocType helper.
62 // Basic checks on the allocation type for values just above and below
63 // the thresholds.
64 TEST_F(MemoryProfileInfoTest, GetAllocType) {
65   const uint64_t AllocCount = 2;
66   // To be cold we require that
67   // ((float)TotalLifetimeAccessDensity) / AllocCount / 100 <
68   //    MemProfLifetimeAccessDensityColdThreshold
69   // so compute the ColdTotalLifetimeAccessDensityThreshold at the threshold.
70   const uint64_t ColdTotalLifetimeAccessDensityThreshold =
71       (uint64_t)(MemProfLifetimeAccessDensityColdThreshold * AllocCount * 100);
72   // To be cold we require that
73   // ((float)TotalLifetime) / AllocCount >=
74   //    MemProfAveLifetimeColdThreshold * 1000
75   // so compute the TotalLifetime right at the threshold.
76   const uint64_t ColdTotalLifetimeThreshold =
77       MemProfAveLifetimeColdThreshold * AllocCount * 1000;
78   // To be hot we require that
79   // ((float)TotalLifetimeAccessDensity) / AllocCount / 100 >
80   //    MemProfMinAveLifetimeAccessDensityHotThreshold
81   // so compute the HotTotalLifetimeAccessDensityThreshold  at the threshold.
82   const uint64_t HotTotalLifetimeAccessDensityThreshold =
83       (uint64_t)(MemProfMinAveLifetimeAccessDensityHotThreshold * AllocCount * 100);
84 
85 
86   // Test Hot
87   // More accesses per byte per sec than hot threshold is hot.
88   EXPECT_EQ(getAllocType(HotTotalLifetimeAccessDensityThreshold + 1, AllocCount,
89                          ColdTotalLifetimeThreshold + 1),
90             AllocationType::Hot);
91 
92   // Test Cold
93   // Long lived with less accesses per byte per sec than cold threshold is cold.
94   EXPECT_EQ(getAllocType(ColdTotalLifetimeAccessDensityThreshold - 1, AllocCount,
95                          ColdTotalLifetimeThreshold + 1),
96             AllocationType::Cold);
97 
98   // Test NotCold
99   // Long lived with more accesses per byte per sec than cold threshold is not cold.
100   EXPECT_EQ(getAllocType(ColdTotalLifetimeAccessDensityThreshold + 1, AllocCount,
101                          ColdTotalLifetimeThreshold + 1),
102             AllocationType::NotCold);
103   // Short lived with more accesses per byte per sec than cold threshold is not cold.
104   EXPECT_EQ(getAllocType(ColdTotalLifetimeAccessDensityThreshold + 1, AllocCount,
105                          ColdTotalLifetimeThreshold - 1),
106             AllocationType::NotCold);
107   // Short lived with less accesses per byte per sec than cold threshold is not cold.
108   EXPECT_EQ(getAllocType(ColdTotalLifetimeAccessDensityThreshold - 1, AllocCount,
109                          ColdTotalLifetimeThreshold - 1),
110             AllocationType::NotCold);
111 }
112 
113 // Test the hasSingleAllocType helper.
114 TEST_F(MemoryProfileInfoTest, SingleAllocType) {
115   uint8_t NotCold = (uint8_t)AllocationType::NotCold;
116   uint8_t Cold = (uint8_t)AllocationType::Cold;
117   uint8_t Hot = (uint8_t)AllocationType::Hot;
118   EXPECT_TRUE(hasSingleAllocType(NotCold));
119   EXPECT_TRUE(hasSingleAllocType(Cold));
120   EXPECT_TRUE(hasSingleAllocType(Hot));
121   EXPECT_FALSE(hasSingleAllocType(NotCold | Cold));
122   EXPECT_FALSE(hasSingleAllocType(NotCold | Hot));
123   EXPECT_FALSE(hasSingleAllocType(Cold | Hot));
124   EXPECT_FALSE(hasSingleAllocType(NotCold | Cold | Hot));
125 }
126 
127 // Test buildCallstackMetadata helper.
128 TEST_F(MemoryProfileInfoTest, BuildCallStackMD) {
129   LLVMContext C;
130   MDNode *CallStack = buildCallstackMetadata({1, 2, 3}, C);
131   ASSERT_EQ(CallStack->getNumOperands(), 3u);
132   unsigned ExpectedId = 1;
133   for (auto &Op : CallStack->operands()) {
134     auto *StackId = mdconst::dyn_extract<ConstantInt>(Op);
135     EXPECT_EQ(StackId->getZExtValue(), ExpectedId++);
136   }
137 }
138 
139 // Test CallStackTrie::addCallStack interface taking allocation type and list of
140 // call stack ids.
141 // Check that allocations with a single allocation type along all call stacks
142 // get an attribute instead of memprof metadata.
143 TEST_F(MemoryProfileInfoTest, Attribute) {
144   LLVMContext C;
145   std::unique_ptr<Module> M = makeLLVMModule(C,
146                                              R"IR(
147 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
148 target triple = "x86_64-pc-linux-gnu"
149 define i32* @test() {
150 entry:
151   %call1 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
152   %0 = bitcast i8* %call1 to i32*
153   %call2 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
154   %1 = bitcast i8* %call2 to i32*
155   %call3 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
156   %2 = bitcast i8* %call3 to i32*
157   ret i32* %1
158 }
159 declare dso_local noalias noundef i8* @malloc(i64 noundef)
160 )IR");
161 
162   Function *Func = M->getFunction("test");
163 
164   // First call has all cold contexts.
165   CallStackTrie Trie1;
166   Trie1.addCallStack(AllocationType::Cold, {1, 2});
167   Trie1.addCallStack(AllocationType::Cold, {1, 3, 4});
168   CallBase *Call1 = findCall(*Func, "call1");
169   Trie1.buildAndAttachMIBMetadata(Call1);
170 
171   EXPECT_FALSE(Call1->hasMetadata(LLVMContext::MD_memprof));
172   EXPECT_TRUE(Call1->hasFnAttr("memprof"));
173   EXPECT_EQ(Call1->getFnAttr("memprof").getValueAsString(), "cold");
174 
175   // Second call has all non-cold contexts.
176   CallStackTrie Trie2;
177   Trie2.addCallStack(AllocationType::NotCold, {5, 6});
178   Trie2.addCallStack(AllocationType::NotCold, {5, 7, 8});
179   CallBase *Call2 = findCall(*Func, "call2");
180   Trie2.buildAndAttachMIBMetadata(Call2);
181 
182   EXPECT_FALSE(Call2->hasMetadata(LLVMContext::MD_memprof));
183   EXPECT_TRUE(Call2->hasFnAttr("memprof"));
184   EXPECT_EQ(Call2->getFnAttr("memprof").getValueAsString(), "notcold");
185 
186   // Third call has all hot contexts.
187   CallStackTrie Trie3;
188   Trie3.addCallStack(AllocationType::Hot, {9, 10});
189   Trie3.addCallStack(AllocationType::Hot, {9, 11, 12});
190   CallBase *Call3 = findCall(*Func, "call3");
191   Trie3.buildAndAttachMIBMetadata(Call3);
192 
193   EXPECT_FALSE(Call3->hasMetadata(LLVMContext::MD_memprof));
194   EXPECT_TRUE(Call3->hasFnAttr("memprof"));
195   EXPECT_EQ(Call3->getFnAttr("memprof").getValueAsString(), "hot");
196 }
197 
198 // Test CallStackTrie::addCallStack interface taking allocation type and list of
199 // call stack ids.
200 // Test that an allocation call reached by both cold and non cold call stacks
201 // gets memprof metadata representing the different allocation type contexts.
202 TEST_F(MemoryProfileInfoTest, ColdAndNotColdMIB) {
203   LLVMContext C;
204   std::unique_ptr<Module> M = makeLLVMModule(C,
205                                              R"IR(
206 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
207 target triple = "x86_64-pc-linux-gnu"
208 define i32* @test() {
209 entry:
210   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
211   %0 = bitcast i8* %call to i32*
212   ret i32* %0
213 }
214 declare dso_local noalias noundef i8* @malloc(i64 noundef)
215 )IR");
216 
217   Function *Func = M->getFunction("test");
218 
219   CallStackTrie Trie;
220   Trie.addCallStack(AllocationType::Cold, {1, 2});
221   Trie.addCallStack(AllocationType::NotCold, {1, 3});
222 
223   CallBase *Call = findCall(*Func, "call");
224   Trie.buildAndAttachMIBMetadata(Call);
225 
226   EXPECT_FALSE(Call->hasFnAttr("memprof"));
227   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
228   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
229   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
230   for (auto &MIBOp : MemProfMD->operands()) {
231     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
232     MDNode *StackMD = getMIBStackNode(MIB);
233     ASSERT_NE(StackMD, nullptr);
234     ASSERT_EQ(StackMD->getNumOperands(), 2u);
235     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
236     ASSERT_EQ(StackId->getZExtValue(), 1u);
237     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
238     if (StackId->getZExtValue() == 2u)
239       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
240     else {
241       ASSERT_EQ(StackId->getZExtValue(), 3u);
242       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
243     }
244   }
245 }
246 
247 // Test CallStackTrie::addCallStack interface taking allocation type and list of
248 // call stack ids.
249 // Test that an allocation call reached by both cold and hot call stacks
250 // gets memprof metadata representing the different allocation type contexts.
251 TEST_F(MemoryProfileInfoTest, ColdAndHotMIB) {
252   LLVMContext C;
253   std::unique_ptr<Module> M = makeLLVMModule(C,
254                                              R"IR(
255 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
256 target triple = "x86_64-pc-linux-gnu"
257 define i32* @test() {
258 entry:
259   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
260   %0 = bitcast i8* %call to i32*
261   ret i32* %0
262 }
263 declare dso_local noalias noundef i8* @malloc(i64 noundef)
264 )IR");
265 
266   Function *Func = M->getFunction("test");
267 
268   CallStackTrie Trie;
269   Trie.addCallStack(AllocationType::Cold, {1, 2});
270   Trie.addCallStack(AllocationType::Hot, {1, 3});
271 
272   CallBase *Call = findCall(*Func, "call");
273   Trie.buildAndAttachMIBMetadata(Call);
274 
275   EXPECT_FALSE(Call->hasFnAttr("memprof"));
276   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
277   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
278   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
279   for (auto &MIBOp : MemProfMD->operands()) {
280     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
281     MDNode *StackMD = getMIBStackNode(MIB);
282     ASSERT_NE(StackMD, nullptr);
283     ASSERT_EQ(StackMD->getNumOperands(), 2u);
284     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
285     ASSERT_EQ(StackId->getZExtValue(), 1u);
286     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
287     if (StackId->getZExtValue() == 2u)
288       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
289     else {
290       ASSERT_EQ(StackId->getZExtValue(), 3u);
291       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Hot);
292     }
293   }
294 }
295 
296 // Test CallStackTrie::addCallStack interface taking allocation type and list of
297 // call stack ids.
298 // Test that an allocation call reached by both non cold and hot call stacks
299 // gets memprof metadata representing the different allocation type contexts.
300 TEST_F(MemoryProfileInfoTest, NotColdAndHotMIB) {
301   LLVMContext C;
302   std::unique_ptr<Module> M = makeLLVMModule(C,
303                                              R"IR(
304 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
305 target triple = "x86_64-pc-linux-gnu"
306 define i32* @test() {
307 entry:
308   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
309   %0 = bitcast i8* %call to i32*
310   ret i32* %0
311 }
312 declare dso_local noalias noundef i8* @malloc(i64 noundef)
313 )IR");
314 
315   Function *Func = M->getFunction("test");
316 
317   CallStackTrie Trie;
318   Trie.addCallStack(AllocationType::NotCold, {1, 2});
319   Trie.addCallStack(AllocationType::Hot, {1, 3});
320 
321   CallBase *Call = findCall(*Func, "call");
322   Trie.buildAndAttachMIBMetadata(Call);
323 
324   EXPECT_FALSE(Call->hasFnAttr("memprof"));
325   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
326   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
327   ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
328   for (auto &MIBOp : MemProfMD->operands()) {
329     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
330     MDNode *StackMD = getMIBStackNode(MIB);
331     ASSERT_NE(StackMD, nullptr);
332     ASSERT_EQ(StackMD->getNumOperands(), 2u);
333     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
334     ASSERT_EQ(StackId->getZExtValue(), 1u);
335     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
336     if (StackId->getZExtValue() == 2u)
337       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
338     else {
339       ASSERT_EQ(StackId->getZExtValue(), 3u);
340       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Hot);
341     }
342   }
343 }
344 
345 // Test CallStackTrie::addCallStack interface taking allocation type and list of
346 // call stack ids.
347 // Test that an allocation call reached by both cold, non cold and hot call
348 // stacks gets memprof metadata representing the different allocation type
349 // contexts.
350 TEST_F(MemoryProfileInfoTest, ColdAndNotColdAndHotMIB) {
351   LLVMContext C;
352   std::unique_ptr<Module> M = makeLLVMModule(C,
353                                              R"IR(
354 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
355 target triple = "x86_64-pc-linux-gnu"
356 define i32* @test() {
357 entry:
358   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
359   %0 = bitcast i8* %call to i32*
360   ret i32* %0
361 }
362 declare dso_local noalias noundef i8* @malloc(i64 noundef)
363 )IR");
364 
365   Function *Func = M->getFunction("test");
366 
367   CallStackTrie Trie;
368   Trie.addCallStack(AllocationType::Cold, {1, 2});
369   Trie.addCallStack(AllocationType::NotCold, {1, 3});
370   Trie.addCallStack(AllocationType::Hot, {1, 4});
371 
372   CallBase *Call = findCall(*Func, "call");
373   Trie.buildAndAttachMIBMetadata(Call);
374 
375   EXPECT_FALSE(Call->hasFnAttr("memprof"));
376   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
377   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
378   ASSERT_EQ(MemProfMD->getNumOperands(), 3u);
379   for (auto &MIBOp : MemProfMD->operands()) {
380     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
381     MDNode *StackMD = getMIBStackNode(MIB);
382     ASSERT_NE(StackMD, nullptr);
383     ASSERT_EQ(StackMD->getNumOperands(), 2u);
384     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
385     ASSERT_EQ(StackId->getZExtValue(), 1u);
386     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
387     if (StackId->getZExtValue() == 2u) {
388       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
389     } else if (StackId->getZExtValue() == 3u) {
390       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
391     } else {
392       ASSERT_EQ(StackId->getZExtValue(), 4u);
393       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Hot);
394     }
395   }
396 }
397 
398 // Test CallStackTrie::addCallStack interface taking allocation type and list of
399 // call stack ids.
400 // Test that an allocation call reached by multiple call stacks has memprof
401 // metadata with the contexts trimmed to the minimum context required to
402 // identify the allocation type.
403 TEST_F(MemoryProfileInfoTest, TrimmedMIBContext) {
404   LLVMContext C;
405   std::unique_ptr<Module> M = makeLLVMModule(C,
406                                              R"IR(
407 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
408 target triple = "x86_64-pc-linux-gnu"
409 define i32* @test() {
410 entry:
411   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
412   %0 = bitcast i8* %call to i32*
413   ret i32* %0
414 }
415 declare dso_local noalias noundef i8* @malloc(i64 noundef)
416 )IR");
417 
418   Function *Func = M->getFunction("test");
419 
420   CallStackTrie Trie;
421   // We should be able to trim the following two and combine into a single MIB
422   // with the cold context {1, 2}.
423   Trie.addCallStack(AllocationType::Cold, {1, 2, 3});
424   Trie.addCallStack(AllocationType::Cold, {1, 2, 4});
425   // We should be able to trim the following two and combine into a single MIB
426   // with the non-cold context {1, 5}.
427   Trie.addCallStack(AllocationType::NotCold, {1, 5, 6});
428   Trie.addCallStack(AllocationType::NotCold, {1, 5, 7});
429   // We should be able to trim the following two and combine into a single MIB
430   // with the hot context {1, 8}.
431   Trie.addCallStack(AllocationType::Hot, {1, 8, 9});
432   Trie.addCallStack(AllocationType::Hot, {1, 8, 10});
433 
434   CallBase *Call = findCall(*Func, "call");
435   Trie.buildAndAttachMIBMetadata(Call);
436 
437   EXPECT_FALSE(Call->hasFnAttr("memprof"));
438   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
439   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
440   ASSERT_EQ(MemProfMD->getNumOperands(), 3u);
441   for (auto &MIBOp : MemProfMD->operands()) {
442     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
443     MDNode *StackMD = getMIBStackNode(MIB);
444     ASSERT_NE(StackMD, nullptr);
445     ASSERT_EQ(StackMD->getNumOperands(), 2u);
446     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
447     EXPECT_EQ(StackId->getZExtValue(), 1u);
448     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
449     if (StackId->getZExtValue() == 2u)
450       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
451     else if (StackId->getZExtValue() == 5u)
452       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
453     else {
454       ASSERT_EQ(StackId->getZExtValue(), 8u);
455       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Hot);
456     }
457   }
458 }
459 
460 // Test CallStackTrie::addCallStack interface taking memprof MIB metadata.
461 // Check that allocations annotated with memprof metadata with a single
462 // allocation type get simplified to an attribute.
463 TEST_F(MemoryProfileInfoTest, SimplifyMIBToAttribute) {
464   LLVMContext C;
465   std::unique_ptr<Module> M = makeLLVMModule(C,
466                                              R"IR(
467 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
468 target triple = "x86_64-pc-linux-gnu"
469 define i32* @test() {
470 entry:
471   %call1 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !0
472   %0 = bitcast i8* %call1 to i32*
473   %call2 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !3
474   %1 = bitcast i8* %call2 to i32*
475   %call3 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !6
476   %2 = bitcast i8* %call3 to i32*
477   ret i32* %1
478 }
479 declare dso_local noalias noundef i8* @malloc(i64 noundef)
480 !0 = !{!1}
481 !1 = !{!2, !"cold"}
482 !2 = !{i64 1, i64 2, i64 3}
483 !3 = !{!4}
484 !4 = !{!5, !"notcold"}
485 !5 = !{i64 4, i64 5, i64 6, i64 7}
486 !6 = !{!7}
487 !7 = !{!8, !"hot"}
488 !8 = !{i64 8, i64 9, i64 10}
489 )IR");
490 
491   Function *Func = M->getFunction("test");
492 
493   // First call has all cold contexts.
494   CallStackTrie Trie1;
495   CallBase *Call1 = findCall(*Func, "call1");
496   MDNode *MemProfMD1 = Call1->getMetadata(LLVMContext::MD_memprof);
497   ASSERT_EQ(MemProfMD1->getNumOperands(), 1u);
498   MDNode *MIB1 = dyn_cast<MDNode>(MemProfMD1->getOperand(0));
499   Trie1.addCallStack(MIB1);
500   Trie1.buildAndAttachMIBMetadata(Call1);
501 
502   EXPECT_TRUE(Call1->hasFnAttr("memprof"));
503   EXPECT_EQ(Call1->getFnAttr("memprof").getValueAsString(), "cold");
504 
505   // Second call has all non-cold contexts.
506   CallStackTrie Trie2;
507   CallBase *Call2 = findCall(*Func, "call2");
508   MDNode *MemProfMD2 = Call2->getMetadata(LLVMContext::MD_memprof);
509   ASSERT_EQ(MemProfMD2->getNumOperands(), 1u);
510   MDNode *MIB2 = dyn_cast<MDNode>(MemProfMD2->getOperand(0));
511   Trie2.addCallStack(MIB2);
512   Trie2.buildAndAttachMIBMetadata(Call2);
513 
514   EXPECT_TRUE(Call2->hasFnAttr("memprof"));
515   EXPECT_EQ(Call2->getFnAttr("memprof").getValueAsString(), "notcold");
516 
517   // Third call has all hot contexts.
518   CallStackTrie Trie3;
519   CallBase *Call3 = findCall(*Func, "call3");
520   MDNode *MemProfMD3 = Call3->getMetadata(LLVMContext::MD_memprof);
521   ASSERT_EQ(MemProfMD2->getNumOperands(), 1u);
522   MDNode *MIB3 = dyn_cast<MDNode>(MemProfMD3->getOperand(0));
523   Trie3.addCallStack(MIB3);
524   Trie3.buildAndAttachMIBMetadata(Call3);
525 
526   EXPECT_TRUE(Call3->hasFnAttr("memprof"));
527   EXPECT_EQ(Call3->getFnAttr("memprof").getValueAsString(), "hot");
528 }
529 
530 // Test CallStackTrie::addCallStack interface taking memprof MIB metadata.
531 // Test that allocations annotated with memprof metadata with multiple call
532 // stacks gets new memprof metadata with the contexts trimmed to the minimum
533 // context required to identify the allocation type.
534 TEST_F(MemoryProfileInfoTest, ReTrimMIBContext) {
535   LLVMContext C;
536   std::unique_ptr<Module> M = makeLLVMModule(C,
537                                              R"IR(
538 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
539 target triple = "x86_64-pc-linux-gnu"
540 define i32* @test() {
541 entry:
542   %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !0
543   %0 = bitcast i8* %call to i32*
544   ret i32* %0
545 }
546 declare dso_local noalias noundef i8* @malloc(i64 noundef)
547 !0 = !{!1, !3, !5, !7, !9, !11}
548 !1 = !{!2, !"cold"}
549 !2 = !{i64 1, i64 2, i64 3}
550 !3 = !{!4, !"cold"}
551 !4 = !{i64 1, i64 2, i64 4}
552 !5 = !{!6, !"notcold"}
553 !6 = !{i64 1, i64 5, i64 6}
554 !7 = !{!8, !"notcold"}
555 !8 = !{i64 1, i64 5, i64 7}
556 !9 = !{!10, !"hot"}
557 !10 = !{i64 1, i64 8, i64 9}
558 !11 = !{!12, !"hot"}
559 !12 = !{i64 1, i64 8, i64 10}
560 )IR");
561 
562   Function *Func = M->getFunction("test");
563 
564   CallStackTrie Trie;
565   ASSERT_TRUE(Trie.empty());
566   CallBase *Call = findCall(*Func, "call");
567   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
568   for (auto &MIBOp : MemProfMD->operands()) {
569     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
570     Trie.addCallStack(MIB);
571   }
572   ASSERT_FALSE(Trie.empty());
573   Trie.buildAndAttachMIBMetadata(Call);
574 
575   // We should be able to trim the first two and combine into a single MIB
576   // with the cold context {1, 2}.
577   // We should be able to trim the second two and combine into a single MIB
578   // with the non-cold context {1, 5}.
579 
580   EXPECT_FALSE(Call->hasFnAttr("memprof"));
581   EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
582   MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
583   ASSERT_EQ(MemProfMD->getNumOperands(), 3u);
584   for (auto &MIBOp : MemProfMD->operands()) {
585     MDNode *MIB = dyn_cast<MDNode>(MIBOp);
586     MDNode *StackMD = getMIBStackNode(MIB);
587     ASSERT_NE(StackMD, nullptr);
588     ASSERT_EQ(StackMD->getNumOperands(), 2u);
589     auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
590     EXPECT_EQ(StackId->getZExtValue(), 1u);
591     StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
592     if (StackId->getZExtValue() == 2u)
593       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
594     else if (StackId->getZExtValue() == 5u)
595       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
596     else {
597       ASSERT_EQ(StackId->getZExtValue(), 8u);
598       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Hot);
599     }
600   }
601 }
602 
603 TEST_F(MemoryProfileInfoTest, CallStackTestIR) {
604   LLVMContext C;
605   std::unique_ptr<Module> M = makeLLVMModule(C,
606                                              R"IR(
607 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
608 target triple = "x86_64-pc-linux-gnu"
609 define ptr @test() {
610 entry:
611   %call = call noalias noundef nonnull dereferenceable(10) ptr @_Znam(i64 noundef 10), !memprof !1, !callsite !6
612   ret ptr %call
613 }
614 declare noundef nonnull ptr @_Znam(i64 noundef)
615 !1 = !{!2, !4, !7}
616 !2 = !{!3, !"notcold"}
617 !3 = !{i64 1, i64 2, i64 3, i64 4}
618 !4 = !{!5, !"cold"}
619 !5 = !{i64 1, i64 2, i64 3, i64 5}
620 !6 = !{i64 1}
621 !7 = !{!8, !"hot"}
622 !8 = !{i64 1, i64 2, i64 3, i64 6}
623 )IR");
624 
625   Function *Func = M->getFunction("test");
626   CallBase *Call = findCall(*Func, "call");
627 
628   CallStack<MDNode, MDNode::op_iterator> InstCallsite(
629       Call->getMetadata(LLVMContext::MD_callsite));
630 
631   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
632   uint Idx = 0;
633   for (auto &MIBOp : MemProfMD->operands()) {
634     auto *MIBMD = cast<const MDNode>(MIBOp);
635     MDNode *StackNode = getMIBStackNode(MIBMD);
636     CallStack<MDNode, MDNode::op_iterator> StackContext(StackNode);
637     EXPECT_EQ(StackContext.back(), 4 + Idx);
638     std::vector<uint64_t> StackIds;
639     for (auto ContextIter = StackContext.beginAfterSharedPrefix(InstCallsite);
640          ContextIter != StackContext.end(); ++ContextIter)
641       StackIds.push_back(*ContextIter);
642     if (Idx == 0) {
643       std::vector<uint64_t> Expected = {2, 3, 4};
644       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
645     } else if (Idx == 1) {
646       std::vector<uint64_t> Expected = {2, 3, 5};
647       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
648     } else {
649       std::vector<uint64_t> Expected = {2, 3, 6};
650       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
651     }
652     Idx++;
653   }
654 }
655 
656 TEST_F(MemoryProfileInfoTest, CallStackTestSummary) {
657   std::unique_ptr<ModuleSummaryIndex> Index = makeLLVMIndex(R"Summary(
658 ^0 = module: (path: "test.o", hash: (0, 0, 0, 0, 0))
659 ^1 = gv: (guid: 23, summaries: (function: (module: ^0, flags: (linkage: external, visibility: default, notEligibleToImport: 0, live: 0, dsoLocal: 1, canAutoHide: 0), insts: 2, funcFlags: (readNone: 0, readOnly: 0, noRecurse: 0, returnDoesNotAlias: 0, noInline: 1, alwaysInline: 0, noUnwind: 0, mayThrow: 0, hasUnknownCall: 0, mustBeUnreachable: 0), allocs: ((versions: (none), memProf: ((type: notcold, stackIds: (1, 2, 3, 4)), (type: cold, stackIds: (1, 2, 3, 5)), (type: hot, stackIds: (1, 2, 3, 6))))))))
660 ^2 = gv: (guid: 25, summaries: (function: (module: ^0, flags: (linkage: external, visibility: default, notEligibleToImport: 0, live: 0, dsoLocal: 1, canAutoHide: 0), insts: 22, funcFlags: (readNone: 0, readOnly: 0, noRecurse: 1, returnDoesNotAlias: 0, noInline: 1, alwaysInline: 0, noUnwind: 0, mayThrow: 0, hasUnknownCall: 0, mustBeUnreachable: 0), calls: ((callee: ^1)), callsites: ((callee: ^1, clones: (0), stackIds: (3, 4)), (callee: ^1, clones: (0), stackIds: (3, 5)), (callee: ^1, clones: (0), stackIds: (3, 6))))))
661 )Summary");
662 
663   ASSERT_NE(Index, nullptr);
664   auto *CallsiteSummary =
665       cast<FunctionSummary>(Index->getGlobalValueSummary(/*guid=*/25));
666   uint Idx = 0;
667   for (auto &CI : CallsiteSummary->callsites()) {
668     CallStack<CallsiteInfo, SmallVector<unsigned>::const_iterator> InstCallsite(
669         &CI);
670     std::vector<uint64_t> StackIds;
671     for (auto StackIdIndex : InstCallsite)
672       StackIds.push_back(Index->getStackIdAtIndex(StackIdIndex));
673     if (Idx == 0) {
674       std::vector<uint64_t> Expected = {3, 4};
675       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
676     } else if (Idx == 1) {
677       std::vector<uint64_t> Expected = {3, 5};
678       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
679     } else {
680       std::vector<uint64_t> Expected = {3, 6};
681       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
682     }
683     Idx++;
684   }
685 
686   auto *AllocSummary =
687       cast<FunctionSummary>(Index->getGlobalValueSummary(/*guid=*/23));
688   for (auto &AI : AllocSummary->allocs()) {
689     uint Idx = 0;
690     for (auto &MIB : AI.MIBs) {
691       CallStack<MIBInfo, SmallVector<unsigned>::const_iterator> StackContext(
692           &MIB);
693       EXPECT_EQ(Index->getStackIdAtIndex(StackContext.back()), 4 + Idx);
694       std::vector<uint64_t> StackIds;
695       for (auto StackIdIndex : StackContext)
696         StackIds.push_back(Index->getStackIdAtIndex(StackIdIndex));
697       if (Idx == 0) {
698         std::vector<uint64_t> Expected = {1, 2, 3, 4};
699         EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
700       } else if (Idx == 1) {
701         std::vector<uint64_t> Expected = {1, 2, 3, 5};
702         EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
703       } else {
704         std::vector<uint64_t> Expected = {1, 2, 3, 6};
705         EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
706       }
707       Idx++;
708     }
709   }
710 }
711 } // end anonymous namespace
712