From b6ae2fd71d142625da4fb3180be25185c2d278ee Mon Sep 17 00:00:00 2001
From: Dingming Wu <dmwu0506@gmail.com>
Date: Fri, 10 Jan 2025 07:57:46 -0800
Subject: [PATCH] improving kernel traces on opCount bits and adding channelId
 in ncclCollTrace (#1485)

[ROCm/rccl commit: 69d0134ed2b6a57191594ea0966428423ccbeb62]
---
 projects/rccl/src/device/common.h  | 3 ++-
 projects/rccl/src/include/device.h | 5 +++--
 projects/rccl/src/init.cc          | 8 ++++----
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/projects/rccl/src/device/common.h b/projects/rccl/src/device/common.h
index 44e628b57c..390608a9d1 100644
--- a/projects/rccl/src/device/common.h
+++ b/projects/rccl/src/device/common.h
@@ -46,7 +46,8 @@
     uint32_t pos = atomicAdd(&ncclShmem.collTraceTail->tail, 1)%COLLTRACE_NUM_ITEMS; \
     struct ncclCollTrace* collTrace = ncclShmem.collTrace+pos; \
     collTrace->timeStamp = wall_clock64(); \
-    collTrace->bid = blockIdx.x;
+    collTrace->bid = blockIdx.x; \
+    collTrace->channelId = ncclShmem.channelId;
     // TODO: switch to atomicInc after llvm crash is fixed
     // uint32_t pos = atomicInc(&ncclShmem.collTraceTail->tail, COLLTRACE_NUM_ITEMS)
 
diff --git a/projects/rccl/src/include/device.h b/projects/rccl/src/include/device.h
index b30a4c23c4..983601e82f 100644
--- a/projects/rccl/src/include/device.h
+++ b/projects/rccl/src/include/device.h
@@ -298,7 +298,7 @@ struct ncclWorkElemP2p {
     };
   };
   uint8_t reg:1;
-  uint16_t opCount:12;
+  uint16_t opCount:15;
   // Important not to use any fields with greater than 4-byte alignment since
   // we need sizeof(ncclWorkElemP2p)==28, but that would be padded up to 32 if
   // there were 8-byte fields.
@@ -378,7 +378,8 @@ struct ncclCollTrace {
   uint8_t bid;
   int16_t funcIndex;
   uint32_t data_0;
-  uint64_t timeStamp;
+  uint8_t channelId;
+  uint64_t timeStamp:56;
   union {
     uint64_t opCount;
     uint32_t p2pOpCount[2];
diff --git a/projects/rccl/src/init.cc b/projects/rccl/src/init.cc
index 2819703a5e..9bec5637fa 100644
--- a/projects/rccl/src/init.cc
+++ b/projects/rccl/src/init.cc
@@ -250,14 +250,14 @@ void *ncclCommThreadMain(void *arg) {
         int offset = 0;
         uint16_t fIdx = td->funcIndex;
         if (type == ncclCollTraceDataType) {
-          sprintf(line, "## [%012.6f] [%02d:%02d] L:%04d DT %08x %016lx %016lx",
-            (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid,
+          sprintf(line, "## [%012.6f] [%02d:%02d:%02d] L:%04d DT %08x %016lx %016lx",
+            (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->channelId,
             fIdx, td->data_0, td->opCount, td->data_1);
         } else {
           if (type == ncclCollTraceP2pElemType)
-            sprintf(line, "## [%012.6f] [%02d:%02d] %06x-%06x", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->p2pOpCount[0], td->p2pOpCount[1]);
+            sprintf(line, "## [%012.6f] [%02d:%02d:%02d] %06x-%06x", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->channelId, td->p2pOpCount[0], td->p2pOpCount[1]);
           else
-            sprintf(line, "## [%012.6f] [%02d:%02d] %06lx", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->opCount);
+            sprintf(line, "## [%012.6f] [%02d:%02d:%02d] %06lx", (double)(td->timeStamp)/vega_gpu_rtc_freq, comm->rank, td->bid, td->channelId, td->opCount);
           offset = strlen(line);
           if (type == ncclCollTraceCollElemType) {
             sprintf(line+offset, " CE %s nw %d bi %d nc %d busId %lx nRanks %d", funcNames[fIdx], td->coll.nWarps, td->coll.bid, td->coll.nChannels, comm->busId, comm->nRanks);