From 7987015a19c43c2b07d7dd2c32d52eb049eef85e Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 24 Jan 2024 08:43:50 -0800
Subject: [PATCH] Revert "msccl: build same number of kernels as in ROCm 5.7"
 (#1058)

This reverts commit f960174d03be7e5174baa83b256526d388a38842.
---
 cmake/Generator.cmake                      | 9 +++------
 src/collectives/device/msccl_kernel_impl.h | 4 +---
 src/include/msccl/msccl_kernel.h           | 4 +---
 src/misc/msccl/msccl_lifecycle.cc          | 5 ++---
 src/misc/msccl/msccl_setup.cc              | 6 ++----
 5 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/cmake/Generator.cmake b/cmake/Generator.cmake
index db7682bd93..782dfb0980 100644
--- a/cmake/Generator.cmake
+++ b/cmake/Generator.cmake
@@ -297,11 +297,9 @@ endfunction()
 ## Function to generate MSCCL Kernels
 ###########################################################################################################
 function(gen_msccl_kernels)
-  foreach(REDOP_CURRENT IN LISTS ALL_REDOPS)
+  set(MSCCL_REDOP Sum Prod Max Min)
+  foreach(REDOP_CURRENT IN LISTS MSCCL_REDOP)
     foreach(DATA_TYPE ${ALL_TYPES})
-      if (REDOP_CURRENT STREQUAL "SumPostDiv" AND DATA_TYPE IN_LIST FLOATS_LIST)
-        continue()  # Skip the iteration for FLOATS_LIST when REDOP_CURRENT is SumPostDiv
-      endif()
       set(FILE_NAME "${HIPIFY_DIR}/src/collectives/device/msccl_kernel_${REDOP_CURRENT}_${DATA_TYPE}.cpp")
       message(STATUS "Generating ${FILE_NAME}")
       file(WRITE ${FILE_NAME}
@@ -309,8 +307,7 @@ function(gen_msccl_kernels)
         #include \"primitives.h\"
         #include \"collectives.h\"
         #include \"devcomm.h\"
-        MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE}, false);
-        MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE}, true);")
+        MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE}, false);")
       list(APPEND HIP_SOURCES ${FILE_NAME})
     endforeach()
   endforeach()
diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h
index afddb8a1dc..11894644a6 100644
--- a/src/collectives/device/msccl_kernel_impl.h
+++ b/src/collectives/device/msccl_kernel_impl.h
@@ -428,8 +428,6 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, false) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, false) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Max, false) \
-  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, false) \
-  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(PreMulSum, false) \
-  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(SumPostDiv, false)
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, false)
 
 #endif
diff --git a/src/include/msccl/msccl_kernel.h b/src/include/msccl/msccl_kernel.h
index 4962fdcb69..3536a8eb94 100644
--- a/src/include/msccl/msccl_kernel.h
+++ b/src/include/msccl/msccl_kernel.h
@@ -40,9 +40,7 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto, fullOps)(struct n
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, false) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, false) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Max, false) \
-  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, false) \
-  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(PreMulSum, false) \
-  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(SumPostDiv, false)
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, false)
 
 MSCCL_DECL_KERNEL_ENTRY_FUNC()
 
diff --git a/src/misc/msccl/msccl_lifecycle.cc b/src/misc/msccl/msccl_lifecycle.cc
index 105d093716..034a597769 100644
--- a/src/misc/msccl/msccl_lifecycle.cc
+++ b/src/misc/msccl/msccl_lifecycle.cc
@@ -270,11 +270,10 @@ static ncclResult_t mscclInternalSchedulerSelectAlgo(struct mscclSchedulerParam*
   mscclStatus& status = mscclGetStatus();
   param->scheduled = false;
 
-
-  /*// Current MSCCL doesn't support pre/post op
+  // Current MSCCL doesn't support pre/post op
   if (param->op >= ncclAvg) {
     return ncclSuccess;
-  }*/
+  }
 
   // Whether the algorithm is in-place
   bool isInPlace = false;
diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc
index 48c6d60789..1fe62ac7c6 100644
--- a/src/misc/msccl/msccl_setup.cc
+++ b/src/misc/msccl/msccl_setup.cc
@@ -318,12 +318,10 @@ static ncclResult_t hostToDevRedOp(
   MSCCL_KERNEL_ENTRY_DEVREDOP(Sum, false), \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Prod, false), \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Max, false), \
-  MSCCL_KERNEL_ENTRY_DEVREDOP(Min, false), \
-  MSCCL_KERNEL_ENTRY_DEVREDOP(PreMulSum, false), \
-  MSCCL_KERNEL_ENTRY_DEVREDOP_NOFLOAT(SumPostDiv, false)
+  MSCCL_KERNEL_ENTRY_DEVREDOP(Min, false)
 
 // Except for ncclDevPreMulSum and ncclDevSumPostDiv required by ncclAvg
-void* mscclKernelEntries[ncclNumDevRedOps * ncclNumTypes * NCCL_NUM_PROTOCOLS] = {
+void* mscclKernelEntries[(ncclNumDevRedOps-2) * ncclNumTypes * NCCL_NUM_PROTOCOLS] = {
 #ifdef COMPILE_MSCCL_KERNEL
   MSCCL_KERNEL_ENTRY()
 #endif