From 9fa7a738dab6e5131960f9031246a7d2d68589bd Mon Sep 17 00:00:00 2001
From: isaki001 <36317038+isaki001@users.noreply.github.com>
Date: Wed, 10 Sep 2025 14:25:23 -0500
Subject: [PATCH] add reduce/broadcast algo/proto selection table for
 multi-node gfx940 (#1889)

[ROCm/rccl commit: 9c364393542b346bad2231573ac61b171cf4fa1d]
---
 projects/rccl/src/graph/tuning.cc           | 4 ++++
 projects/rccl/src/include/rccl_common.h     | 2 +-
 projects/rccl/src/rccl_wrap.cc              | 2 +-
 projects/rccl/tools/topo_expl/topo_expl.cpp | 4 +++-
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/projects/rccl/src/graph/tuning.cc b/projects/rccl/src/graph/tuning.cc
index 63dfef4485..e459507029 100644
--- a/projects/rccl/src/graph/tuning.cc
+++ b/projects/rccl/src/graph/tuning.cc
@@ -349,6 +349,10 @@ static struct tuningModel tuning_model_5 {
     {/*LL (min/max/factor/thread_threshold)*/ {0, 98304,  1, 16}, /*LL64/128 (min/max/factor/thread_threshold)*/ {98304, 5046272, 1, 64}},
     /*AllReduce*/
     {/*LL (min/max/factor/thread_threshold)*/ {0, 1048576, 1, 0},/*LL64/128 (min/max/factor/thread_threshold)*/ {1048576, 9437184, 3145728, 0}},
+    /*Reduce*/
+    {/*LL (min/max/factor/thread_threshold)*/ {0, 4096, 1, 0},/*LL64/128 (min/max/factor/thread_threshold)*/ {4096, 16777216, 1, 0}},
+    /*Broadcast*/
+    {/*LL (min/max/factor/thread_threshold)*/ {0, 8192, 1, 0},/*LL64/128 (min/max/factor/thread_threshold)*/ {8192, 33554432, 1, 0}},
   },
 };
 
diff --git a/projects/rccl/src/include/rccl_common.h b/projects/rccl/src/include/rccl_common.h
index 5c124b6c31..542149f6b0 100644
--- a/projects/rccl/src/include/rccl_common.h
+++ b/projects/rccl/src/include/rccl_common.h
@@ -79,7 +79,7 @@ inline size_t rcclGetSizePerRank(ncclFunc_t const& func, size_t const& nBytes, i
   // For AG, this is the send size per rank
   // For RS, this is the recv size per rank
   // For AR, this is the send/recv size per rank
-  return (func == ncclFuncReduceScatter || func == ncclFuncAllGather || func == ncclFuncBroadcast) ? nBytes / nRanks : nBytes;
+  return (func == ncclFuncReduceScatter || func == ncclFuncAllGather || func == ncclFuncBroadcast || func == ncclFuncReduce) ? nBytes / nRanks : nBytes;
 }
 ncclResult_t rcclGetAlgoProtoIndex(const char *envStr, const char* algoProtoString[], int nEntries, int& result);
 ncclResult_t rcclOverrideProtocol(const char* ncclProtoStr[], float table[][NCCL_NUM_PROTOCOLS], struct ncclTaskColl* info);
diff --git a/projects/rccl/src/rccl_wrap.cc b/projects/rccl/src/rccl_wrap.cc
index 83988dfe74..55f6e40dc3 100644
--- a/projects/rccl/src/rccl_wrap.cc
+++ b/projects/rccl/src/rccl_wrap.cc
@@ -41,7 +41,7 @@ void rcclUpdateCollectiveProtocol(struct ncclComm* comm, size_t const& nBytes, s
     userProtocolInput = !protoStr ? 0 : 1;
   }
 
-  if(!userProtocolInput && comm->nNodes >= 2 && (info->func == ncclFuncReduceScatter || info->func == ncclFuncAllGather || info->func == ncclFuncAllReduce || info->func == ncclFuncBroadcast)) {
+  if(!userProtocolInput && comm->nNodes >= 2 && (info->func == ncclFuncReduceScatter || info->func == ncclFuncAllGather || info->func == ncclFuncAllReduce || info->func == ncclFuncBroadcast || info->func == ncclFuncReduce)) {
     auto tunableIndex = rcclGetTunableIndex(info->func);
     auto llMin = comm->minMaxLLRange[tunableIndex][NCCL_PROTO_LL][RCCL_PROTOCOL_MIN_IDX];
     auto llMax = comm->minMaxLLRange[tunableIndex][NCCL_PROTO_LL][RCCL_PROTOCOL_MAX_IDX];
diff --git a/projects/rccl/tools/topo_expl/topo_expl.cpp b/projects/rccl/tools/topo_expl/topo_expl.cpp
index 20222c4a23..5a9d8cf151 100644
--- a/projects/rccl/tools/topo_expl/topo_expl.cpp
+++ b/projects/rccl/tools/topo_expl/topo_expl.cpp
@@ -291,7 +291,9 @@ int main(int argc,char* argv[])
   std::vector<ncclFunc_t> ncclFuncTypes = {
     ncclFuncAllReduce,
     ncclFuncReduceScatter,
-    ncclFuncAllGather
+    ncclFuncAllGather,
+    ncclFuncReduce,
+    ncclFuncBroadcast
   };
 
   std::cout << "Running fp32 production choices for algorithm/protocol/maxChannels" << std::endl;