From da85abab543827156631c7d89ffbb81df40ed38d Mon Sep 17 00:00:00 2001
From: BertanDogancay <bertan.dogancay@gmail.com>
Date: Wed, 31 Jan 2024 17:09:07 -0800
Subject: [PATCH] Fix stack size

---
 src/device/common.cu | 10 +++++++---
 src/device/common.h  |  8 +-------
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/device/common.cu b/src/device/common.cu
index a939408344..174934cd22 100644
--- a/src/device/common.cu
+++ b/src/device/common.cu
@@ -17,13 +17,17 @@ struct RunWorkNop {
   __device__ void run(ncclWork *w) {}
 };
 
-__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
   ncclKernelMain<-1, RunWorkNop, false>(comm, channelMask, workHead);
 }
 #ifdef ENABLE_COLLTRACE
-__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
   ncclKernelMain<-1, RunWorkNop, true>(comm, channelMask, workHead);
 }
 #endif
 
-__device__ void ncclDevFunc_Nop() {}
+#ifdef USE_INDIRECT_FUNCTION_CALL
+__device__ void ncclDevFunc_Nop();
+#else
+__device__ __attribute__((noinline)) void ncclDevFunc_Nop();
+#endif
diff --git a/src/device/common.h b/src/device/common.h
index 2f65e54cfb..5ec92c696e 100644
--- a/src/device/common.h
+++ b/src/device/common.h
@@ -225,7 +225,7 @@ static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we
 }
 
 template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE>
-__device__ void ncclKernelMain(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+__forceinline__ __device__ void ncclKernelMain(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
   const int tid = threadIdx.x;
   int x = tid;
   switch (tid/WARP_SIZE) {
@@ -365,12 +365,6 @@ __global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channel
 __global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
 #endif
 
-#ifdef USE_INDIRECT_FUNCTION_CALL
-__device__ void ncclDevFunc_Nop();
-#else
-__device__ __attribute__((noinline)) void ncclDevFunc_Nop();
-#endif
-
 #ifdef ENABLE_COLLTRACE
 #define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
   __global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \