[XRay][tools] Updated stacks tool with flamegraph output.

[kpw]

Summary

As the first step to allow analysis and visualization of xray collected data,
allow using the llvm-xray stacks tool to emit a complete listing of stacks in
the format consumable by a flamegraph tool.

Possible follow up formats include chrome trace viewer format and sql load
files.

As a POC, I'm able to generate flamegraphs of an xray instrumented llc compiling
hello world.

[dberris]

LGTM with some nits and style issues.

Thanks, @kpw!

I suppose you're updating the documentation on how to generate the flame graphs too, at some point?

tools/llvm-xray/xray-stacks.cc

336

In LLVM's style guide, we don't put names to TODO's (just // TODO:).

594

Same here (TODO:).

599–600

Maybe lift out of the loop and clear, instead of creating and destroying one every iteration?

621

Do you actually need a std::stack, as opposed to just using a std::vector, or a SmallVector?

738–742

nit: Enclosing curlies for single-statement bodies not required in LLVM style.

[kpw]

Updating with Dean's review comments.

[dberris]

LGTM -- thanks @kpw!

tools/llvm-xray/xray-stacks.cc

SeparateThreadStacks("per-thread-stacks",
                          cl::desc("Report top stacks within each thread id"),
                          cl::sub(Stack), cl::init(false));
 
 static cl::opt<bool>
     AggregateThreads("aggregate-threads",
                      cl::desc("Aggregate stack times across threads"),
                      cl::sub(Stack), cl::init(false));
 
