From 48859a97b1196ddea613b232bfd8a6ae599ecf5c Mon Sep 17 00:00:00 2001
From: Nusrat Islam <Nusrat.Islam@amd.com>
Date: Wed, 22 May 2024 13:37:15 -0500
Subject: [PATCH] device: update the logic for channelId assignment

---
 src/device/common.h | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/device/common.h b/src/device/common.h
index 093f7172ff..138a42ce09 100644
--- a/src/device/common.h
+++ b/src/device/common.h
@@ -235,23 +235,29 @@ __forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, struct
 
   switch (tid/WARP_SIZE) {
   case 0:
-	  ncclShmem.channelId = blockIdx.x;
-  /*for (int i = 0; i < num; i++) {
-    if (channelMask.masks[i] & (1ull<<x)) {
-      y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
-      y = total + y;
-      if (blockIdx.x == y) ncclShmem.channelId = x;
-    }
-    if (WARP_SIZE < 64) {
-      x = WARP_SIZE + tid;
+	//ncclShmem.channelId = blockIdx.x;
+    for (int i = 0; i < num; i++) {
       if (channelMask.masks[i] & (1ull<<x)) {
-	y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
-	y = y + total;
-        if (blockIdx.x == y) ncclShmem.channelId = x;
+        y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
+        y = total + y;
+        if (blockIdx.x == y) {
+          ncclShmem.channelId = y;
+	  break;
+        }
       }
+      if (WARP_SIZE < 64) {
+        x = WARP_SIZE + tid;
+        if (channelMask.masks[i] & (1ull<<x)) {
+	  y = __popcll(channelMask.masks[i] & ((1ull<<x)-1));
+	  y = y + total;
+          if (blockIdx.x == y) {
+	    ncclShmem.channelId = y;
+	    break;
+	  }
+        }
+      }
+      total = total + __popcll(channelMask.masks[i]);
     }
-    total = __popcll(channelMask.masks[i]);
-  }*/
     break;
   case 1:
     if (tid < WARP_SIZE + NCCL_MAX_GROUPS)