From 478d8312b825eb6d54b5d87f620ef73bd5e1578d Mon Sep 17 00:00:00 2001
From: Ziyue Yang <ziyyang@microsoft.com>
Date: Sat, 6 Aug 2022 10:40:19 +0800
Subject: [PATCH] Improve alignment and tuning for Pivot A2A algorithm (#593)

* Improve alignment and tuning for Pivot A2A algorithm

* enable pivot a2a by default

[ROCm/rccl commit: f6b9686482f70509144715dcba2242ed188e23d5]
---
 projects/rccl/src/collectives/all_to_all_api.cc       | 7 ++++---
 projects/rccl/src/collectives/device/alltoall_pivot.h | 2 +-
 projects/rccl/src/init.cc                             | 2 +-
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/projects/rccl/src/collectives/all_to_all_api.cc b/projects/rccl/src/collectives/all_to_all_api.cc
index d5f83e1893..bbf503b61e 100644
--- a/projects/rccl/src/collectives/all_to_all_api.cc
+++ b/projects/rccl/src/collectives/all_to_all_api.cc
@@ -13,9 +13,11 @@ NCCL_API(ncclResult_t, ncclAllToAll, const void* sendbuff, void* recvbuff, size_
   ncclComm_t comm, hipStream_t stream);
 ncclResult_t ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
   ncclComm_t comm, hipStream_t stream) {
+  size_t rankOffset = count * ncclTypeSize(datatype);
+  size_t rankAlign = rankOffset & ((~rankOffset) + 1);
   // Determine Pivot A2A support now that we know number of channels
-  comm->topo->pivotA2AEnabled = comm->topo->pivotA2AEnabled && comm->nChannels >= comm->topo->pivotA2ANumBiRings * 2;
-  if (comm->topo->pivotA2AEnabled) {
+  if (comm->topo->pivotA2AEnabled && comm->nChannels >= comm->topo->pivotA2ANumBiRings * 2 &&
+      rankOffset >= 744 * 1024 && rankAlign != 4) {
     struct ncclInfo info = { ncclFuncAllToAllPivot, "AllToAllPivot",
       sendbuff, recvbuff, count, datatype, ncclSum, 0, comm, stream, /* Args */
       ALLTOALL_PIVOT_CHUNKSTEPS, ALLTOALL_PIVOT_SLICESTEPS };
@@ -23,7 +25,6 @@ ncclResult_t ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count, nc
   } else {
     int nRanks;
     NCCLCHECK(ncclCommCount(comm, &nRanks));
-    size_t rankOffset = count * ncclTypeSize(datatype);
     if (count == 0) return ncclSuccess;
     NCCLCHECK(ncclGroupStart());
     for (int r=0; r<nRanks; r++) {
diff --git a/projects/rccl/src/collectives/device/alltoall_pivot.h b/projects/rccl/src/collectives/device/alltoall_pivot.h
index 4b00ddbca8..8d5c0eb362 100644
--- a/projects/rccl/src/collectives/device/alltoall_pivot.h
+++ b/projects/rccl/src/collectives/device/alltoall_pivot.h
@@ -20,7 +20,7 @@ namespace {
     const int num_uni_rings = num_bi_rings * 2;
     const int num_chunks = args->nChannels / 2;
     const int chunk_id = (bid % num_bi_rings) + (bid / num_uni_rings * num_bi_rings);
-    const int elem_size = args->count % 256 ? 1 : 256;
+    const int elem_size = min(256, args->count & (~(args->count) + 1));
     const ssize_t num_elems = args->count / elem_size;
     const int num_padding_chunks = num_elems % num_chunks;
     const ssize_t chunk_offset = elem_size * (num_elems / num_chunks * chunk_id + (chunk_id < num_padding_chunks ? chunk_id : num_padding_chunks));
diff --git a/projects/rccl/src/init.cc b/projects/rccl/src/init.cc
index bc334d4a62..5a8f0bcf73 100644
--- a/projects/rccl/src/init.cc
+++ b/projects/rccl/src/init.cc
@@ -379,7 +379,7 @@ static ncclResult_t commFree(ncclComm_t comm) {
 
 RCCL_PARAM(CliqueIgnoreTopo, "CLIQUE_IGNORE_TOPO", 0);
 RCCL_PARAM(P2pNetDisable, "P2P_NET_DISABLE", 0);
-RCCL_PARAM(PivotAlltoallEnable, "PIVOT_ALLTOALL_ENABLE", 0);
+RCCL_PARAM(PivotAlltoallEnable, "PIVOT_ALLTOALL_ENABLE", 1);
 NCCL_PARAM(AggChannelSize, "AGG_CHANNEL_SIZE", -2);
 NCCL_PARAM(DisableGraphHelper, "GRAPH_HELPER_DISABLE", 0);
 NCCL_PARAM(GraphRegister, "GRAPH_REGISTER", 0);