-/// A helper struct to work with formatv and XRayRecords. Makes it easier to use
-/// instrumentation map names or addresses in formatted output.
+static cl::opt<bool>
+    DumpAllStacks("all-stacks",
+                  cl::desc("Dump sum of timings for all stacks. "
+                           "By default separates stacks per-thread."),
+                  cl::sub(Stack), cl::init(false));
+static cl::alias DumpAllStacksShort("all", cl::aliasopt(DumpAllStacks),
+                                    cl::desc("Alias for -all-stacks"),
+                                    cl::sub(Stack));
+
+// TODO(kpw): Add other interesting formats. Perhaps chrome trace viewer format
+// possibly with aggregations or just a linear trace of timings.
+enum StackOutputFormat { HUMAN, FLAMETOOL };
+
+static cl::opt<StackOutputFormat> StacksOutputFormat(
+    "stack-format",
+    cl::desc("The format that output stacks should be "
+             "output in. Only applies with all-stacks."),
+    cl::values(
+        clEnumValN(HUMAN, "human",
+                   "Human readable output. Only valid without -all-stacks."),
+        clEnumValN(FLAMETOOL, "flame",
+                   "Format consumable by Brendan Gregg's FlameGraph tool. "
+                   "Only valid with -all-stacks.")),
+    cl::sub(Stack), cl::init(HUMAN));
+
+// Types of values for each stack in a CallTrie.
+enum class AggregationType {
+  TOTAL_TIME,      // The total time spent in a stack and its callees.
+  INVOCATION_COUNT // The number of times the stack was invoked.
+};
+
+static cl::opt<AggregationType> RequestedAggregation(
+    "aggregation-type",
+    cl::desc("The type of aggregation to do on call stacks."),
+    cl::values(
+        clEnumValN(
+            AggregationType::TOTAL_TIME, "time",
+            "Capture the total time spent in an all invocations of a stack."),
+        clEnumValN(AggregationType::INVOCATION_COUNT, "count",
+                   "Capture the number of times a stack was invoked. "
+                   "In flamegraph mode, this count also includes invocations "
+                   "of all callees.")),
+    cl::sub(Stack), cl::init(AggregationType::TOTAL_TIME));
+
+/// A helper struct to work with formatv and XRayRecords. Makes it easier to
+/// use instrumentation map names or addresses in formatted output.
 struct format_xray_record : public FormatAdapter<XRayRecord> {
   explicit format_xray_record(XRayRecord record,
                               const FuncIdConversionHelper &conv)
       : FormatAdapter<XRayRecord>(std::move(record)), Converter(&conv) {}
   void format(raw_ostream &Stream, StringRef Style) override {
     Stream << formatv(
         "{FuncId: \"{0}\", ThreadId: \"{1}\", RecordType: \"{2}\"}",
         Converter->SymbolOrNumber(Item.FuncId), Item.TId,
TrieNode *mergeTrieNodes(const TrieNode &Left, const TrieNode &Right,
   }
   for (auto duration : Right.IntermediateDurations) {
     Node->IntermediateDurations.push_back(duration);
   }
 
   return Node;
 }
 
+template <AggregationType AggType>
+std::size_t GetValueForStack(const TrieNode *Node);
+
+// When computing total time spent in a stack, we're adding the timings from
+// its callees and the timings from when it was a leaf.
+template <>
+std::size_t
+GetValueForStack<AggregationType::TOTAL_TIME>(const TrieNode *Node) {
+  auto TopSum = std::accumulate(Node->TerminalDurations.begin(),
+                                Node->TerminalDurations.end(), 0uLL);
+  return std::accumulate(Node->IntermediateDurations.begin(),
+                         Node->IntermediateDurations.end(), TopSum);
+}
+
+// Calculates how many times a function was invoked.
+// TODO(kpw): Hook up option to produce stacks
+template <>
+std::size_t
+GetValueForStack<AggregationType::INVOCATION_COUNT>(const TrieNode *Node) {
+  return Node->TerminalDurations.size() + Node->IntermediateDurations.size();
+}
+
+// Make sure there are implementations for each enum value.
+template <AggregationType T> struct DependentFalseType : std::false_type {};
+
+template <AggregationType AggType>
+std::size_t GetValueForStack(const TrieNode *Node) {
+  static_assert(DependentFalseType<AggType>::value,
+                "No implementation found for aggregation type provided.");
+  return 0;
+}
+
 class StackTrie {
+  // Avoid the magic number of 4 propagated through the code with an alias.
+  // We use this SmallVector to track the root nodes in a call graph.
+  using RootVector = SmallVector<TrieNode *, 4>;
 
   // We maintain pointers to the roots of the tries we see.
-  DenseMap<uint32_t, SmallVector<TrieNode *, 4>> Roots;
+  DenseMap<uint32_t, RootVector> Roots;
 
   // We make sure all the nodes are accounted for in this list.
   std::forward_list<TrieNode> NodeStore;
 
   // A map of thread ids to pairs call stack trie nodes and their start times.
   DenseMap<uint32_t, SmallVector<std::pair<TrieNode *, uint64_t>, 8>>
       ThreadStackMap;
 
public:
   void printPerThread(raw_ostream &OS, FuncIdConversionHelper &FN) {
     for (auto iter : Roots) {
       OS << "Thread " << iter.first << ":\n";
       print(OS, FN, iter.second);
       OS << "\n";
     }
   }
 
+  /// Prints timing sums for each stack in each threads.
+  template <AggregationType AggType>
+  void printAllPerThread(raw_ostream &OS, FuncIdConversionHelper &FN,
+                         StackOutputFormat format) {
+    for (auto iter : Roots) {
+      uint32_t threadId = iter.first;
+      RootVector &perThreadRoots = iter.second;
+      bool reportThreadId = true;
+      printAll<AggType>(OS, FN, perThreadRoots, threadId, reportThreadId);
+    }
+  }
+
   /// Prints top stacks from looking at all the leaves and ignoring thread IDs.
   /// Stacks that consist of the same function IDs but were called in different
   /// thread IDs are not considered unique in this printout.
   void printIgnoringThreads(raw_ostream &OS, FuncIdConversionHelper &FN) {
-    SmallVector<TrieNode *, 4> RootValues;
+    RootVector RootValues;
 
     // Function to pull the values out of a map iterator.
     using RootsType = decltype(Roots.begin())::value_type;
     auto MapValueFn = [](const RootsType &Value) { return Value.second; };
 
     for (const auto &RootNodeRange :
          make_range(map_iterator(Roots.begin(), MapValueFn),
                     map_iterator(Roots.end(), MapValueFn))) {
       for (auto *RootNode : RootNodeRange)
         RootValues.push_back(RootNode);
     }
 
     print(OS, FN, RootValues);
   }
 
-  /// Merges the trie by thread id before printing top stacks.
-  void printAggregatingThreads(raw_ostream &OS, FuncIdConversionHelper &FN) {
-    std::forward_list<TrieNode> AggregatedNodeStore;
-    SmallVector<TrieNode *, 4> RootValues;
+  /// Creates a merged list of Tries for unique stacks that disregards their
+  /// thread IDs.
+  RootVector mergeAcrossThreads(std::forward_list<TrieNode> &NodeStore) {
+    RootVector MergedByThreadRoots;
     for (auto MapIter : Roots) {
       const auto &RootNodeVector = MapIter.second;
       for (auto *Node : RootNodeVector) {
-        auto MaybeFoundIter = find_if(RootValues, [Node](TrieNode *elem) {
+        auto MaybeFoundIter =
+            find_if(MergedByThreadRoots, [Node](TrieNode *elem) {
               return Node->FuncId == elem->FuncId;
             });
-        if (MaybeFoundIter == RootValues.end()) {
-          RootValues.push_back(Node);
+        if (MaybeFoundIter == MergedByThreadRoots.end()) {
+          MergedByThreadRoots.push_back(Node);
         } else {
-          RootValues.push_back(mergeTrieNodes(**MaybeFoundIter, *Node, nullptr,
-                                              AggregatedNodeStore));
-          RootValues.erase(MaybeFoundIter);
+          MergedByThreadRoots.push_back(
+              mergeTrieNodes(**MaybeFoundIter, *Node, nullptr, NodeStore));
+          MergedByThreadRoots.erase(MaybeFoundIter);
         }
       }
     }
-    print(OS, FN, RootValues);
+    return MergedByThreadRoots;
+  }
+
+  /// Print timing sums for all stacks merged by Thread ID.
+  template <AggregationType AggType>
+  void printAllAggregatingThreads(raw_ostream &OS, FuncIdConversionHelper &FN,
+                                  StackOutputFormat format) {
+    std::forward_list<TrieNode> AggregatedNodeStore;
+    RootVector MergedByThreadRoots = mergeAcrossThreads(AggregatedNodeStore);
+    bool reportThreadId = false;
+    printAll<AggType>(OS, FN, MergedByThreadRoots,
+                      /*threadId*/ 0, reportThreadId);
+  }
+
+  /// Merges the trie by thread id before printing top stacks.
+  void printAggregatingThreads(raw_ostream &OS, FuncIdConversionHelper &FN) {
+    std::forward_list<TrieNode> AggregatedNodeStore;
+    RootVector MergedByThreadRoots = mergeAcrossThreads(AggregatedNodeStore);
+    print(OS, FN, MergedByThreadRoots);
+  }
+
+  // TODO(kpw): Add a format option when more than one are supported.
+  template <AggregationType AggType>
+  void printAll(raw_ostream &OS, FuncIdConversionHelper &FN,
+                RootVector RootValues, uint32_t ThreadId, bool ReportThread) {
+    for (const auto *N : RootValues) {
+      SmallVector<const TrieNode *, 16> S;
+      S.emplace_back(N);
+      while (!S.empty()) {
+        auto *Top = S.pop_back_val();
+        printSingleStack<AggType>(OS, FN, ReportThread, ThreadId, Top);
+        for (const auto *C : Top->Callees)
+          S.push_back(C);
+      }
+    }
+  }
+
+  /// Prints values for stacks in a format consumable for the flamegraph.pl
+  /// tool. This is a line based format that lists each level in the stack
+  /// hierarchy in a semicolon delimited form followed by a space and a numeric
+  /// value. If breaking down by thread, the thread ID will be added as the
+  /// root level of the stack.
+  template <AggregationType AggType>
+  void printSingleStack(raw_ostream &OS, FuncIdConversionHelper &Converter,
+                        bool ReportThread, uint32_t ThreadId,
+                        const TrieNode *Node) {
+    if (ReportThread)
+      OS << "thread_" << ThreadId << ";";
+    std::stack<const TrieNode *> lineage{};
+    lineage.push(Node);
+    while (lineage.top()->Parent != nullptr)
+      lineage.push(lineage.top()->Parent);
+    while (!lineage.empty()) {
+      OS << Converter.SymbolOrNumber(lineage.top()->FuncId) << ";";
+      lineage.pop();
+    }
+    OS << " " << GetValueForStack<AggType>(Node) << "\n";
   }
 
   void print(raw_ostream &OS, FuncIdConversionHelper &FN,
-             SmallVector<TrieNode *, 4> RootValues) {
+             RootVector RootValues) {
     // Go through each of the roots, and traverse the call stack, producing the
     // aggregates as you go along. Remember these aggregates and stacks, and
     // show summary statistics about:
     //
     //   - Total number of unique stacks
     //   - Top 10 stacks by count
     //   - Top 10 stacks by aggregate duration
     SmallVector<std::pair<const TrieNode *, uint64_t>, 11> TopStacksByCount;
     SmallVector<std::pair<const TrieNode *, uint64_t>, 11> TopStacksBySum;
     auto greater_second = [](const std::pair<const TrieNode *, uint64_t> &A,
                              const std::pair<const TrieNode *, uint64_t> &B) {
       return A.second > B.second;
     };
     uint64_t UniqueStacks = 0;
     for (const auto *N : RootValues) {
       SmallVector<const TrieNode *, 16> S;
       S.emplace_back(N);
 
       while (!S.empty()) {
-        auto Top = S.pop_back_val();
+        auto *Top = S.pop_back_val();
 
         // We only start printing the stack (by walking up the parent pointers)
         // when we get to a leaf function.
         if (!Top->TerminalDurations.empty()) {
           ++UniqueStacks;
           auto TopSum = std::accumulate(Top->TerminalDurations.begin(),
                                         Top->TerminalDurations.end(), 0uLL);
           {
static CommandRegistration Unused(&Stack, []() -> Error {
   }
 
   if (SeparateThreadStacks && AggregateThreads)
     return make_error<StringError>(
         Twine("Can't specify options for per thread reporting and reporting "
               "that aggregates threads."),
         std::make_error_code(std::errc::invalid_argument));
 
+  if (!DumpAllStacks && StacksOutputFormat != HUMAN) {
+    return make_error<StringError>(
+        Twine("Can't specify a non-human format without -all-stacks."),
+        std::make_error_code(std::errc::invalid_argument));
+  }
+  if (DumpAllStacks && StacksOutputFormat == HUMAN) {
+    return make_error<StringError>(
+        Twine("You must specify a non-human format when reporting with "
+              "-all-stacks."),
+        std::make_error_code(std::errc::invalid_argument));
+  }
+
   symbolize::LLVMSymbolizer::Options Opts(
       symbolize::FunctionNameKind::LinkageName, true, true, false, "");
   symbolize::LLVMSymbolizer Symbolizer(Opts);
   FuncIdConversionHelper FuncIdHelper(StacksInstrMap, Symbolizer,
                                       Map.getFunctionAddresses());
   // TODO: Someday, support output to files instead of just directly to
   // standard output.
   for (const auto &Filename : StackInputs) {
for (const auto &Record : T) {
       }
     }
   }
   if (ST.isEmpty()) {
     return make_error<StringError>(
         "No instrumented calls were accounted in the input file.",
         make_error_code(errc::result_out_of_range));
   }
+
+  // Report the stacks in a long form mode for another tool to analyze.
+  if (DumpAllStacks) {
+    if (AggregateThreads) {
+      switch (RequestedAggregation) {
+      case AggregationType::TOTAL_TIME:
+        ST.printAllAggregatingThreads<AggregationType::TOTAL_TIME>(
+            outs(), FuncIdHelper, StacksOutputFormat);
+        break;
+      case AggregationType::INVOCATION_COUNT:
+        ST.printAllAggregatingThreads<AggregationType::INVOCATION_COUNT>(
+            outs(), FuncIdHelper, StacksOutputFormat);
+        break;
+      default:
+        return make_error<StringError>(
+            "Illegal value for aggregation-type.",
+            make_error_code(errc::result_out_of_range));
+      }
+    } else {
+      switch (RequestedAggregation) {
+      case AggregationType::TOTAL_TIME:
+        ST.printAllPerThread<AggregationType::TOTAL_TIME>(outs(), FuncIdHelper,
+                                                          StacksOutputFormat);
+        break;
+      case AggregationType::INVOCATION_COUNT:
+        ST.printAllPerThread<AggregationType::INVOCATION_COUNT>(
+            outs(), FuncIdHelper, StacksOutputFormat);
+        break;
+      default:
+        return make_error<StringError>(
+            "Illegal value for aggregation-type.",
+            make_error_code(errc::result_out_of_range));
+      }
+    }
+    return Error::success();
+  }
+
+  // We're only outputting top stacks.
   if (AggregateThreads) {
     ST.printAggregatingThreads(outs(), FuncIdHelper);
   } else if (SeparateThreadStacks) {
     ST.printPerThread(outs(), FuncIdHelper);
   } else {
     ST.printIgnoringThreads(outs(), FuncIdHelper);
   }
   return Error::success();
 });