AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reduce NPKit latency overhead in MSCCL kernel (#893)

Browse Source 
* Reduce NPKit latency overhead in MSCCL kernel

* Fix build error without NPKit enable
This commit is contained in:
Wenkai Du
2023-09-15 13:28:26 -07:00
committed by GitHub
parent 16dd05a58a
commit 26e982d913
3 changed files with 39 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/include/npkit/npkit.h
+15
View File
@@ -13,9 +13,11 @@
#include "npkit/npkit_event.h"
#include "npkit/npkit_struct.h"
#include "common.h"
#define NPKIT_GET_GPU_TIMESTAMP wall_clock64
class NpKit {
public:
static const uint64_t kNumGpuEventBuffers = 512;
@@ -43,6 +45,19 @@ class NpKit {
}
}
static inline __device__ void CollectGpuEventLDS(uint8_t type, int64_t size, uint32_t rsvd, uint64_t timestamp) {
#if defined(ENABLE_NPKIT)
if (ncclShmem.event_buffer_head < LDS_NUM_EVENTS) {
NpKitEvent& event = ncclShmem.event_buffer[ncclShmem.event_buffer_head];
event.fields.type = type;
event.fields.size = size < 0 ? 0 : size;
event.fields.rsvd = rsvd;
event.fields.timestamp = timestamp;
ncclShmem.event_buffer_head++;
}
#endif
}
static void CollectCpuEvent(uint8_t type, int64_t size, uint32_t rsvd, uint64_t timestamp, int channel_id);
static uint64_t *GetCpuTimestamp();