llvm/unittests/Transforms/Instrumentation/MemProfUseTest.cpp

   1 //===- MemProfUseTest.cpp - MemProf use 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/TargetLibraryInfo.h"
  10 #include "llvm/AsmParser/Parser.h"
  11 #include "llvm/IR/LLVMContext.h"
  12 #include "llvm/IR/Module.h"
  13 #include "llvm/Passes/PassBuilder.h"
  14 #include "llvm/ProfileData/InstrProfReader.h"
  15 #include "llvm/ProfileData/InstrProfWriter.h"
  16 #include "llvm/ProfileData/MemProf.h"
  17 #include "llvm/Support/SourceMgr.h"
  18 #include "llvm/Testing/Support/Error.h"
  19 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
  20
  21 #include "gmock/gmock.h"
  22 #include "gtest/gtest.h"
  23
  24 namespace {
  25 using namespace llvm;
  26 using namespace llvm::memprof;
  27 using testing::Contains;
  28 using testing::ElementsAre;
  29 using testing::FieldsAre;
  30 using testing::Pair;
  31 using testing::SizeIs;
  32 using testing::UnorderedElementsAre;
  33
  34 TEST(MemProf, ExtractDirectCallsFromIR) {
  35   // The following IR is generated from:
  36   //
  37   // void f1();
  38   // void f2();
  39   // void f3();
  40   //
  41   // void foo() {
  42   //   f1();
  43   //   f2(); f3();
  44   // }
  45   StringRef IR = R"IR(
  46 define dso_local void @_Z3foov() !dbg !10 {
  47 entry:
  48   call void @_Z2f1v(), !dbg !13
  49   call void @_Z2f2v(), !dbg !14
  50   call void @_Z2f3v(), !dbg !15
  51   ret void, !dbg !16
  52 }
  53
  54 declare !dbg !17 void @_Z2f1v()
  55
  56 declare !dbg !18 void @_Z2f2v()
  57
  58 declare !dbg !19 void @_Z2f3v()
  59
  60 !llvm.dbg.cu = !{!0}
  61 !llvm.module.flags = !{!2, !3, !4, !5, !6, !7, !8}
  62 !llvm.ident = !{!9}
  63
  64 !0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, debugInfoForProfiling: true, nameTableKind: None)
  65 !1 = !DIFile(filename: "foobar.cc", directory: "/")
  66 !2 = !{i32 7, !"Dwarf Version", i32 5}
  67 !3 = !{i32 2, !"Debug Info Version", i32 3}
  68 !4 = !{i32 1, !"wchar_size", i32 4}
  69 !5 = !{i32 1, !"MemProfProfileFilename", !"memprof.profraw"}
  70 !6 = !{i32 8, !"PIC Level", i32 2}
  71 !7 = !{i32 7, !"PIE Level", i32 2}
  72 !8 = !{i32 7, !"uwtable", i32 2}
  73 !9 = !{!"clang"}
  74 !10 = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov", scope: !1, file: !1, line: 5, type: !11, scopeLine: 5, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
  75 !11 = !DISubroutineType(types: !12)
  76 !12 = !{}
  77 !13 = !DILocation(line: 6, column: 3, scope: !10)
  78 !14 = !DILocation(line: 7, column: 3, scope: !10)
  79 !15 = !DILocation(line: 7, column: 9, scope: !10)
  80 !16 = !DILocation(line: 8, column: 1, scope: !10)
  81 !17 = !DISubprogram(name: "f1", linkageName: "_Z2f1v", scope: !1, file: !1, line: 1, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
  82 !18 = !DISubprogram(name: "f2", linkageName: "_Z2f2v", scope: !1, file: !1, line: 2, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
  83 !19 = !DISubprogram(name: "f3", linkageName: "_Z2f3v", scope: !1, file: !1, line: 3, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
  84 )IR";
  85
  86   LLVMContext Ctx;
  87   SMDiagnostic Err;
  88   std::unique_ptr<Module> M = parseAssemblyString(IR, Err, Ctx);
  89   ASSERT_TRUE(M);
  90
  91   auto *F = M->getFunction("_Z3foov");
  92   ASSERT_NE(F, nullptr);
  93
  94   TargetLibraryInfoWrapperPass WrapperPass;
  95   auto &TLI = WrapperPass.getTLI(*F);
  96   auto Calls = extractCallsFromIR(*M, TLI);
  97
  98   // Expect exactly one caller.
  99   ASSERT_THAT(Calls, SizeIs(1));
 100
 101   auto It = Calls.begin();
 102   ASSERT_NE(It, Calls.end());
 103
 104   const auto &[CallerGUID, CallSites] = *It;
 105   EXPECT_EQ(CallerGUID, IndexedMemProfRecord::getGUID("_Z3foov"));
 106   ASSERT_THAT(CallSites, SizeIs(3));
 107
 108   // Verify that call sites show up in the ascending order of their source
 109   // locations.
 110   EXPECT_THAT(CallSites[0],
 111               Pair(FieldsAre(1U, 3U), IndexedMemProfRecord::getGUID("_Z2f1v")));
 112   EXPECT_THAT(CallSites[1],
 113               Pair(FieldsAre(2U, 3U), IndexedMemProfRecord::getGUID("_Z2f2v")));
 114   EXPECT_THAT(CallSites[2],
 115               Pair(FieldsAre(2U, 9U), IndexedMemProfRecord::getGUID("_Z2f3v")));
 116 }
 117
 118 TEST(MemProf, ExtractDirectCallsFromIRInline) {
 119   // The following IR is generated from:
 120   //
 121   // void f1();
 122   // static inline void f2() {
 123   //   // For an interesting line number.
 124   //   f1();
 125   // }
 126   // static inline void f3() {
 127   //   /****/ f2();  // For an interesting column number.
 128   // }
 129   //
 130   // void g1();
 131   // void g2();
 132   // static inline void g3() {
 133   //   /**/ g1();  // For an interesting column number.
 134   //   g2();
 135   // }
 136   //
 137   // void foo() {
 138   //   f3();
 139   //   /***/ g3();  // For an interesting column number.
 140   // }
 141   StringRef IR = R"IR(
 142 define dso_local void @_Z3foov() local_unnamed_addr !dbg !10 {
 143 entry:
 144   tail call void @_Z2f1v(), !dbg !13
 145   tail call void @_Z2g1v(), !dbg !18
 146   tail call void @_Z2g2v(), !dbg !21
 147   ret void, !dbg !22
 148 }
 149
 150 declare !dbg !23 void @_Z2f1v() local_unnamed_addr
 151
 152 declare !dbg !24 void @_Z2g1v() local_unnamed_addr
 153
 154 declare !dbg !25 void @_Z2g2v() local_unnamed_addr
 155
 156 !llvm.dbg.cu = !{!0}
 157 !llvm.module.flags = !{!2, !3, !4, !5, !6, !7, !8}
 158 !llvm.ident = !{!9}
 159
 160 !0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, debugInfoForProfiling: true, nameTableKind: None)
 161 !1 = !DIFile(filename: "foobar.cc", directory: "/")
 162 !2 = !{i32 7, !"Dwarf Version", i32 5}
 163 !3 = !{i32 2, !"Debug Info Version", i32 3}
 164 !4 = !{i32 1, !"wchar_size", i32 4}
 165 !5 = !{i32 1, !"MemProfProfileFilename", !"memprof.profraw"}
 166 !6 = !{i32 8, !"PIC Level", i32 2}
 167 !7 = !{i32 7, !"PIE Level", i32 2}
 168 !8 = !{i32 7, !"uwtable", i32 2}
 169 !9 = !{!"clang"}
 170 !10 = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov", scope: !1, file: !1, line: 17, type: !11, scopeLine: 17, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 171 !11 = !DISubroutineType(types: !12)
 172 !12 = !{}
 173 !13 = !DILocation(line: 4, column: 3, scope: !14, inlinedAt: !15)
 174 !14 = distinct !DISubprogram(name: "f2", linkageName: "_ZL2f2v", scope: !1, file: !1, line: 2, type: !11, scopeLine: 2, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 175 !15 = distinct !DILocation(line: 7, column: 10, scope: !16, inlinedAt: !17)
 176 !16 = distinct !DISubprogram(name: "f3", linkageName: "_ZL2f3v", scope: !1, file: !1, line: 6, type: !11, scopeLine: 6, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 177 !17 = distinct !DILocation(line: 18, column: 3, scope: !10)
 178 !18 = !DILocation(line: 13, column: 8, scope: !19, inlinedAt: !20)
 179 !19 = distinct !DISubprogram(name: "g3", linkageName: "_ZL2g3v", scope: !1, file: !1, line: 12, type: !11, scopeLine: 12, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 180 !20 = distinct !DILocation(line: 19, column: 9, scope: !10)
 181 !21 = !DILocation(line: 14, column: 3, scope: !19, inlinedAt: !20)
 182 !22 = !DILocation(line: 20, column: 1, scope: !10)
 183 !23 = !DISubprogram(name: "f1", linkageName: "_Z2f1v", scope: !1, file: !1, line: 1, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
 184 !24 = !DISubprogram(name: "g1", linkageName: "_Z2g1v", scope: !1, file: !1, line: 10, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
 185 !25 = !DISubprogram(name: "g2", linkageName: "_Z2g2v", scope: !1, file: !1, line: 11, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
 186 )IR";
 187
 188   LLVMContext Ctx;
 189   SMDiagnostic Err;
 190   std::unique_ptr<Module> M = parseAssemblyString(IR, Err, Ctx);
 191   ASSERT_TRUE(M);
 192
 193   auto *F = M->getFunction("_Z3foov");
 194   ASSERT_NE(F, nullptr);
 195
 196   TargetLibraryInfoWrapperPass WrapperPass;
 197   auto &TLI = WrapperPass.getTLI(*F);
 198   auto Calls = extractCallsFromIR(*M, TLI);
 199
 200   // Expect exactly 4 callers.
 201   ASSERT_THAT(Calls, SizeIs(4));
 202
 203   // Verify each key-value pair.
 204
 205   auto FooIt = Calls.find(IndexedMemProfRecord::getGUID("_Z3foov"));
 206   ASSERT_NE(FooIt, Calls.end());
 207   const auto &[FooCallerGUID, FooCallSites] = *FooIt;
 208   EXPECT_EQ(FooCallerGUID, IndexedMemProfRecord::getGUID("_Z3foov"));
 209   ASSERT_THAT(FooCallSites, SizeIs(2));
 210   EXPECT_THAT(FooCallSites[0], Pair(FieldsAre(1U, 3U),
 211                                     IndexedMemProfRecord::getGUID("_ZL2f3v")));
 212   EXPECT_THAT(FooCallSites[1], Pair(FieldsAre(2U, 9U),
 213                                     IndexedMemProfRecord::getGUID("_ZL2g3v")));
 214
 215   auto F2It = Calls.find(IndexedMemProfRecord::getGUID("_ZL2f2v"));
 216   ASSERT_NE(F2It, Calls.end());
 217   const auto &[F2CallerGUID, F2CallSites] = *F2It;
 218   EXPECT_EQ(F2CallerGUID, IndexedMemProfRecord::getGUID("_ZL2f2v"));
 219   ASSERT_THAT(F2CallSites, SizeIs(1));
 220   EXPECT_THAT(F2CallSites[0],
 221               Pair(FieldsAre(2U, 3U), IndexedMemProfRecord::getGUID("_Z2f1v")));
 222
 223   auto F3It = Calls.find(IndexedMemProfRecord::getGUID("_ZL2f3v"));
 224   ASSERT_NE(F3It, Calls.end());
 225   const auto &[F3CallerGUID, F3CallSites] = *F3It;
 226   EXPECT_EQ(F3CallerGUID, IndexedMemProfRecord::getGUID("_ZL2f3v"));
 227   ASSERT_THAT(F3CallSites, SizeIs(1));
 228   EXPECT_THAT(F3CallSites[0], Pair(FieldsAre(1U, 10U),
 229                                    IndexedMemProfRecord::getGUID("_ZL2f2v")));
 230
 231   auto G3It = Calls.find(IndexedMemProfRecord::getGUID("_ZL2g3v"));
 232   ASSERT_NE(G3It, Calls.end());
 233   const auto &[G3CallerGUID, G3CallSites] = *G3It;
 234   EXPECT_EQ(G3CallerGUID, IndexedMemProfRecord::getGUID("_ZL2g3v"));
 235   ASSERT_THAT(G3CallSites, SizeIs(2));
 236   EXPECT_THAT(G3CallSites[0],
 237               Pair(FieldsAre(1U, 8U), IndexedMemProfRecord::getGUID("_Z2g1v")));
 238   EXPECT_THAT(G3CallSites[1],
 239               Pair(FieldsAre(2U, 3U), IndexedMemProfRecord::getGUID("_Z2g2v")));
 240 }
 241
 242 TEST(MemProf, ExtractDirectCallsFromIRCallingNew) {
 243   // The following IR is generated from:
 244   //
 245   // int *foo() {
 246   //   return ::new (int);
 247   // }
 248   StringRef IR = R"IR(
 249 define dso_local noundef ptr @_Z3foov() #0 !dbg !10 {
 250 entry:
 251   %call = call noalias noundef nonnull ptr @_Znwm(i64 noundef 4) #2, !dbg !13
 252   ret ptr %call, !dbg !14
 253 }
 254
 255 ; Function Attrs: nobuiltin allocsize(0)
 256 declare noundef nonnull ptr @_Znwm(i64 noundef) #1
 257
 258 attributes #0 = { mustprogress uwtable "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
 259 attributes #1 = { nobuiltin allocsize(0) "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
 260 attributes #2 = { builtin allocsize(0) }
 261
 262 !llvm.dbg.cu = !{!0}
 263 !llvm.module.flags = !{!2, !3, !4, !5, !6, !7, !8}
 264 !llvm.ident = !{!9}
 265
 266 !0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, debugInfoForProfiling: true, nameTableKind: None)
 267 !1 = !DIFile(filename: "foobar.cc", directory: "/")
 268 !2 = !{i32 7, !"Dwarf Version", i32 5}
 269 !3 = !{i32 2, !"Debug Info Version", i32 3}
 270 !4 = !{i32 1, !"wchar_size", i32 4}
 271 !5 = !{i32 1, !"MemProfProfileFilename", !"memprof.profraw"}
 272 !6 = !{i32 8, !"PIC Level", i32 2}
 273 !7 = !{i32 7, !"PIE Level", i32 2}
 274 !8 = !{i32 7, !"uwtable", i32 2}
 275 !9 = !{!"clang"}
 276 !10 = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov", scope: !1, file: !1, line: 1, type: !11, scopeLine: 1, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 277 !11 = !DISubroutineType(types: !12)
 278 !12 = !{}
 279 !13 = !DILocation(line: 2, column: 10, scope: !10)
 280 !14 = !DILocation(line: 2, column: 3, scope: !10)
 281 )IR";
 282
 283   LLVMContext Ctx;
 284   SMDiagnostic Err;
 285   std::unique_ptr<Module> M = parseAssemblyString(IR, Err, Ctx);
 286   ASSERT_TRUE(M);
 287
 288   auto *F = M->getFunction("_Z3foov");
 289   ASSERT_NE(F, nullptr);
 290
 291   TargetLibraryInfoWrapperPass WrapperPass;
 292   auto &TLI = WrapperPass.getTLI(*F);
 293   auto Calls = extractCallsFromIR(*M, TLI);
 294
 295   // Expect exactly one caller.
 296   ASSERT_THAT(Calls, SizeIs(1));
 297
 298   // Verify each key-value pair.
 299
 300   auto FooIt = Calls.find(IndexedMemProfRecord::getGUID("_Z3foov"));
 301   ASSERT_NE(FooIt, Calls.end());
 302   const auto &[FooCallerGUID, FooCallSites] = *FooIt;
 303   EXPECT_EQ(FooCallerGUID, IndexedMemProfRecord::getGUID("_Z3foov"));
 304   ASSERT_THAT(FooCallSites, SizeIs(1));
 305   EXPECT_THAT(FooCallSites[0], Pair(FieldsAre(1U, 10U), 0));
 306 }
 307
 308 // Populate those fields returned by getHotColdSchema.
 309 MemInfoBlock makePartialMIB() {
 310   MemInfoBlock MIB;
 311   MIB.AllocCount = 1;
 312   MIB.TotalSize = 5;
 313   MIB.TotalLifetime = 10;
 314   MIB.TotalLifetimeAccessDensity = 23;
 315   return MIB;
 316 }
 317
 318 IndexedMemProfRecord
 319 makeRecordV2(std::initializer_list<::llvm::memprof::CallStackId> AllocFrames,
 320              std::initializer_list<::llvm::memprof::CallStackId> CallSiteFrames,
 321              const MemInfoBlock &Block, const memprof::MemProfSchema &Schema) {
 322   llvm::memprof::IndexedMemProfRecord MR;
 323   for (const auto &CSId : AllocFrames) {
 324     // We don't populate IndexedAllocationInfo::CallStack because we use it only
 325     // in Version1.
 326     MR.AllocSites.push_back({{}, CSId, Block, Schema});
 327   }
 328   for (const auto &CSId : CallSiteFrames)
 329     MR.CallSiteIds.push_back(CSId);
 330   return MR;
 331 }
 332
 333 static const auto Err = [](Error E) {
 334   FAIL() << E;
 335   consumeError(std::move(E));
 336 };
 337
 338 // Make sure that we can undrift direct calls.
 339 TEST(MemProf, ComputeUndriftingMap) {
 340   // Suppose that the source code has changed from:
 341   //
 342   //   void bar();
 343   //   void baz();
 344   //   void zzz();
 345   //
 346   //   void foo() {
 347   //     /**/ bar();  // LineLocation(1, 8)
 348   //     zzz();       // LineLocation(2, 3)
 349   //     baz();       // LineLocation(3, 3)
 350   //   }
 351   //
 352   // to:
 353   //
 354   //   void bar();
 355   //   void baz();
 356   //
 357   //   void foo() {
 358   //     bar();        // LineLocation(1, 3)
 359   //     /**/ baz();   // LineLocation(2, 8)
 360   //   }
 361   //
 362   // Notice that the calls to bar and baz have drifted while zzz has been
 363   // removed.
 364   StringRef IR = R"IR(
 365 define dso_local void @_Z3foov() #0 !dbg !10 {
 366 entry:
 367   call void @_Z3barv(), !dbg !13
 368   call void @_Z3bazv(), !dbg !14
 369   ret void, !dbg !15
 370 }
 371
 372 declare !dbg !16 void @_Z3barv() #1
 373
 374 declare !dbg !17 void @_Z3bazv() #1
 375
 376 attributes #0 = { mustprogress uwtable "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
 377 attributes #1 = { "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cmov,+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
 378
 379 !llvm.dbg.cu = !{!0}
 380 !llvm.module.flags = !{!2, !3, !4, !5, !6, !7, !8}
 381 !llvm.ident = !{!9}
 382
 383 !0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "clang", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, splitDebugInlining: false, debugInfoForProfiling: true, nameTableKind: None)
 384 !1 = !DIFile(filename: "foobar.cc", directory: "/")
 385 !2 = !{i32 7, !"Dwarf Version", i32 5}
 386 !3 = !{i32 2, !"Debug Info Version", i32 3}
 387 !4 = !{i32 1, !"wchar_size", i32 4}
 388 !5 = !{i32 1, !"MemProfProfileFilename", !"memprof.profraw"}
 389 !6 = !{i32 8, !"PIC Level", i32 2}
 390 !7 = !{i32 7, !"PIE Level", i32 2}
 391 !8 = !{i32 7, !"uwtable", i32 2}
 392 !9 = !{!"clang"}
 393 !10 = distinct !DISubprogram(name: "foo", linkageName: "_Z3foov", scope: !1, file: !1, line: 4, type: !11, scopeLine: 4, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
 394 !11 = !DISubroutineType(types: !12)
 395 !12 = !{}
 396 !13 = !DILocation(line: 5, column: 3, scope: !10)
 397 !14 = !DILocation(line: 6, column: 8, scope: !10)
 398 !15 = !DILocation(line: 7, column: 1, scope: !10)
 399 !16 = !DISubprogram(name: "bar", linkageName: "_Z3barv", scope: !1, file: !1, line: 1, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
 400 !17 = !DISubprogram(name: "baz", linkageName: "_Z3bazv", scope: !1, file: !1, line: 2, type: !11, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
 401 )IR";
 402
 403   LLVMContext Ctx;
 404   SMDiagnostic SMErr;
 405   std::unique_ptr<Module> M = parseAssemblyString(IR, SMErr, Ctx);
 406   ASSERT_TRUE(M);
 407
 408   auto *F = M->getFunction("_Z3foov");
 409   ASSERT_NE(F, nullptr);
 410
 411   TargetLibraryInfoWrapperPass WrapperPass;
 412   auto &TLI = WrapperPass.getTLI(*F);
 413   auto Calls = extractCallsFromIR(*M, TLI);
 414
 415   uint64_t GUIDFoo = IndexedMemProfRecord::getGUID("_Z3foov");
 416   uint64_t GUIDBar = IndexedMemProfRecord::getGUID("_Z3barv");
 417   uint64_t GUIDBaz = IndexedMemProfRecord::getGUID("_Z3bazv");
 418   uint64_t GUIDZzz = IndexedMemProfRecord::getGUID("_Z3zzzv");
 419
 420   // Verify that extractCallsFromIR extracts caller-callee pairs as expected.
 421   EXPECT_THAT(Calls,
 422               UnorderedElementsAre(Pair(
 423                   GUIDFoo, ElementsAre(Pair(LineLocation(1, 3), GUIDBar),
 424                                        Pair(LineLocation(2, 8), GUIDBaz)))));
 425
 426   llvm::InstrProfWriter Writer;
 427   std::unique_ptr<IndexedInstrProfReader> Reader;
 428
 429   const MemInfoBlock MIB = makePartialMIB();
 430
 431   Writer.setMemProfVersionRequested(memprof::Version3);
 432   Writer.setMemProfFullSchema(false);
 433
 434   ASSERT_THAT_ERROR(Writer.mergeProfileKind(InstrProfKind::MemProf),
 435                     Succeeded());
 436
 437   const std::pair<memprof::FrameId, memprof::Frame> Frames[] = {
 438       // The call sites within foo.
 439       {0, {GUIDFoo, 1, 8, false}},
 440       {1, {GUIDFoo, 2, 3, false}},
 441       {2, {GUIDFoo, 3, 3, false}},
 442       // Line/column numbers below don't matter.
 443       {3, {GUIDBar, 9, 9, false}},
 444       {4, {GUIDZzz, 9, 9, false}},
 445       {5, {GUIDBaz, 9, 9, false}}};
 446   for (const auto &[FrameId, Frame] : Frames)
 447     Writer.addMemProfFrame(FrameId, Frame, Err);
 448
 449   const std::pair<memprof::CallStackId, SmallVector<memprof::FrameId>>
 450       CallStacks[] = {
 451           {0x111, {3, 0}}, // bar called by foo
 452           {0x222, {4, 1}}, // zzz called by foo
 453           {0x333, {5, 2}}  // baz called by foo
 454       };
 455   for (const auto &[CSId, CallStack] : CallStacks)
 456     Writer.addMemProfCallStack(CSId, CallStack, Err);
 457
 458   const IndexedMemProfRecord IndexedMR = makeRecordV2(
 459       /*AllocFrames=*/{0x111, 0x222, 0x333},
 460       /*CallSiteFrames=*/{}, MIB, memprof::getHotColdSchema());
 461   Writer.addMemProfRecord(/*Id=*/0x9999, IndexedMR);
 462
 463   auto Profile = Writer.writeBuffer();
 464
 465   auto ReaderOrErr =
 466       IndexedInstrProfReader::create(std::move(Profile), nullptr);
 467   EXPECT_THAT_ERROR(ReaderOrErr.takeError(), Succeeded());
 468   Reader = std::move(ReaderOrErr.get());
 469
 470   // Verify that getMemProfCallerCalleePairs extracts caller-callee pairs as
 471   // expected.
 472   auto Pairs = Reader->getMemProfCallerCalleePairs();
 473   ASSERT_THAT(Pairs, SizeIs(4));
 474   ASSERT_THAT(
 475       Pairs,
 476       Contains(Pair(GUIDFoo, ElementsAre(Pair(LineLocation(1, 8), GUIDBar),
 477                                          Pair(LineLocation(2, 3), GUIDZzz),
 478                                          Pair(LineLocation(3, 3), GUIDBaz)))));
 479
 480   // Verify that computeUndriftMap identifies undrifting opportunities:
 481   //
 482   //   Profile                 IR
 483   //   (Line: 1, Column: 8) -> (Line: 1, Column: 3)
 484   //   (Line: 3, Column: 3) -> (Line: 2, Column: 8)
 485   auto UndriftMap = computeUndriftMap(*M, Reader.get(), TLI);
 486   ASSERT_THAT(UndriftMap,
 487               UnorderedElementsAre(Pair(
 488                   GUIDFoo, UnorderedElementsAre(
 489                                Pair(LineLocation(1, 8), LineLocation(1, 3)),
 490                                Pair(LineLocation(3, 3), LineLocation(2, 8))))));
 491 }
 492 } // namespace