Add Feature - Add NPKit Support in RCCL (#564)

* apply npkit * fix bug * add npkit in readme
2022-06-21 05:30:19 +08:00
parent f274c865c1
commit 6e93fafdc3
14 changed files with 1236 additions and 8 deletions
@@ -11,6 +11,9 @@
 #include "nccl.h"
 #include "rccl_bfloat16.h"
 #include "align.h"
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit_struct.h"
+#endif
 #include <stdint.h>
 // [RCCL] Support for clique-based kernels
 //#include "clique/CliqueCommon.h"
@@ -419,6 +422,11 @@ struct ncclDevComm {
  // Channels, device side
  struct ncclChannel* channels;

+#if defined(ENABLE_NPKIT)
+  NpKitEventCollectContext* npKitEventCollectContexts;
+  uint64_t* cpuTimestamp;
+#endif
+
 #ifdef ENABLE_PROFILING
  // Profiling counters
  struct ncclProf devProf;
@@ -0,0 +1,65 @@
+#ifndef NPKIT_H_
+#define NPKIT_H_
+
+#include <string>
+#include <thread>
+
+#include <hip/hip_runtime.h>
+
+#include "npkit/npkit_event.h"
+#include "npkit/npkit_struct.h"
+
+class NpKit {
+ public:
+  static const uint64_t kNumGpuEventBuffers = 512;
+
+  static const uint64_t kNumCpuEventBuffers = 32;
+
+  static ncclResult_t Init(int rank);
+
+  static ncclResult_t Dump(const std::string& dump_dir);
+
+  static ncclResult_t Shutdown();
+
+  static NpKitEventCollectContext* GetGpuEventCollectContexts();
+
+  static inline __device__ void CollectGpuEvent(uint8_t type, uint32_t size, uint32_t rsvd, uint64_t timestamp,
+                                                NpKitEventCollectContext* ctx) {
+    uint64_t event_buffer_head = ctx->event_buffer_head;
+    if (event_buffer_head < kMaxNumGpuEventsPerBuffer) {
+      NpKitEvent& event = ctx->event_buffer[event_buffer_head];
+      event.fields.type = type;
+      event.fields.size = size;
+      event.fields.rsvd = rsvd;
+      event.fields.timestamp = timestamp;
+      ctx->event_buffer_head++;
+    }
+  }
+
+  static void CollectCpuEvent(uint8_t type, uint32_t size, uint32_t rsvd, uint64_t timestamp, int channel_id);
+
+  static uint64_t* GetCpuTimestamp();
+
+ private:
+  static void CpuTimestampUpdateThread();
+
+  // 64K * 512 * 16B = 512MB per GPU
+  static const uint64_t kMaxNumGpuEventsPerBuffer = 1ULL << 16;
+
+  // 64K * 2 (send/recv) * (512/32) = 2M, 2M * 32 * 16B = 1GB per CPU
+  static const uint64_t kMaxNumCpuEventsPerBuffer = 1ULL << 21;
+
+  static NpKitEvent** gpu_event_buffers_;
+  static NpKitEvent** cpu_event_buffers_;
+
+  static NpKitEventCollectContext* gpu_collect_contexts_;
+  static NpKitEventCollectContext* cpu_collect_contexts_;
+  static uint64_t* cpu_timestamp_;
+
+  static uint64_t rank_;
+
+  static std::thread* cpu_timestamp_update_thread_;
+  static volatile bool cpu_timestamp_update_thread_should_stop_;
+};
+
+#endif
@@ -0,0 +1,98 @@
+#ifndef NPKIT_EVENT_H_
+#define NPKIT_EVENT_H_
+
+#define NPKIT_EVENT_INVALID                                     0x0
+
+#define NPKIT_EVENT_ALL_REDUCE_RING_ENTRY                       0x1
+#define NPKIT_EVENT_ALL_REDUCE_RING_EXIT                        0x2
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_ENTRY                0x3
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_EXIT                 0x4
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_ENTRY                 0x5
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_EXIT                  0x6
+
+#define NPKIT_EVENT_COPY_SEND_ENTRY                             0x7
+#define NPKIT_EVENT_COPY_SEND_EXIT                              0x8
+#define NPKIT_EVENT_DIRECT_COPY_SEND_ENTRY                      0x9
+#define NPKIT_EVENT_DIRECT_COPY_SEND_EXIT                       0xA
+#define NPKIT_EVENT_DIRECT_RECV_ENTRY                           0xB
+#define NPKIT_EVENT_DIRECT_RECV_EXIT                            0xC
+#define NPKIT_EVENT_DIRECT_RECV_COPY_SEND_ENTRY                 0xD
+#define NPKIT_EVENT_DIRECT_RECV_COPY_SEND_EXIT                  0xE
+#define NPKIT_EVENT_DIRECT_RECV_REDUCE_COPY_SEND_ENTRY          0xF
+#define NPKIT_EVENT_DIRECT_RECV_REDUCE_COPY_SEND_EXIT           0x10
+#define NPKIT_EVENT_DIRECT_SEND_ENTRY                           0x11
+#define NPKIT_EVENT_DIRECT_SEND_EXIT                            0x12
+#define NPKIT_EVENT_DIRECT_SEND_FROM_OUTPUT_ENTRY               0x13
+#define NPKIT_EVENT_DIRECT_SEND_FROM_OUTPUT_EXIT                0x14
+#define NPKIT_EVENT_RECV_ENTRY                                  0x15
+#define NPKIT_EVENT_RECV_EXIT                                   0x16
+#define NPKIT_EVENT_RECV_COPY_SEND_ENTRY                        0x17
+#define NPKIT_EVENT_RECV_COPY_SEND_EXIT                         0x18
+#define NPKIT_EVENT_RECV_REDUCE_COPY_ENTRY                      0x19
+#define NPKIT_EVENT_RECV_REDUCE_COPY_EXIT                       0x1A
+#define NPKIT_EVENT_RECV_REDUCE_COPY_SEND_ENTRY                 0x1B
+#define NPKIT_EVENT_RECV_REDUCE_COPY_SEND_EXIT                  0x1C
+#define NPKIT_EVENT_RECV_REDUCE_SEND_ENTRY                      0x1D
+#define NPKIT_EVENT_RECV_REDUCE_SEND_EXIT                       0x1E
+#define NPKIT_EVENT_SEND_ENTRY                                  0x1F
+#define NPKIT_EVENT_SEND_EXIT                                   0x20
+#define NPKIT_EVENT_SEND_FROM_OUTPUT_ENTRY                      0x21
+#define NPKIT_EVENT_SEND_FROM_OUTPUT_EXIT                       0x22
+
+#define NPKIT_EVENT_PRIM_SIMPLE_WAIT_PEER_ENTRY                 0x23
+#define NPKIT_EVENT_PRIM_SIMPLE_WAIT_PEER_EXIT                  0x24
+#define NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY      0x25
+#define NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT       0x26
+
+#define NPKIT_EVENT_PRIM_LL_WAIT_SEND_ENTRY                     0x27
+#define NPKIT_EVENT_PRIM_LL_WAIT_SEND_EXIT                      0x28
+#define NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY                  0x29
+#define NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT                   0x2A
+
+#define NPKIT_EVENT_PRIM_LL128_WAIT_SEND_ENTRY                  0x2B
+#define NPKIT_EVENT_PRIM_LL128_WAIT_SEND_EXIT                   0x2C
+#define NPKIT_EVENT_PRIM_LL128_DATA_PROCESS_ENTRY               0x2D
+#define NPKIT_EVENT_PRIM_LL128_DATA_PROCESS_EXIT                0x2E
+
+#define NPKIT_EVENT_NET_SEND_ENTRY                              0x2F
+#define NPKIT_EVENT_NET_SEND_EXIT                               0x30
+
+#define NPKIT_EVENT_NET_RECV_ENTRY                              0x31
+#define NPKIT_EVENT_NET_RECV_EXIT                               0x32
+
+#define NPKIT_EVENT_TIME_SYNC_GPU                               0x33
+#define NPKIT_EVENT_TIME_SYNC_CPU                               0x34
+
+#define NPKIT_EVENT_ALL_REDUCE_RING_SEND_ENTRY                  0x35
+#define NPKIT_EVENT_ALL_REDUCE_RING_SEND_EXIT                   0x36
+#define NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_ENTRY      0x37
+#define NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_EXIT       0x38
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_ENTRY  0x39
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_EXIT   0x3A
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_ENTRY 0x3B
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_EXIT  0x3C
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_ENTRY           0x3D
+#define NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_EXIT            0x3E
+
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_ENTRY         0x3F
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_EXIT          0x40
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_ENTRY      0x41
+#define NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_EXIT       0x42
+
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_ENTRY    0x43
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_EXIT     0x44
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_ENTRY          0x45
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_EXIT           0x46
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_ENTRY       0x47
+#define NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_EXIT        0x48
+
+#define NPKIT_EVENT_SEND_RECV_LOCAL_COPY_ENTRY                  0x49
+#define NPKIT_EVENT_SEND_RECV_LOCAL_COPY_EXIT                   0x4A
+#define NPKIT_EVENT_SEND_RECV_SEND_ENTRY                        0x4B
+#define NPKIT_EVENT_SEND_RECV_SEND_EXIT                         0x4C
+#define NPKIT_EVENT_SEND_RECV_RECV_ENTRY                        0x4D
+#define NPKIT_EVENT_SEND_RECV_RECV_EXIT                         0x4E
+
+#define NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME                    0x4F
+
+#endif
@@ -0,0 +1,25 @@
+#ifndef NPKIT_STRUCT_H_
+#define NPKIT_STRUCT_H_
+
+#include <cstdint>
+
+#pragma pack(push, 1)
+
+union NpKitEvent {
+  uint64_t bits[2];
+  struct {
+    uint64_t type : 8;
+    uint64_t size : 32;
+    uint64_t rsvd : 24;
+    uint64_t timestamp;
+  } fields;
+};
+
+struct NpKitEventCollectContext {
+  NpKitEvent* event_buffer;
+  uint64_t event_buffer_head;
+};
+
+#pragma pack(pop)
+
+#endif
@@ -58,6 +58,10 @@ struct ncclProxySubArgs {
  uint64_t end;
  void* requests[NCCL_STEPS];
  void* profilingEvents[NCCL_STEPS];
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_NET_SEND_ENTRY) && defined(ENABLE_NPKIT_EVENT_NET_SEND_EXIT)
+  int npKitSizesFifo[NCCL_STEPS];
+#endif
 };

 struct ncclProxyArgs {