From 225e871de2f7e8684d16a77a6273ae81c093aec6 Mon Sep 17 00:00:00 2001
From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com>
Date: Wed, 21 Jun 2023 08:26:40 -0600
Subject: [PATCH] fixed npkit size to never be a negative number (#779)

[ROCm/rccl commit: 9bdf6797a5d9a22d7674d778011beb8dc222dd5a]
---
 projects/rccl/src/include/npkit/npkit.h        | 6 +++---
 projects/rccl/src/include/npkit/npkit_struct.h | 2 +-
 projects/rccl/src/misc/npkit.cc                | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/projects/rccl/src/include/npkit/npkit.h b/projects/rccl/src/include/npkit/npkit.h
index 69304eb553..cfec52aa1b 100644
--- a/projects/rccl/src/include/npkit/npkit.h
+++ b/projects/rccl/src/include/npkit/npkit.h
@@ -30,20 +30,20 @@ class NpKit {
 
   static NpKitEventCollectContext* GetGpuEventCollectContexts();
 
-  static inline __device__ void CollectGpuEvent(uint8_t type, uint32_t size, uint32_t rsvd, uint64_t timestamp,
+  static inline __device__ void CollectGpuEvent(uint8_t type, int64_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.size = size < 0 ? 0 : 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 void CollectCpuEvent(uint8_t type, int64_t size, uint32_t rsvd, uint64_t timestamp, int channel_id);
 
   static uint64_t *GetCpuTimestamp();
 
diff --git a/projects/rccl/src/include/npkit/npkit_struct.h b/projects/rccl/src/include/npkit/npkit_struct.h
index a31a0900fd..cbcad2d8a7 100644
--- a/projects/rccl/src/include/npkit/npkit_struct.h
+++ b/projects/rccl/src/include/npkit/npkit_struct.h
@@ -14,7 +14,7 @@ union NpKitEvent {
   uint64_t bits[2];
   struct {
     uint64_t type : 8;
-    uint64_t size : 32;
+    uint32_t size : 32;
     uint64_t rsvd : 24;
     uint64_t timestamp;
   } fields;
diff --git a/projects/rccl/src/misc/npkit.cc b/projects/rccl/src/misc/npkit.cc
index b6c86e81a6..6ef83f0a40 100644
--- a/projects/rccl/src/misc/npkit.cc
+++ b/projects/rccl/src/misc/npkit.cc
@@ -160,12 +160,12 @@ NpKitEventCollectContext* NpKit::GetGpuEventCollectContexts() {
   return gpu_collect_contexts_;
 }
 
-void NpKit::CollectCpuEvent(uint8_t type, uint32_t size, uint32_t rsvd, uint64_t timestamp, int channel_id) {
+void NpKit::CollectCpuEvent(uint8_t type, int64_t size, uint32_t rsvd, uint64_t timestamp, int channel_id) {
   uint64_t event_buffer_head = cpu_collect_contexts_[channel_id].event_buffer_head;
   if (event_buffer_head < kMaxNumCpuEventsPerBuffer) {
     NpKitEvent& event = cpu_collect_contexts_[channel_id].event_buffer[event_buffer_head];
     event.fields.type = type;
-    event.fields.size = size;
+    event.fields.size = size < 0 ? 0 : size;
     event.fields.rsvd = rsvd;
     event.fields.timestamp = timestamp;
     cpu_collect_contexts_[channel_id].event_buffer_head++;