From 734f4d832cd2d3a9520d2384c98e1429b53e0238 Mon Sep 17 00:00:00 2001
From: Hongtao Yu <hoy@fb.com>
Date: Wed, 22 Sep 2021 09:11:53 -0700
Subject: [PATCH] [llvm-profgen] An option to dump disasm of specified symbols

For large app, dumping disasm of the whole program can be slow and result in gianant output. Adding a switch to dump specific symbols only.

Reviewed By: wlei

Differential Revision: https://reviews.llvm.org/D110079
---
 .../llvm-profgen/pseudoprobe-decoding.test    |  6 ++++++
 llvm/tools/llvm-profgen/ProfiledBinary.cpp    | 20 ++++++++++++++++---
 llvm/tools/llvm-profgen/ProfiledBinary.h      |  3 +++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/llvm/test/tools/llvm-profgen/pseudoprobe-decoding.test b/llvm/test/tools/llvm-profgen/pseudoprobe-decoding.test
index 45be13c5d39e..cd69ea82c09e 100644
--- a/llvm/test/tools/llvm-profgen/pseudoprobe-decoding.test
+++ b/llvm/test/tools/llvm-profgen/pseudoprobe-decoding.test
@@ -1,4 +1,5 @@
 ; RUN: llvm-profgen --format=text --perfscript=%s  --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-pseudo-probe --show-disassembly-only | FileCheck %s
+; RUN: llvm-profgen --format=text --perfscript=%s  --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-pseudo-probe --show-disassembly-only --disassemble-functions=main,foo  | FileCheck %s -check-prefix=SYM
 
 PERF_RECORD_MMAP2 2854748/2854748: [0x400000(0x1000) @ 0 00:1d 123291722 526021]: r-xp /home/inline-cs-pseudoprobe.perfbin
 
@@ -79,6 +80,11 @@ PERF_RECORD_MMAP2 2854748/2854748: [0x400000(0x1000) @ 0 00:1d 123291722 526021]
 ; CHECK:       [Probe]: FUNC: foo Index: 9  Type: DirectCall  Inlined: @ main:2
 ; CHECK-NEXT:      865: callq 0x930
 
+; SYM-NOT: <bar>:
+; SYM: <foo>:
+; SYM: <main>:
+
+
 
 ; clang -O3 -fexperimental-new-pass-manager -fuse-ld=lld -fpseudo-probe-for-profiling
 ; -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -Xclang -mdisable-tail-calls
diff --git a/llvm/tools/llvm-profgen/ProfiledBinary.cpp b/llvm/tools/llvm-profgen/ProfiledBinary.cpp
index 9d17d7858237..f817d0c12e40 100644
--- a/llvm/tools/llvm-profgen/ProfiledBinary.cpp
+++ b/llvm/tools/llvm-profgen/ProfiledBinary.cpp
@@ -38,6 +38,11 @@ cl::opt<bool> ShowPseudoProbe(
     "show-pseudo-probe", cl::ReallyHidden, cl::init(false), cl::ZeroOrMore,
     cl::desc("Print pseudo probe section and disassembled info."));
 
+static cl::list<std::string> DisassembleFunctions(
+    "disassemble-functions", cl::CommaSeparated,
+    cl::desc("List of functions to print disassembly for. Accept demangled "
+             "names only. Only work with show-disassembly-only"));
+
 namespace llvm {
 namespace sampleprof {
 
@@ -298,7 +303,10 @@ bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
       ShowCanonicalFnName
           ? FunctionSamples::getCanonicalFnName(Symbols[SI].Name)
           : Symbols[SI].Name;
-  if (ShowDisassemblyOnly)
+  bool ShowDisassembly =
+      ShowDisassemblyOnly && (DisassembleFunctionSet.empty() ||
+                              DisassembleFunctionSet.count(SymbolName));
+  if (ShowDisassembly)
     outs() << '<' << SymbolName << ">:\n";
 
   auto WarnInvalidInsts = [](uint64_t Start, uint64_t End) {
@@ -321,7 +329,7 @@ bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
     if (Size == 0)
       Size = 1;
 
-    if (ShowDisassemblyOnly) {
+    if (ShowDisassembly) {
       if (ShowPseudoProbe) {
         ProbeDecoder.printProbeForAddress(outs(),
                                           Offset + getPreferredBaseAddress());
@@ -385,7 +393,7 @@ bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
   if (InvalidInstLength)
     WarnInvalidInsts(Offset - InvalidInstLength, Offset - 1);
 
-  if (ShowDisassemblyOnly)
+  if (ShowDisassembly)
     outs() << "\n";
 
   FuncStartAddrMap[StartOffset] = Symbols[SI].Name.str();
@@ -452,6 +460,12 @@ void ProfiledBinary::disassemble(const ELFObjectFileBase *Obj) {
   for (std::pair<const SectionRef, SectionSymbolsTy> &SecSyms : AllSymbols)
     stable_sort(SecSyms.second);
 
+  DisassembleFunctionSet.insert(DisassembleFunctions.begin(),
+                                DisassembleFunctions.end());
+  assert((DisassembleFunctionSet.empty() || ShowDisassemblyOnly) &&
+         "Functions to disassemble should be only specified together with "
+         "--show-disassembly-only");
+
   if (ShowDisassemblyOnly)
     outs() << "\nDisassembly of " << FileName << ":\n";
 
diff --git a/llvm/tools/llvm-profgen/ProfiledBinary.h b/llvm/tools/llvm-profgen/ProfiledBinary.h
index fabc73d3b567..9e4d11bf19ba 100644
--- a/llvm/tools/llvm-profgen/ProfiledBinary.h
+++ b/llvm/tools/llvm-profgen/ProfiledBinary.h
@@ -185,6 +185,9 @@ class ProfiledBinary {
   // String table owning function name strings created from the symbolizer.
   std::unordered_set<std::string> NameStrings;
 
+  // A collection of functions to print disassembly for.
+  StringSet<> DisassembleFunctionSet;
+
   // Pseudo probe decoder
   MCPseudoProbeDecoder ProbeDecoder;