[PROTON] Significantly reduce graph launching overhead with static context

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,76 @@
 #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);
-    }
-  }
+  /// `data` points to the owning data object for this entry.
+  Data *data{nullptr};
+  /// Per-entry storage for direct and linked metric maps.
+  std::reference_wrapper<MetricSet> metricSet;
+
+  explicit DataEntry(size_t id, size_t phase, Data *data, MetricSet &metricSet)
+      : id(id), phase(phase), data(data), 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 +96,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 +129,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 +153,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 +190,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,8 +213,6 @@ class Data : public ScopeInterface {
   void *currentPhasePtr{};
 };
 
-typedef std::map<Data *, DataEntry> DataToEntryMap;
-
 OutputFormat parseOutputFormat(const std::string &outputFormat);
 
 const std::string outputFormatToString(OutputFormat 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/Cupti/CuptiPCSampling.h

@@ -109,7 +109,7 @@ class CuptiPCSampling : public Singleton<CuptiPCSampling> {
 
   void start(CUcontext context);
 
-  void stop(CUcontext context, const DataToEntryMap &dataToEntry);
+  void stop(CUcontext context, const std::vector<DataEntry> &dataEntries);
 
   void finalize(CUcontext context);
 
@@ -123,7 +123,7 @@ class CuptiPCSampling : public Singleton<CuptiPCSampling> {
   CubinData *getCubinData(uint64_t cubinCrc);
 
   void processPCSamplingData(ConfigureData *configureData,
-                             const DataToEntryMap &dataToEntry);
+                             const std::vector<DataEntry> &dataEntries);
 
   ThreadSafeMap<uint32_t, ConfigureData> contextIdToConfigureData;
   // In case the same cubin is loaded multiple times, we need to keep track of

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

@@ -60,7 +60,8 @@ class GPUProfiler : public Profiler,
 
   struct ExternIdState {
     // ----non-graph launch fields----
-    DataToEntryMap dataToEntry;
+    // Active entries for each data sink associated with this extern launch.
+    std::vector<DataEntry> dataEntries;
     // Sometimes the kernel name cannot be retrieved in application threads
     // for reasons like uninitialize CUDA context.
     bool isMissingName{true};
@@ -71,32 +72,29 @@ 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{};
 
-      void setEntry(Data *data, const DataEntry &entry) {
-        dataToEntry.insert_or_assign(data, entry);
-      }
+      // 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(); }
 
-      const DataEntry *findEntry(Data *data) const {
-        auto it = dataToEntry.find(data);
-        if (it == dataToEntry.end())
-          return nullptr;
-        return &it->second;
+      void addEntry(DataEntry &&entry) {
+        dataEntries.emplace_back(std::move(entry));
       }
 
       template <typename FnT> void forEachEntry(FnT &&fn) {
-        for (auto &[data, entry] : dataToEntry)
-          fn(data, entry);
+        for (auto &entry : dataEntries)
+          fn(entry);
       }
 
-      DataToEntryMap dataToEntry;
+      // Entries in the launched graph phase for each data sink.
+      std::vector<DataEntry> dataEntries;
     };
 
     using GraphNodeStateTable = RangeTable<GraphNodeState>;
 
-    // graphNodeId -> (per-Data entry)
+    // graphNodeId -> per-node entries across active data sinks
     GraphNodeStateTable graphNodeIdToState;
   };
 
@@ -109,14 +107,13 @@ class GPUProfiler : public Profiler,
   void startOp(const Scope &scope) override {
     this->threadState.scopeStack.push_back(scope);
     for (auto *data : dataSet) {
-      auto entry = data->addOp(scope.name);
-      threadState.dataToEntry.insert_or_assign(data, entry);
+      threadState.dataEntries.push_back(data->addOp(scope.name));
     }
   }
 
   void stopOp(const Scope &scope) override {
     this->threadState.scopeStack.pop_back();
-    threadState.dataToEntry.clear();
+    threadState.dataEntries.clear();
   }
 
   void flushDataPhases(
@@ -145,7 +142,8 @@ class GPUProfiler : public Profiler,
     ConcreteProfilerT &profiler;
     SessionManager &sessionManager = SessionManager::instance();
     std::vector<Scope> scopeStack; // Used for nvtx range or triton op tracking
-    DataToEntryMap dataToEntry;
+    // Active entries for the currently open op, one per data sink.
+    std::vector<DataEntry> dataEntries;
     bool isApiExternOp{false};
     bool isStreamCapturing{false};
     bool isMetricKernelLaunching{false};
@@ -200,11 +198,12 @@ class GPUProfiler : public Profiler,
 
     // Correlate the correlationId with the last externId
     void correlate(uint64_t correlationId, size_t externId, size_t numNodes,
-                   bool isMissingName, const DataToEntryMap &dataToEntry) {
+                   bool isMissingName,
+                   const std::vector<DataEntry> &dataEntries) {
       corrIdToExternId.insert(correlationId, externId);
       externIdToState.upsert(externId, [&](ExternIdState &state) {
         state.numNodes = numNodes;
-        state.dataToEntry = dataToEntry;
+        state.dataEntries = dataEntries;
         state.isMissingName = isMissingName;
       });
     }
@@ -269,18 +268,18 @@ class GPUProfiler : public Profiler,
         auto tensorMetricsHost =
             collectTensorMetrics(profiler.metricBuffer->getRuntime(),
                                  tensorMetrics, profiler.metricKernelStream);
-        auto &dataToEntry = threadState.dataToEntry;
-        if (dataToEntry.empty()) {
+        auto &dataEntries = threadState.dataEntries;
+        if (dataEntries.empty()) {
           // Add metrics to a specific scope
           for (auto *data : profiler.dataSet) {
             data->addMetrics(scopeId, scalarMetrics);
             data->addMetrics(scopeId, tensorMetricsHost);
           }
         } 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 &entry : dataEntries) {
+            entry.upsertFlexibleMetrics(scalarMetrics);
+            entry.upsertFlexibleMetrics(tensorMetricsHost);
           }
         }
       }