From e7f27c66e0984c4bafbd94c865d5df4e3c094bb9 Mon Sep 17 00:00:00 2001
From: Pedram Alizadeh <pmohamma@amd.com>
Date: Tue, 29 Aug 2023 12:49:06 -0400
Subject: [PATCH] optimizing COLL_UNROLL for MI100 machines (#863)

---
 src/collectives/device/common.h   | 4 ++++
 src/collectives/device/sendrecv.h | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/src/collectives/device/common.h b/src/collectives/device/common.h
index cedb3bb855..49eea4ee6e 100644
--- a/src/collectives/device/common.h
+++ b/src/collectives/device/common.h
@@ -11,7 +11,11 @@
 #include "collectives.h"
 #include "devcomm.h"
 
+#if defined(__gfx908__)
+#define COLL_UNROLL 2
+#else
 #define COLL_UNROLL 4
+#endif
 
 #define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1)  // Using balanced tree instead of split tree
 
diff --git a/src/collectives/device/sendrecv.h b/src/collectives/device/sendrecv.h
index 5b6cf4b647..86c72f457a 100644
--- a/src/collectives/device/sendrecv.h
+++ b/src/collectives/device/sendrecv.h
@@ -204,6 +204,8 @@ struct RunWork<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
       } else {
 #if defined(__gfx90a__)
         runRecv<ProtoSimple<1,1,8>>(tid, nthreads, group, args);
+#elif defined(__gfx908__)
+        runRecv<ProtoSimple<1,1,4>>(tid, nthreads, group, args);
 #else
         runRecv<ProtoSimple<1,1>>(tid, nthreads, group, args);
 #endif
@@ -214,6 +216,8 @@ struct RunWork<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
       } else {
 #if defined(__gfx90a__)
         runSend<ProtoSimple<1,1,8>>(tid, nthreads, group, args);
+#elif defined(__gfx908__)
+        runSend<ProtoSimple<1,1,4>>(tid, nthreads, group, args);
 #else
         runSend<ProtoSimple<1,1>>(tid, nthreads, group, args);
 #endif