From 56ce9ef05f7f2e653eaa0c8ce7ef8298ac1f54b1 Mon Sep 17 00:00:00 2001
From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com>
Date: Tue, 28 Nov 2023 16:33:47 -0700
Subject: [PATCH] recreated pr 914 to work with current develop branch (#979)

---
 cmake/Dependencies.cmake                   |  2 +-
 src/collectives/device/msccl_kernel_impl.h | 12 +++++++++++-
 src/include/msccl/msccl_kernel.h           | 12 +++++++++++-
 src/misc/msccl/msccl_lifecycle.cc          |  5 +++--
 src/misc/msccl/msccl_setup.cc              | 18 ++++++++++++++++--
 5 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
index 5b4c76a25c..25862b3acd 100644
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@@ -88,7 +88,7 @@ set(DATATYPES_FLOAT
 function(expand_collectives FILE FUNC)
   set(REDOP Sum Prod Min Max PreMulSum SumPostDiv)
   if (FUNC STREQUAL "MscclKernel")
-    set(REDOP_FILTERED Sum Prod Min Max)
+    set(REDOP_FILTERED Sum Prod Min Max PreMulSum SumPostDiv)
   else()
     set(REDOP_FILTERED ${REDOP})
   endif()
diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h
index bbaf9d5dac..cd21e70965 100644
--- a/src/collectives/device/msccl_kernel_impl.h
+++ b/src/collectives/device/msccl_kernel_impl.h
@@ -504,6 +504,14 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, double, fullOps) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps)
 
+#define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(devredop) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int64_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint64_t, fullOps)
+
 #define MSCCL_IMPL_KERNEL_ENTRY_FUNC() \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, false) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, false) \
@@ -512,6 +520,8 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, true) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, true) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Max, true) \
-  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, true)
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, true) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(PreMulSum, true) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(SumPostDiv, true)
 
 #endif
diff --git a/src/include/msccl/msccl_kernel.h b/src/include/msccl/msccl_kernel.h
index 5e65bcf103..d8519b7db4 100644
--- a/src/include/msccl/msccl_kernel.h
+++ b/src/include/msccl/msccl_kernel.h
@@ -28,6 +28,14 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto, fullOps)(struct n
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, double, fullOps) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps)
 
+#define MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(devredop, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint8_t, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int32_t, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint32_t, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int64_t, fullOps) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint64_t, fullOps)
+
 #define MSCCL_DECL_KERNEL_ENTRY_FUNC() \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, false) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, false) \
@@ -36,7 +44,9 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto, fullOps)(struct n
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, true) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, true) \
   MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Max, true) \
-  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, true)
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(Min, true) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP(PreMulSum, true) \
+  MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(SumPostDiv, true)
 
 MSCCL_DECL_KERNEL_ENTRY_FUNC()
 
diff --git a/src/misc/msccl/msccl_lifecycle.cc b/src/misc/msccl/msccl_lifecycle.cc
index 01d605c945..7bf97f17b6 100644
--- a/src/misc/msccl/msccl_lifecycle.cc
+++ b/src/misc/msccl/msccl_lifecycle.cc
@@ -215,10 +215,11 @@ 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 606501af4e..b57b539d4b 100644
--- a/src/misc/msccl/msccl_setup.cc
+++ b/src/misc/msccl/msccl_setup.cc
@@ -309,6 +309,18 @@ static ncclResult_t hostToDevRedOp(
   MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, double, fullOps), \
   MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps)
 
+#define MSCCL_KERNEL_ENTRY_DEVREDOP_NOFLOAT(devredop, fullOps) \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, int8_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, uint8_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, int32_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, uint32_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, int64_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, uint64_t, fullOps), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_NULL()
+
 #define MSCCL_KERNEL_ENTRY() \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Sum, false), \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Prod, false), \
@@ -317,10 +329,12 @@ static ncclResult_t hostToDevRedOp(
   MSCCL_KERNEL_ENTRY_DEVREDOP(Sum, true), \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Prod, true), \
   MSCCL_KERNEL_ENTRY_DEVREDOP(Max, true), \
-  MSCCL_KERNEL_ENTRY_DEVREDOP(Min, true)
+  MSCCL_KERNEL_ENTRY_DEVREDOP(Min, true), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP(PreMulSum, true), \
+  MSCCL_KERNEL_ENTRY_DEVREDOP_NOFLOAT(SumPostDiv, true)
 
 // Except for ncclDevPreMulSum and ncclDevSumPostDiv required by ncclAvg
-void* mscclKernelEntries[(ncclNumDevRedOps - 2) * ncclNumTypes * NCCL_NUM_PROTOCOLS * 2] = {
+void* mscclKernelEntries[ncclNumDevRedOps * ncclNumTypes * NCCL_NUM_PROTOCOLS * 2] = {
 #ifdef COMPILE_MSCCL_KERNEL
   MSCCL_KERNEL_ENTRY()
 #endif