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

To avoid races in memory access in GPU

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<uint64_t*>((uint8_t *)ncclShmem.comm.cpuTimestamp + 128*blockIdx.x), __ATOMIC_RELAXED)
 
 
 class NpKit {