NPKit: separate time stamps for GPU access from different blocks (#1229)

To avoid races in memory access in GPU
2024-06-28 08:00:22 -07:00
parent b09ea29d66
commit 9d8f68b4ee
7 changed files with 17 additions and 21 deletions
@@ -36,8 +36,7 @@ namespace {

 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
    if (tid == 0) {
-      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
-      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, NPKIT_GET_CPU_TIMESTAMP_FROM_BLOCK,
          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
    }
 #endif