[PROTON] Significantly reduce graph launching overhead with static context (#9405)

Author: Jokeren

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

@@ -6,6 +6,7 @@
 #include "PhaseStore.h"
 #include <atomic>
 #include <cstdint>
+#include <functional>
 #include <limits>
 #include <map>
 #include <memory>
@@ -15,48 +16,74 @@
 #include <shared_mutex>
 #include <stdexcept>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
 namespace proton {
 
 enum class OutputFormat { Hatchet, HatchetMsgPack, ChromeTrace, Count };
 
+class Data;
+
 /// An "entry" is a data specific unit of operation, e.g., a node in a tree
 /// data structure or an event in a trace data structure.
 struct DataEntry {
-  /// `entryId` is a unique identifier for the entry in the data.
+  using MetricMap = std::map<MetricKind, std::unique_ptr<Metric>>;
+  using FlexibleMetricMap = std::map<std::string, FlexibleMetric>;
+  using LinkedMetricMap = std::unordered_map<size_t, MetricMap>;
+  using LinkedFlexibleMetricMap = std::unordered_map<size_t, FlexibleMetricMap>;
+  struct MetricSet {
+    // Direct metrics associated with this entry.
+    MetricMap metrics{};
+    // Direct flexible metrics associated with this entry.
+    FlexibleMetricMap flexibleMetrics{};
+    // Metrics associated with linked entries.
+    LinkedMetricMap linkedMetrics{};
+    // Flexible metrics associated with linked entries.
+    LinkedFlexibleMetricMap linkedFlexibleMetrics{};
+  };
+
+  /// `id` is a unique identifier for the entry in the data.
+  /// When `phase` is a virtual phase, `id` refers to the linked entry id
+  /// for the node entry.
   size_t id{Scope::DummyScopeId};
   /// `phase` indicates which phase the entry belongs to.
   size_t phase{0};
-  /// `metrics` is a map from metric kind to metric accumulator associated
-  /// with the entry.
-  /// Flexible metrics cannot be directly stored here since they maybe added by
-  /// both the frontend and the backend.
-  /// Use `Data::addMetrics` and `Data::addMetrics` to add flexible
-  /// metrics.
-  std::reference_wrapper<std::map<MetricKind, std::unique_ptr<Metric>>> metrics;
-
-  explicit DataEntry(size_t id, size_t phase,
-                     std::map<MetricKind, std::unique_ptr<Metric>> &metrics)
-      : id(id), phase(phase), metrics(metrics) {}
-
-  void upsertMetric(std::unique_ptr<Metric> metric) {
-    if (!metric)
-      return;
-    auto &metricsMap = metrics.get();
-    auto it = metricsMap.find(metric->getKind());
-    if (it == metricsMap.end()) {
-      metricsMap.emplace(metric->getKind(), std::move(metric));
-    } else {
-      it->second->updateMetric(*metric);
-    }
-  }
+  /// Per-entry storage for direct and linked metric maps.
+  std::reference_wrapper<MetricSet> metricSet;
+
+  explicit DataEntry(size_t id, size_t phase, MetricSet &metricSet)
+      : id(id), phase(phase), metricSet(metricSet) {}
+
+  void upsertMetric(std::unique_ptr<Metric> metric) const;
+
+  void upsertLinkedMetric(std::unique_ptr<Metric> metric,
+                          size_t linkedId) const;
+
+  void upsertFlexibleMetric(const std::string &metricName,
+                            const MetricValueType &metricValue) const;
+
+  void upsertFlexibleMetrics(
+      const std::map<std::string, MetricValueType> &metrics) const;
+
+  void upsertLinkedFlexibleMetric(const std::string &metricName,
+                                  const MetricValueType &metricValue,
+                                  size_t linkedId) const;
+
+  void upsertLinkedFlexibleMetrics(
+      const std::map<std::string, MetricValueType> &metrics,
+      size_t linkedId) const;
 };
 
 class Data : public ScopeInterface {
 public:
   static constexpr size_t kNoCompletePhase = std::numeric_limits<size_t>::max();
+  // A special phase used for static/captured graph metadata.
+  static constexpr size_t kVirtualPhase =
+      std::numeric_limits<size_t>::max() - 1;
+  // Sentinel root id used when adding an op from the root.
+  static constexpr size_t kRootEntryId = Scope::DummyScopeId;
 
   struct PhaseInfo {
     size_t current{0};
@@ -67,7 +94,7 @@ class Data : public ScopeInterface {
     }
   };
 
-  Data(const std::string &path, ContextSource *contextSource = nullptr)
+  Data(const std::string &path, ContextSource *contextSource)
       : path(path), contextSource(contextSource) {}
   virtual ~Data() = default;
 
@@ -100,7 +127,7 @@ class Data : public ScopeInterface {
   /// If `opName` is empty, just use the current context as is.
   /// Otherwise obtain the current context and append `opName` to it. Return the
   /// entry id of the added op.
-  virtual DataEntry addOp(const std::string &opName = {}) = 0;
+  DataEntry addOp(const std::string &opName = {});
 
   /// Add an op with custom contexts to the data.
   /// This is often used when context source is not available or when
@@ -124,17 +151,6 @@ class Data : public ScopeInterface {
   addMetrics(size_t scopeId,
              const std::map<std::string, MetricValueType> &metrics) = 0;
 
-  /// Record a batch of named metrics for an entry.
-  ///
-  /// This is primarily intended for user-defined metrics defined in Python and
-  /// added lazily by the backend profiler.
-  /// `metrics` is a map from metric name to value to be applied to `entryId`.
-  ///
-  /// The same as `addOp`, `phase` is important for asynchronous profilers.
-  virtual void
-  addMetrics(size_t phase, size_t entryId,
-             const std::map<std::string, MetricValueType> &metrics) = 0;
-
   /// To Json
   virtual std::string toJsonString(size_t phase) const = 0;
 
@@ -172,6 +188,16 @@ class Data : public ScopeInterface {
     return lock;
   }
 
+  [[nodiscard]] std::unique_lock<std::shared_mutex>
+  lockIfCurrentOrStaticPhase(size_t phase) {
+    std::unique_lock<std::shared_mutex> lock(mutex, std::defer_lock);
+    const auto currentPhaseValue = currentPhase.load(std::memory_order_relaxed);
+    if (phase == currentPhaseValue || phase == kVirtualPhase) {
+      lock.lock();
+    }
+    return lock;
+  }
+
   std::atomic<std::size_t> currentPhase{0};
   std::size_t completeUpToPhase{kNoCompletePhase};
   std::set<size_t> activePhases{};
@@ -185,7 +211,7 @@ class Data : public ScopeInterface {
   void *currentPhasePtr{};
 };
 
-typedef std::map<Data *, DataEntry> DataToEntryMap;
+using DataToEntryMap = std::map<Data *, DataEntry>;
 
 OutputFormat parseOutputFormat(const std::string &outputFormat);
 

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

@@ -16,19 +16,13 @@ class TraceData : public Data {
 
   std::vector<uint8_t> toMsgPack(size_t phase) const override;
 
-  DataEntry addOp(const std::string &name) override;
-
   DataEntry addOp(size_t phase, size_t eventId,
                   const std::vector<Context> &contexts) override;
 
   void
   addMetrics(size_t scopeId,
              const std::map<std::string, MetricValueType> &metrics) override;
 
-  void
-  addMetrics(size_t phase, size_t entryId,
-             const std::map<std::string, MetricValueType> &metrics) override;
-
   class Trace;
 
 protected:

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

@@ -24,19 +24,13 @@ class TreeData : public Data {
 
   std::vector<uint8_t> toMsgPack(size_t phase) const override;
 
-  DataEntry addOp(const std::string &name) override;
-
   DataEntry addOp(size_t phase, size_t contextId,
                   const std::vector<Context> &contexts) override;
 
   void
   addMetrics(size_t scopeId,
              const std::map<std::string, MetricValueType> &metrics) override;
 
-  void
-  addMetrics(size_t phase, size_t entryId,
-             const std::map<std::string, MetricValueType> &metrics) override;
-
 protected:
   // ScopeInterface
   void enterScope(const Scope &scope) override;
@@ -48,8 +42,9 @@ class TreeData : public Data {
   // the background threads concurrently, so methods that access them should be
   // protected by a (shared) mutex.
   class Tree;
-  json buildHatchetJson(TreeData::Tree *tree) const;
-  std::vector<uint8_t> buildHatchetMsgPack(TreeData::Tree *tree) const;
+  json buildHatchetJson(TreeData::Tree *tree, TreeData::Tree *staticTree) const;
+  std::vector<uint8_t> buildHatchetMsgPack(TreeData::Tree *tree,
+                                           TreeData::Tree *staticTree) const;
 
   // Data
   void doDump(std::ostream &os, OutputFormat outputFormat,

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

@@ -71,9 +71,12 @@ class GPUProfiler : public Profiler,
     size_t numNodes{1};
 
     struct GraphNodeState {
-      // If the node is launched as a metric kernel, ignore it's timing data.
-      bool isMetricNode{false};
-      bool isMissingName{true};
+      // Per-node launch status bits (missing-name / metric-node).
+      NodeStatus status{};
+
+      // If the node is launched as a metric kernel, ignore its timing data.
+      bool isMetricNode() const { return status.isMetricNode(); }
+      bool isMissingName() const { return status.isMissingName(); }
 
       void setEntry(Data *data, const DataEntry &entry) {
         dataToEntry.insert_or_assign(data, entry);
@@ -96,7 +99,7 @@ class GPUProfiler : public Profiler,
 
     using GraphNodeStateTable = RangeTable<GraphNodeState>;
 
-    // graphNodeId -> (per-Data entry)
+    // graphNodeId -> per-node entries across active data sinks
     GraphNodeStateTable graphNodeIdToState;
   };
 
@@ -278,9 +281,10 @@ class GPUProfiler : public Profiler,
           }
         } else {
           // Add metrics to the current op
-          for (auto [data, entry] : dataToEntry) {
-            data->addMetrics(entry.phase, entry.id, scalarMetrics);
-            data->addMetrics(entry.phase, entry.id, tensorMetricsHost);
+          for (const auto &entryIt : dataToEntry) {
+            const auto &entry = entryIt.second;
+            entry.upsertFlexibleMetrics(scalarMetrics);
+            entry.upsertFlexibleMetrics(tensorMetricsHost);
           }
         }
       }