From 77615cce287147d8f68e83d3b4a820930f3c66fc Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 6 Mar 2024 11:11:04 -0800
Subject: [PATCH] msccl: fix scratch memory allocation after API change (#1103)

---
 src/misc/msccl/msccl_setup.cc | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc
index b418dabcd5..f8ee6be403 100644
--- a/src/misc/msccl/msccl_setup.cc
+++ b/src/misc/msccl/msccl_setup.cc
@@ -445,7 +445,11 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
         }
         sizeRounded *= 2;
       }
-      NCCLCHECK(ncclCudaMalloc((char**)&scratchBuffer, sizeRounded, true));
+#if defined(HIP_UNCACHED_MEMORY)
+      NCCLCHECK(ncclCudaMalloc((char**)&scratchBuffer, sizeRounded, hipDeviceMallocUncached));
+#else
+      NCCLCHECK(ncclCudaMalloc((char**)&scratchBuffer, sizeRounded, hipDeviceMallocFinegrained));
+#endif
       work.scratchBuffer = status.scratchBuffers[sizeRounded] = scratchBuffer;
       INFO(NCCL_INIT, "MSCCL: Allocated scratch buffer of size %lu on request (%lu)", sizeRounded, sizeNeeded);
     } else {
@@ -533,7 +537,11 @@ error:
 
 ncclResult_t mscclInitWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus) {
   workFifoStatus->workFifoDepth = rcclParamMscclWorkFifoDepth();
-  NCCLCHECK(ncclCudaMalloc(&(workFifoStatus->workFifo), workFifoStatus->workFifoDepth, true));
+#if defined(HIP_UNCACHED_MEMORY)
+  NCCLCHECK(ncclCudaMalloc(&(workFifoStatus->workFifo), workFifoStatus->workFifoDepth, hipDeviceMallocUncached));
+#else
+  NCCLCHECK(ncclCudaMalloc(&(workFifoStatus->workFifo), workFifoStatus->workFifoDepth, hipDeviceMallocFinegrained));
+#endif
   NCCLCHECK(ncclCudaHostCalloc(&(workFifoStatus->workFifoDone), MSCCL_MAX_NUM_THREAD_BLOCKS));
   workFifoStatus->workFifoSent = 0;
   for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) {