From ba6e2db70d3af315594ab00d3338f991bd468cbb Mon Sep 17 00:00:00 2001
From: Wenkai Du <wenkai.du@amd.com>
Date: Wed, 14 Sep 2022 00:10:19 +0000
Subject: [PATCH] Misc fixes and disable binTree

[ROCm/rccl commit: a06e14e39b96c0b802041ee080e085f1154b710a]
---
 projects/rccl/src/collectives/device/all_reduce.h | 3 ++-
 projects/rccl/src/include/comm.h                  | 1 +
 projects/rccl/src/include/devcomm.h               | 5 +++--
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/projects/rccl/src/collectives/device/all_reduce.h b/projects/rccl/src/collectives/device/all_reduce.h
index ea90cbd37f..cdf1eaba4c 100644
--- a/projects/rccl/src/collectives/device/all_reduce.h
+++ b/projects/rccl/src/collectives/device/all_reduce.h
@@ -376,7 +376,8 @@ namespace {
     const int nthreads = args->nWarps*WARP_SIZE;
     const int bid = args->bid;
     const int nChannels = args->nChannels;
-    ncclTree *tree = (args->pad_0 == 2) ? &ncclShmem->channel.binTree : &ncclShmem->channel.tree;
+    ncclTree *tree = &ncclShmem->channel.tree;
+    //ncclTree *tree = (args->pad_0 == 2) ? &ncclShmem->channel.binTree : &ncclShmem->channel.tree;
     ssize_t chunkSize = int(
       Proto::Id != NCCL_PROTO_LL ? args->lastChunkSize
                                  : Proto::calcBytePerStep()/sizeof(T));
diff --git a/projects/rccl/src/include/comm.h b/projects/rccl/src/include/comm.h
index 844eea5bcd..83a8ddeac9 100644
--- a/projects/rccl/src/include/comm.h
+++ b/projects/rccl/src/include/comm.h
@@ -107,6 +107,7 @@ struct ncclChannel {
   struct ncclRing ring;
   int* devRingUserRanks;
   struct ncclTree tree;
+  struct ncclTree binTree;
   struct ncclDirect collTree;
   int id; // index of this channel
   uint32_t workFifoSent; // last used work index+1
diff --git a/projects/rccl/src/include/devcomm.h b/projects/rccl/src/include/devcomm.h
index 3a8949460c..d47854b0b4 100644
--- a/projects/rccl/src/include/devcomm.h
+++ b/projects/rccl/src/include/devcomm.h
@@ -203,14 +203,15 @@ struct ncclWorkElem {
   union {
     uint8_t flagBits;
     struct {
-      uint8_t isUsed:1, redOpArgIsPtr:1, regUsed:1, pad_0:1, nWarps:4;
+      uint8_t isUsed:1, redOpArgIsPtr:1, regUsed:1, nWarps:5;
     };
   };
   uint8_t direct;
   uint8_t bid;
   uint8_t nChannels;
   struct {
-    uint32_t root:30;
+    uint32_t root:28;
+    uint32_t pad_0:2;
     uint32_t connIndex:2;
   };