[PROTON] Fix metric buffer deadlock and support multi-device metric profiling (#8943)

Jokeren · web-flow · commit 1c15a2942816 · 2025-12-08T21:58:13.000-05:00
diff --git a/third_party/proton/csrc/include/Data/Metric.h b/third_party/proton/csrc/include/Data/Metric.h
@@ -407,9 +407,9 @@ class MetricBuffer {
   void reserve() { getOrCreateBuffer(); }
 
   template <typename Func> void flush(Func callback, bool flushAll = false) {
-    std::lock_guard<std::mutex> lock(bufferMutex);
     std::vector<DeviceBuffer> buffersToFlush;
     if (flushAll) {
+      std::lock_guard<std::mutex> lock(bufferMutex);
       for (auto &[device, buffer] : deviceBuffers) {
         buffersToFlush.push_back(buffer);
       }
diff --git a/third_party/proton/csrc/lib/Data/TreeData.cpp b/third_party/proton/csrc/lib/Data/TreeData.cpp
@@ -244,7 +244,7 @@ void TreeData::dumpHatchet(std::ostream &os) const {
             kernelMetric->getValueName(KernelMetric::Duration));
         inclusiveValueNames.insert(
             kernelMetric->getValueName(KernelMetric::Invocations));
-        deviceIds.insert({deviceType, {deviceId}});
+        deviceIds[deviceType].insert(deviceId);
       } else if (metricKind == MetricKind::PCSampling) {
         auto pcSamplingMetric =
             std::dynamic_pointer_cast<PCSamplingMetric>(metric);
@@ -276,7 +276,7 @@ void TreeData::dumpHatchet(std::ostream &os) const {
         (*jsonNode)["metrics"]
                    [cycleMetric->getValueName(CycleMetric::DeviceType)] =
                        std::to_string(deviceType);
-        deviceIds.insert({deviceType, {deviceId}});
+        deviceIds[deviceType].insert(deviceId);
       } else if (metricKind == MetricKind::Flexible) {
         // Flexible metrics are handled in a different way
       } else {
@@ -313,12 +313,12 @@ void TreeData::dumpHatchet(std::ostream &os) const {
   // problems
   output.push_back(json::object());
   auto &deviceJson = output.back();
-  for (auto [deviceType, deviceIds] : deviceIds) {
+  for (auto [deviceType, deviceIdSet] : deviceIds) {
     auto deviceTypeName =
         getDeviceTypeString(static_cast<DeviceType>(deviceType));
     if (!deviceJson.contains(deviceTypeName))
       deviceJson[deviceTypeName] = json::object();
-    for (auto deviceId : deviceIds) {
+    for (auto deviceId : deviceIdSet) {
       Device device = getDevice(static_cast<DeviceType>(deviceType), deviceId);
       deviceJson[deviceTypeName][std::to_string(deviceId)] = {
           {"clock_rate", device.clockRate},
diff --git a/third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp b/third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp
@@ -268,53 +268,73 @@ struct GraphState {
   size_t numInstances{1};
 };
 
-// FIXME: it should be a per-stream queue in case we capture graphs from
-// different streams or different devices
+// Track pending graphs per device so flushing a single device won't drain
+// graphs from other devices.
 class PendingGraphQueue {
 public:
+  explicit PendingGraphQueue(Runtime *runtime) : runtime(runtime) {}
+
   struct PendingGraph {
     size_t externId;
     std::map<Data *, std::vector<std::pair<bool, size_t>>> dataToScopeIds;
     size_t numMetricNodes;
   };
   using PopResult = std::pair<size_t, std::vector<PendingGraph>>;
 
-  PendingGraphQueue() = default;
-
   void push(size_t externId,
             const std::map<Data *, std::vector<std::pair<bool, size_t>>>
                 &dataToScopeIds,
             size_t numNodes) {
     std::lock_guard<std::mutex> lock(mutex);
-    pendingGraphs.push_back(PendingGraph{externId, dataToScopeIds, numNodes});
-    this->totalNumNodes += numNodes;
+    auto device = runtime->getDevice();
+    auto &queue = deviceQueues[device];
+    queue.pendingGraphs.push_back(
+        PendingGraph{externId, dataToScopeIds, numNodes});
+    queue.totalNumNodes += numNodes;
   }
 
-  PopResult popAllIfReachCapacity(size_t numNewNodes, size_t capacity) {
+  PopResult pop(size_t numNewNodes, size_t capacity) {
     std::lock_guard<std::mutex> lock(mutex);
-    if ((this->totalNumNodes + numNewNodes) * 2 * sizeof(uint64_t) <=
+    if (deviceQueues.empty()) {
+      return {0, {}};
+    }
+    auto device = runtime->getDevice();
+    auto &queue = deviceQueues[device];
+    if ((queue.totalNumNodes + numNewNodes) * 2 * sizeof(uint64_t) <=
         capacity) {
       return {0, {}};
     }
-    return popAllLocked();
+    return popLocked(queue);
   }
 
-  PopResult popAll() {
+  std::vector<PopResult> popAll() {
     std::lock_guard<std::mutex> lock(mutex);
-    return popAllLocked();
+    if (deviceQueues.empty()) {
+      return {{0, {}}};
+    }
+    std::vector<PopResult> results;
+    for (auto &[device, queue] : deviceQueues) {
+      results.emplace_back(popLocked(queue));
+    }
+    return results;
   }
 
 private:
-  PopResult popAllLocked() {
+  struct Queue {
+    size_t totalNumNodes{};
+    std::vector<PendingGraph> pendingGraphs;
+  };
+
+  PopResult popLocked(Queue &queue) {
     std::vector<PendingGraph> items;
-    items.swap(pendingGraphs);
-    size_t numNodes = totalNumNodes;
-    totalNumNodes = 0;
+    items.swap(queue.pendingGraphs);
+    size_t numNodes = queue.totalNumNodes;
+    queue.totalNumNodes = 0;
     return {numNodes, items};
   }
 
-  size_t totalNumNodes{};
-  std::vector<PendingGraph> pendingGraphs;
+  Runtime *runtime{};
+  std::map<void *, Queue> deviceQueues;
   mutable std::mutex mutex;
 };
 
@@ -323,7 +343,8 @@ class PendingGraphQueue {
 struct CuptiProfiler::CuptiProfilerPimpl
     : public GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface {
   CuptiProfilerPimpl(CuptiProfiler &profiler)
-      : GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface(profiler) {
+      : GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface(profiler),
+        pendingGraphQueue(&CudaRuntime::instance()) {
     runtime = &CudaRuntime::instance();
     metricBuffer = std::make_unique<MetricBuffer>(1024 * 1024 * 64, runtime);
   }
@@ -619,8 +640,33 @@ void CuptiProfiler::CuptiProfilerPimpl::callbackFn(void *userData,
               }
             }
           }
+        }
+      }
+      profiler.correlation.correlate(callbackData->correlationId, numInstances);
+      if (profiler.pcSamplingEnabled && isDriverAPILaunch(cbId)) {
+        pImpl->pcSampling.start(callbackData->context);
+      }
+    } else if (callbackData->callbackSite == CUPTI_API_EXIT) {
+      auto externId = profiler.correlation.externIdQueue.back();
+      if (profiler.pcSamplingEnabled && isDriverAPILaunch(cbId)) {
+        // XXX: Conservatively stop every GPU kernel for now
+        pImpl->pcSampling.stop(
+            callbackData->context, externId,
+            profiler.correlation.apiExternIds.contain(externId));
+      }
+      if (cbId == CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch ||
+          cbId == CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch_ptsz) {
+        // Cuda context can be lazily initialized, so we need to call device get
+        // here after the first kernel is launched
+        auto graphExec = static_cast<const cuGraphLaunch_params *>(
+                             callbackData->functionParams)
+                             ->hGraph;
+        uint32_t graphExecId = 0;
+        cupti::getGraphExecId<true>(graphExec, &graphExecId);
+        if (pImpl->graphStates.contain(graphExecId)) {
           std::map<Data *, std::vector<std::pair<bool, size_t>>>
               metricNodeScopes;
+          auto dataSet = profiler.getDataSet();
           for (auto *data : dataSet) {
             auto &nodeToScopeId =
                 profiler.correlation.externIdToGraphNodeScopeId[externId][data];
@@ -639,8 +685,8 @@ void CuptiProfiler::CuptiProfilerPimpl::callbackFn(void *userData,
               pImpl->metricBuffer->getCapacity(); // bytes
           auto metricNodeCount =
               pImpl->graphStates[graphExecId].metricKernelNodeIds.size();
-          auto drained = pImpl->pendingGraphQueue.popAllIfReachCapacity(
-              metricNodeCount, metricBufferCapacity);
+          auto drained = pImpl->pendingGraphQueue.pop(metricNodeCount,
+                                                      metricBufferCapacity);
           if (drained.first != 0) { // Reached capacity
             pImpl->metricBuffer->flush([&](uint8_t *data, size_t dataSize) {
               auto *recordPtr = reinterpret_cast<uint64_t *>(data);
@@ -651,18 +697,6 @@ void CuptiProfiler::CuptiProfilerPimpl::callbackFn(void *userData,
                                         metricNodeCount);
         }
       }
-      profiler.correlation.correlate(callbackData->correlationId, numInstances);
-      if (profiler.pcSamplingEnabled && isDriverAPILaunch(cbId)) {
-        pImpl->pcSampling.start(callbackData->context);
-      }
-    } else if (callbackData->callbackSite == CUPTI_API_EXIT) {
-      if (profiler.pcSamplingEnabled && isDriverAPILaunch(cbId)) {
-        // XXX: Conservatively stop every GPU kernel for now
-        auto scopeId = profiler.correlation.externIdQueue.back();
-        pImpl->pcSampling.stop(
-            callbackData->context, scopeId,
-            profiler.correlation.apiExternIds.contain(scopeId));
-      }
       threadState.exitOp();
       profiler.correlation.submit(callbackData->correlationId);
     }
@@ -713,14 +747,17 @@ void CuptiProfiler::CuptiProfilerPimpl::doFlush() {
   // new activities.
   cupti::activityFlushAll<true>(/*flag=*/CUPTI_ACTIVITY_FLAG_FLUSH_FORCED);
   // Flush the tensor metric buffer
-  auto dataSet = profiler.getDataSet();
   auto popResult = pendingGraphQueue.popAll();
-  metricBuffer->flush(
-      [&](uint8_t *data, size_t dataSize) {
-        auto *recordPtr = reinterpret_cast<uint64_t *>(data);
-        emitMetricRecords(recordPtr, popResult.second);
-      },
-      /*flushAll=*/true);
+  if (!popResult.empty()) {
+    auto resultIdx = 0;
+    metricBuffer->flush(
+        [&](uint8_t *data, size_t dataSize) {
+          auto *recordPtr = reinterpret_cast<uint64_t *>(data);
+          emitMetricRecords(recordPtr, popResult[resultIdx].second);
+          resultIdx++;
+        },
+        /*flushAll=*/true);
+  }
 }
 
 void CuptiProfiler::CuptiProfilerPimpl::doStop() {
diff --git a/third_party/proton/test/test_profile.py b/third_party/proton/test/test_profile.py
@@ -669,3 +669,94 @@ def fn():
     assert scope_a_frame["metrics"]["bytes"] == 160
     assert scope_b_frame is not None
     assert scope_b_frame["metrics"]["sum"] == 40.0
+
+
+@pytest.mark.skipif(is_hip(), reason="HIP backend does not support metrics profiling in cudagraphs")
+def test_tensor_metrics_multi_device_cudagraph(tmp_path: pathlib.Path):
+    if torch.cuda.device_count() < 2:
+        pytest.skip("Requires at least two CUDA devices")
+
+    devices = [torch.device(f"cuda:{i}") for i in range(2)]
+    streams = []
+    for device in devices:
+        with torch.cuda.device(device):
+            streams.append(torch.cuda.Stream(device=device))
+
+    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
+        x = args["x"]
+        x_sum = x.sum()
+        device_idx = x.device.index
+        return {"name": f"foo_test_{device_idx}", "bytes": x.numel() * x.element_size(), "flops": x_sum}
+
+    @triton.jit(launch_metadata=metadata_fn)
+    def foo(x, y, z):
+        tl.store(z, tl.load(y) + tl.load(x))
+
+    def run_on_device(device_id):
+        with proton.scope(f"scope_a_{device_id}", metrics={"bytes": 4 * 4}):
+            a = torch.ones((2, 2), device=f"cuda:{device_id}")
+        with proton.metadata_state():
+            a_sum = a.sum()
+        with proton.scope(f"scope_b_{device_id}", metrics={"sum": a_sum}):
+            b = torch.ones((2, 2), device=f"cuda:{device_id}")
+        c = a + b
+        foo[(1, )](a, b, c)
+
+    temp_file = tmp_path / "test_tensor_metrics_multi_device_cudagraph.hatchet"
+    proton.start(str(temp_file.with_suffix("")), context="shadow", hook="triton")
+
+    graphs = []
+    for device, stream in zip(devices, streams):
+        with torch.cuda.device(device):
+            torch.cuda.set_stream(stream)
+            # warmup
+            run_on_device(device.index)
+            # graph capture
+            g = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g, stream=stream):
+                for _ in range(10):
+                    run_on_device(device.index)
+        graphs.append((device, stream, g))
+
+    for device, stream, graph in graphs:
+        with torch.cuda.device(device):
+            torch.cuda.set_stream(stream)
+            with proton.scope(f"test_device_{device.index}"):
+                graph.replay()
+
+    proton.finalize()
+
+    with temp_file.open() as f:
+        data = json.load(f)
+
+    children = data[0]["children"]
+    for device in devices:
+        device_name = f"test_device_{device.index}"
+        launch_frame = next((child for child in children if child["frame"]["name"] == device_name), None)
+        assert launch_frame is not None
+        capture_at_frame = launch_frame["children"][0]
+        assert capture_at_frame["frame"]["name"] == "<captured_at>"
+
+        foo_frame = None
+        scope_a_frame = None
+        scope_b_frame = None
+        for child in capture_at_frame["children"]:
+            if child["frame"]["name"] == f"foo_test_{device.index}":
+                foo_frame = child
+            if child["frame"]["name"] == f"scope_a_{device.index}":
+                scope_a_frame = child
+            if child["frame"]["name"] == f"scope_b_{device.index}":
+                scope_b_frame = child
+
+        assert foo_frame is not None
+        assert scope_a_frame is not None
+        assert scope_b_frame is not None
+        assert foo_frame["metrics"]["bytes"] == 160
+        assert foo_frame["metrics"]["flops"] == 40
+        assert foo_frame["metrics"]["device_id"] == str(device.index)
+        assert scope_a_frame["metrics"]["bytes"] == 160
+        assert scope_b_frame["metrics"]["sum"] == 40.0
+
+    assert len(data) > 1
+    cuda_devices = data[1].get("CUDA", {})
+    assert len(cuda_devices) >= 2
