[PROTON] Show metadata and graph scopes in the trace (#9930)

We differentiate CPU and GPU events using three different trace lines
per process, including:

1. CPU scopes
2. GPU Graph scopes
3. GPU events

## CPU scopes

Timestamps are recorded when a scope is entered and exited.

Metrics are shown in the CPU scope names.

## GPU graph scopes

Timestamps are aligned with the underlying GPU events.

Metrics are shown in the GPU scope names.  

## GPU events

### Non-graph GPU kernels

Each GPU event has a dependency on the CPU scope where it was launched.

### Graph-based GPU kernels

Only the first kernel of each graph has a dependency on the CPU scope
where it was launched.

We show all kernels launched by either proton or the application in the
traces, including metadata kernels (e.g., data type conversion), metric
recording kernels, and application kernels.

Author: Jokeren

third_party/proton/csrc/include/Data/Metric.h

@@ -194,13 +194,15 @@ class KernelMetric : public Metric {
     DeviceId,
     DeviceType,
     StreamId,
+    IsMetricKernel,
     Count,
   };
 
   KernelMetric() : Metric(MetricKind::Kernel, kernelMetricKind::Count) {}
 
   KernelMetric(uint64_t startTime, uint64_t endTime, uint64_t invocations,
-               uint64_t deviceId, uint64_t deviceType, uint64_t streamId)
+               uint64_t deviceId, uint64_t deviceType, uint64_t streamId,
+               uint64_t isMetricKernel = 0)
       : KernelMetric() {
     this->values[StartTime] = startTime;
     this->values[EndTime] = endTime;
@@ -209,6 +211,7 @@ class KernelMetric : public Metric {
     this->values[DeviceId] = deviceId;
     this->values[DeviceType] = deviceType;
     this->values[StreamId] = streamId;
+    this->values[IsMetricKernel] = isMetricKernel;
   }
 
   const std::string &getName() const override { return name; }
@@ -223,12 +226,12 @@ class KernelMetric : public Metric {
 
 private:
   const static inline bool PROPERTY[kernelMetricKind::Count] = {
-      true, true, false, false, true, true, true};
+      true, true, false, false, true, true, true, true};
   const static inline bool EXCLUSIVE[kernelMetricKind::Count] = {
-      false, false, false, false, true, true, true};
+      false, false, false, false, true, true, true, true};
   const static inline std::string VALUE_NAMES[kernelMetricKind::Count] = {
       "start_time (ns)", "end_time (ns)", "count",     "time (ns)",
-      "device_id",       "device_type",   "stream_id",
+      "device_id",       "device_type",   "stream_id", "is_metric_kernel",
   };
   const static inline std::string name = "KernelMetric";
 };

third_party/proton/csrc/include/Data/TraceData.h

@@ -3,6 +3,7 @@
 
 #include "Data.h"
 #include <memory>
+#include <thread>
 #include <unordered_map>
 
 namespace proton {
@@ -41,10 +42,14 @@ class TraceData : public Data {
   }
 
   void dumpChromeTrace(std::ostream &os, size_t phase) const;
+  size_t getCurrentThreadTraceId();
 
   PhaseStore<Trace> tracePhases;
   // ScopeId -> EventId
   std::unordered_map<size_t, size_t> scopeIdToEventId;
+  // ThreadId -> TraceId
+  std::unordered_map<std::thread::id, uint64_t> threadIdToTraceId;
+  uint64_t nextThreadTraceId = 0;
 };
 
 } // namespace proton

third_party/proton/csrc/include/Profiler/Graph.h

@@ -45,6 +45,8 @@ struct NodeStatus {
 struct GraphState {
   // Capture tag to identify captured call paths
   static constexpr const char *captureTag = "<captured_at>";
+  static constexpr const char *metricTag = "<metric>";
+  static constexpr const char *metadataTag = "__proton_launch_metadata";
   struct NodeState {
     // The graph node id for this node
     uint64_t nodeId{};