From 9d8f68b4eebb5a54a6a5445f5038bc085e757359 Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Fri, 28 Jun 2024 08:00:22 -0700 Subject: [PATCH] NPKit: separate time stamps for GPU access from different blocks (#1229) To avoid races in memory access in GPU --- src/device/all_gather.h | 3 +-- src/device/all_reduce.h | 9 +++------ src/device/broadcast.h | 3 +-- src/device/msccl_kernel_impl.h | 3 +-- src/device/sendrecv.h | 6 ++---- src/include/npkit/npkit.h | 2 ++ src/misc/npkit.cc | 12 +++++++----- 7 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/device/all_gather.h b/src/device/all_gather.h index 051662e668..78193f91c8 100644 --- a/src/device/all_gather.h +++ b/src/device/all_gather.h @@ -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 diff --git a/src/device/all_reduce.h b/src/device/all_reduce.h index c33fedd392..39e7d3d62b 100644 --- a/src/device/all_reduce.h +++ b/src/device/all_reduce.h @@ -40,8 +40,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 @@ -238,8 +237,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 @@ -406,8 +404,7 @@ namespace { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU) if (isNpKitThread) { - 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 diff --git a/src/device/broadcast.h b/src/device/broadcast.h index f65b7fb2e7..07a64c0a14 100644 --- a/src/device/broadcast.h +++ b/src/device/broadcast.h @@ -33,8 +33,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 diff --git a/src/device/msccl_kernel_impl.h b/src/device/msccl_kernel_impl.h index 75168278ec..b3f4b29a6f 100644 --- a/src/device/msccl_kernel_impl.h +++ b/src/device/msccl_kernel_impl.h @@ -185,8 +185,7 @@ __device__ __forceinline__ void mscclRunInterpreter( #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU) if (tid == 0) { - uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp; - NpKit::CollectGpuEventLDS(NPKIT_EVENT_TIME_SYNC_CPU, 0, xcc_id, *cpuTimestamp); + NpKit::CollectGpuEventLDS(NPKIT_EVENT_TIME_SYNC_CPU, 0, xcc_id, NPKIT_GET_CPU_TIMESTAMP_FROM_BLOCK); } #endif diff --git a/src/device/sendrecv.h b/src/device/sendrecv.h index bf3a4785ee..fec708e436 100644 --- a/src/device/sendrecv.h +++ b/src/device/sendrecv.h @@ -26,8 +26,7 @@ struct RunWork { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU) if (isNpKitThread) { - 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 @@ -127,8 +126,7 @@ struct RunWork { #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU) if (isNpKitThread) { - 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 diff --git a/src/include/npkit/npkit.h b/src/include/npkit/npkit.h index 30c3a77ac6..3614f3626b 100644 --- a/src/include/npkit/npkit.h +++ b/src/include/npkit/npkit.h @@ -16,6 +16,8 @@ #include "common.h" #define NPKIT_GET_GPU_TIMESTAMP wall_clock64 +#define NPKIT_GET_CPU_TIMESTAMP_FROM_BLOCK \ + __atomic_load_n(reinterpret_cast((uint8_t *)ncclShmem.comm.cpuTimestamp + 128*blockIdx.x), __ATOMIC_RELAXED) class NpKit { diff --git a/src/misc/npkit.cc b/src/misc/npkit.cc index 0c72779645..d5e3245fd8 100644 --- a/src/misc/npkit.cc +++ b/src/misc/npkit.cc @@ -28,10 +28,12 @@ void NpKit::CpuTimestampUpdateThread() { uint64_t init_system_clock = std::chrono::system_clock::now().time_since_epoch().count(); uint64_t init_steady_clock = std::chrono::steady_clock::now().time_since_epoch().count(); uint64_t curr_steady_clock = 0; - volatile uint64_t* volatile_cpu_timestamp_ = cpu_timestamp_; while (!cpu_timestamp_update_thread_should_stop_) { - curr_steady_clock = std::chrono::steady_clock::now().time_since_epoch().count(); - *volatile_cpu_timestamp_ = init_system_clock + (curr_steady_clock - init_steady_clock); + for (int c = 0; c < MAXCHANNELS; c++) { + volatile uint64_t* volatile_cpu_timestamp_ = (volatile uint64_t*)((uint8_t *)cpu_timestamp_ + 128*c); + curr_steady_clock = std::chrono::steady_clock::now().time_since_epoch().count(); + __atomic_store_n(volatile_cpu_timestamp_, init_system_clock + (curr_steady_clock - init_steady_clock), __ATOMIC_RELAXED); + } } } @@ -58,8 +60,8 @@ ncclResult_t NpKit::Init(int rank) { cpu_collect_contexts_[i] = ctx; } - // Init timestamp - NCCLCHECK(ncclCudaHostCalloc(&cpu_timestamp_, 1)); + // Init timestamp. Allocates MAXCHANNELS*128 bytes buffer for GPU + NCCLCHECK(ncclCudaHostCalloc(&cpu_timestamp_, MAXCHANNELS*128/sizeof(cpu_timestamp_))); volatile uint64_t* volatile_cpu_timestamp = cpu_timestamp_; *volatile_cpu_timestamp = std::chrono::system_clock::now().time_since_epoch().count(); cpu_timestamp_update_thread_should_stop_ = false;