Merge pull request #643 from ROCmSoftwarePlatform/2.15.5

Sync up with NCCL 2.15.5

[ROCm/rccl commit: b4f6eee9b4]

projects/rccl/CMakeLists.txt

@@ -100,6 +100,10 @@ list(APPEND CMAKE_PREFIX_PATH
 
 find_package(hip REQUIRED)
 message(STATUS "HIP compiler: ${HIP_COMPILER}")
+message(STATUS "HIP compiler: ${HIP_COMPILER} version ${HIP_CLANG_PATCH_LEVEL}")
+if (${HIP_CLANG_PATCH_LEVEL} LESS "22362")
+  message(FATAL_ERROR "RCCL requires ROCm 5.3 and above to compile")
+endif()
 message(STATUS "HIP runtime: ${HIP_RUNTIME}")
 
 if(BUILD_STATIC)
@@ -119,8 +123,8 @@ configure_file(src/nccl.h.in ${PROJECT_BINARY_DIR}/include/rccl/nccl.h)
 
 include_directories(${PROJECT_BINARY_DIR}/include) # for generated rccl.h header
 include_directories(${PROJECT_BINARY_DIR}/include/rccl) # for generated rccl.h header
+include_directories(${PROJECT_BINARY_DIR}/src/include) # for hipified header files
 include_directories(src)
-include_directories(src/include)
 include_directories(src/collectives)
 include_directories(src/collectives/device)
 
@@ -153,6 +157,96 @@ foreach(filename ${CU_SOURCES})
   list(APPEND CPP_SOURCES ${cpp_filename})
 endforeach(filename)
 
+set(HEADER_SOURCES
+  src/include/collectives.h
+  src/include/align.h
+  src/include/profiler.h
+  src/include/alloc.h
+  src/include/ibvwrap.h
+  src/include/gdrwrap.h
+  src/include/utils.h
+  src/include/strongstream.h
+  src/include/comm.h
+  src/include/trees.h
+  src/include/rccl_vars.h
+  src/include/checks.h
+  src/include/p2p.h
+  src/include/timer.h
+  src/include/coll_net.h
+  src/include/signals.h
+  src/include/proxy.h
+  src/include/net.h
+  src/include/devcomm.h
+  src/include/enqueue.h
+  src/include/debug.h
+  src/include/argcheck.h
+  src/include/rocm_smi_wrap.h
+  src/include/bootstrap.h
+  src/include/BfdBacktrace.hpp
+  src/include/nccl_net.h
+  src/include/cudawrap.h
+  src/include/rccl_bfloat16.h
+  src/include/shm.h
+  src/include/transport.h
+  src/include/group.h
+  src/include/socket.h
+  src/include/cpuset.h
+  src/include/rocmwrap.h
+  src/include/graph.h
+  src/include/nvmlwrap.h
+  src/include/param.h
+  src/include/channel.h
+  src/include/nvtx_stub.h
+  src/include/core.h
+  src/include/info.h
+  src/include/git_version.h
+  src/include/npkit/npkit_event.h
+  src/include/npkit/npkit.h
+  src/include/npkit/npkit_struct.h
+  src/include/nvtx3/nvtxDetail/nvtxImplCudaRt_v3.h
+  src/include/nvtx3/nvtxDetail/nvtxTypes.h
+  src/include/nvtx3/nvtxDetail/nvtxImpl.h
+  src/include/nvtx3/nvtxDetail/nvtxImplSync_v3.h
+  src/include/nvtx3/nvtxDetail/nvtxInitDecls.h
+  src/include/nvtx3/nvtxDetail/nvtxLinkOnce.h
+  src/include/nvtx3/nvtxDetail/nvtxImplCore.h
+  src/include/nvtx3/nvtxDetail/nvtxInitDefs.h
+  src/include/nvtx3/nvtxDetail/nvtxImplCuda_v3.h
+  src/include/nvtx3/nvtxDetail/nvtxInit.h
+  src/include/nvtx3/nvtxDetail/nvtxImplOpenCL_v3.h
+  src/include/nvtx3/nvToolsExtSync.h
+  src/include/nvtx3/nvToolsExtCudaRt.h
+  src/include/nvtx3/nvToolsExtCuda.h
+  src/include/nvtx3/nvToolsExtOpenCL.h
+  src/graph/rings.h
+  src/graph/rome_models.h
+  src/graph/topo.h
+  src/graph/xml.h)
+foreach(filename ${HEADER_SOURCES})
+  configure_file(${PROJECT_SOURCE_DIR}/${filename} ${filename} COPYONLY)
+endforeach(filename)
+
+set(API_SOURCES
+    src/collectives/all_reduce.cc
+    src/collectives/all_gather.cc
+    src/collectives/all_to_all.cc
+    src/collectives/all_to_allv.cc
+    src/collectives/reduce.cc
+    src/collectives/broadcast.cc
+    src/collectives/reduce_scatter.cc
+    src/collectives/scatter.cc
+    src/collectives/gather.cc
+    src/collectives/sendrecv.cc
+    src/net.cc)
+foreach(filename ${API_SOURCES})
+  string(REPLACE ".cc"
+                 "_api.cpp"
+                 cpp_filename
+                 ${filename})
+  configure_file(${filename} ${cpp_filename} COPYONLY)
+  list(APPEND CPP_SOURCES ${cpp_filename})
+endforeach(filename)
+
 set(CC_SOURCES
     src/init.cc
     src/graph/trees.cc
@@ -164,16 +258,6 @@ set(CC_SOURCES
     src/graph/topo.cc
     src/graph/xml.cc
     src/graph/rome_models.cc
-    src/collectives/all_reduce_api.cc
-    src/collectives/all_gather_api.cc
-    src/collectives/reduce_api.cc
-    src/collectives/broadcast_api.cc
-    src/collectives/reduce_scatter_api.cc
-    src/collectives/sendrecv_api.cc
-    src/collectives/gather_api.cc
-    src/collectives/scatter_api.cc
-    src/collectives/all_to_all_api.cc
-    src/collectives/all_to_allv_api.cc
     src/channel.cc
     src/misc/argcheck.cc
     src/misc/nvmlwrap_stub.cc
@@ -200,16 +284,60 @@ set(CC_SOURCES
     src/group.cc
     src/bootstrap.cc
     src/proxy.cc
-    src/net.cc
-    src/enqueue.cc
-    ${CMAKE_CURRENT_BINARY_DIR}/git_version.cpp)
-
+    src/enqueue.cc)
 foreach(filename ${CC_SOURCES})
-  list(APPEND CPP_SOURCES ${filename})
+  string(REPLACE ".cc"
+                 ".cpp"
+                 cpp_filename
+                 ${filename})
+  configure_file(${filename} ${cpp_filename} COPYONLY)
+  list(APPEND CPP_SOURCES ${cpp_filename})
 endforeach(filename)
 
+list(APPEND CPP_SOURCES  ${CMAKE_CURRENT_BINARY_DIR}/git_version.cpp)
 add_library(rccl ${CPP_SOURCES})
 
+message ("-- Hipifying source")
+set(HIPIFY_SOURCES
+    src/collectives/all_gather_api.cpp
+    src/collectives/all_reduce_api.cpp
+    src/collectives/all_to_all_api.cpp
+    src/collectives/all_to_allv_api.cpp
+    src/collectives/broadcast_api.cpp
+    src/collectives/gather_api.cpp
+    src/collectives/reduce_api.cpp
+    src/collectives/reduce_scatter_api.cpp
+    src/collectives/scatter_api.cpp
+    src/collectives/sendrecv_api.cpp
+    src/debug.cpp
+    src/enqueue.cpp
+    src/graph/xml.cpp
+    src/group.cpp
+    src/include/alloc.h
+    src/include/checks.h
+    src/include/info.h
+    src/include/proxy.h
+    src/include/strongstream.h
+    src/init.cpp
+    src/misc/argcheck.cpp
+    src/misc/shmutils.cpp
+    src/misc/strongstream.cpp
+    src/misc/utils.cpp
+    src/net_api.cpp
+    src/proxy.cpp
+    src/transport.cpp
+    src/transport/coll_net.cpp
+    src/transport/net.cpp
+    src/transport/net_socket.cpp
+    src/transport/p2p.cpp
+    src/transport/shm.cpp)
+find_program( hipify-perl_executable hipify-perl )
+foreach(filename ${HIPIFY_SOURCES})
+  message ("   ${filename}")
+  execute_process(COMMAND bash "-c" "${hipify-perl_executable} -inplace -quiet-warnings ${PROJECT_BINARY_DIR}/${filename}" OUTPUT_VARIABLE HIPIFY_OUTPUT ERROR_VARIABLE HIPIFY_OUTPUT)
+endforeach(filename)
+message ("-- Hipifying source - done")
+
 # Create a custom target that creates/updates git_version.cpp
 # that executes whenever rccl is built
 add_custom_target(git_version_check

projects/rccl/makefiles/common.mk

@@ -31,13 +31,17 @@ CUDA8_GENCODE = -gencode=arch=compute_35,code=sm_35 \
                 -gencode=arch=compute_61,code=sm_61
 CUDA9_GENCODE = -gencode=arch=compute_70,code=sm_70
 CUDA11_GENCODE = -gencode=arch=compute_80,code=sm_80
+CUDA11_8_GENCODE = -gencode=arch=compute_90,code=sm_90
 
 CUDA8_PTX     = -gencode=arch=compute_61,code=compute_61
 CUDA9_PTX     = -gencode=arch=compute_70,code=compute_70
 CUDA11_PTX    = -gencode=arch=compute_80,code=compute_80
+CUDA11_8_PTX  = -gencode=arch=compute_90,code=compute_90
 
-# Include Ampere support if we're using CUDA11 or above
-ifeq ($(shell test "0$(CUDA_MAJOR)" -ge 11; echo $$?),0)
+ifeq ($(shell test "0$(CUDA_MAJOR)" -eq 11 -a "0$(CUDA_MINOR)" -ge 8 -o "0$(CUDA_MAJOR)" -gt 11; echo $$?),0)
+# Include Hopper support if we're using CUDA11.8 or above
+  NVCC_GENCODE ?= $(CUDA8_GENCODE) $(CUDA9_GENCODE) $(CUDA11_GENCODE) $(CUDA11_8_GENCODE) $(CUDA11_8_PTX)
+else ifeq ($(shell test "0$(CUDA_MAJOR)" -ge 11; echo $$?),0)
   NVCC_GENCODE ?= $(CUDA8_GENCODE) $(CUDA9_GENCODE) $(CUDA11_GENCODE) $(CUDA11_PTX)
 # Include Volta support if we're using CUDA9 or above
 else ifeq ($(shell test "0$(CUDA_MAJOR)" -ge 9; echo $$?),0)
@@ -45,7 +49,7 @@ else ifeq ($(shell test "0$(CUDA_MAJOR)" -ge 9; echo $$?),0)
 else
   NVCC_GENCODE ?= $(CUDA8_GENCODE) $(CUDA8_PTX)
 endif
-#$(info NVCC_GENCODE is ${NVCC_GENCODE})
+$(info NVCC_GENCODE is ${NVCC_GENCODE})
 
 CXXFLAGS   := -DCUDA_MAJOR=$(CUDA_MAJOR) -DCUDA_MINOR=$(CUDA_MINOR) -fPIC -fvisibility=hidden \
               -Wall -Wno-unused-function -Wno-sign-compare -std=c++11 -Wvla \

projects/rccl/makefiles/version.mk

@@ -1,6 +1,6 @@
 ##### version
 NCCL_MAJOR   := 2
-NCCL_MINOR   := 14
-NCCL_PATCH   := 3
+NCCL_MINOR   := 15
+NCCL_PATCH   := 5
 NCCL_SUFFIX  :=
 PKG_REVISION := 1

projects/rccl/src/channel.cc

@@ -20,14 +20,14 @@ ncclResult_t initChannel(struct ncclComm* comm, int channelId) {
 
   // The extra on nRanks+1 is for collnet root (i.e. network)
   channel->peers = ncclMemoryStackAlloc<struct ncclChannelPeer>(&comm->memPermanent, nRanks+1);
-  NCCLCHECK(ncclCudaCallocAsync(&channel->devPeers, nRanks+1, comm->deviceStream.stream));
+  NCCLCHECK(ncclCudaCallocAsync(&channel->devPeers, nRanks+1, comm->deviceStream.cudaStream));
   ncclCommPushCudaFree(comm, channel->devPeers);
 
   channel->ring.userRanks = ncclMemoryStackAlloc<int>(&comm->memPermanent, nRanks);
-  NCCLCHECK(ncclCudaCallocAsync(&channel->devRingUserRanks, nRanks, comm->deviceStream.stream));
+  NCCLCHECK(ncclCudaCallocAsync(&channel->devRingUserRanks, nRanks, comm->deviceStream.cudaStream));
   ncclCommPushCudaFree(comm, channel->devRingUserRanks);
 
-  NCCLCHECK(ncclStrongStreamRelease(ncclCudaGraphNull(), &comm->deviceStream));
+  NCCLCHECK(ncclStrongStreamRelease(ncclCudaGraphNone(), &comm->deviceStream));
 
   for (int r=0; r < nRanks+1; ++r) {
     for (int b=0; b < NCCL_MAX_CONNS; b++) {

projects/rccl/src/collectives/all_gather.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,9 +8,9 @@
 #include "collectives.h"
 
 NCCL_API(ncclResult_t, ncclAllGather, const void* sendbuff, void* recvbuff, size_t sendcount,
-    ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream);
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount,
-    ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream) {
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncAllGather, "AllGather",
     sendbuff, recvbuff, sendcount, datatype, ncclSum, 0, comm, stream, /* Args */

projects/rccl/src/collectives/all_reduce.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -8,9 +7,9 @@
 #include "enqueue.h"
 
 NCCL_API(ncclResult_t, ncclAllReduce, const void* sendbuff, void* recvbuff, size_t count,
-    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream);
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream);
 ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count,
-    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream) {
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncAllReduce, "AllReduce",
     sendbuff, recvbuff, count, datatype, op, 0, comm, stream, /* Args */

projects/rccl/src/collectives/broadcast.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,9 +8,9 @@
 #include "collectives.h"
 
 NCCL_API(ncclResult_t, ncclBroadcast, const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root,
-    ncclComm_t comm, hipStream_t stream);
+    ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root,
-    ncclComm_t comm, hipStream_t stream) {
+    ncclComm_t comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncBroadcast, "Broadcast",
     sendbuff, recvbuff, count, datatype, ncclSum, root, comm, stream, /* Args */
@@ -20,9 +19,9 @@ ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, n
 }
 /* Deprecated original "in place" function, similar to MPI */
 NCCL_API(ncclResult_t, ncclBcast, void* buff, size_t count, ncclDataType_t datatype, int root,
-    ncclComm_t comm, hipStream_t stream);
+    ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclBcast(void* buff, size_t count, ncclDataType_t datatype, int root,
-    ncclComm_t comm, hipStream_t stream) {
+    ncclComm_t comm, cudaStream_t stream) {
   return ncclBroadcast(buff, buff, count, datatype, root, comm, stream);
 }
 

projects/rccl/src/collectives/reduce.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,9 +8,9 @@
 #include "collectives.h"
 
 NCCL_API(ncclResult_t, ncclReduce, const void* sendbuff, void* recvbuff, size_t count,
-    ncclDataType_t datatype, ncclRedOp_t op, int root, ncclComm_t comm, hipStream_t stream);
+    ncclDataType_t datatype, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclReduce(const void* sendbuff, void* recvbuff, size_t count,
-    ncclDataType_t datatype, ncclRedOp_t op, int root, ncclComm_t comm, hipStream_t stream) {
+    ncclDataType_t datatype, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncReduce, "Reduce",
     sendbuff, recvbuff, count, datatype, op, root, comm, stream, /* Args */

projects/rccl/src/collectives/reduce_scatter.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,9 +8,9 @@
 #include "collectives.h"
 
 NCCL_API(ncclResult_t, ncclReduceScatter, const void* sendbuff, void* recvbuff, size_t recvcount,
-    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream);
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream);
 ncclResult_t ncclReduceScatter(const void* sendbuff, void* recvbuff, size_t recvcount,
-    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, hipStream_t stream) {
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm* comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncReduceScatter, "ReduceScatter",
     sendbuff, recvbuff, recvcount, datatype, op, 0, comm, stream, /* Args */

projects/rccl/src/collectives/sendrecv.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -10,9 +9,9 @@
 #include "argcheck.h" // Need some checks here since we access comm
 
 NCCL_API(ncclResult_t, ncclSend, const void* sendbuff, size_t count, ncclDataType_t datatype, int peer,
-    ncclComm_t comm, hipStream_t stream);
+    ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer,
-    ncclComm_t comm, hipStream_t stream) {
+    ncclComm_t comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncSend, "Send",
     NULL, (void*)sendbuff, count, datatype, ncclSum, peer, comm, stream, /* Args */
@@ -25,9 +24,9 @@ ncclResult_t ncclSend(const void* sendbuff, size_t count, ncclDataType_t datatyp
 }
 
 NCCL_API(ncclResult_t, ncclRecv, void* recvbuff, size_t count, ncclDataType_t datatype, int peer,
-    ncclComm_t comm, hipStream_t stream);
+    ncclComm_t comm, cudaStream_t stream);
 ncclResult_t ncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer,
-    ncclComm_t comm, hipStream_t stream) {
+    ncclComm_t comm, cudaStream_t stream) {
   NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclInfo info = { ncclFuncRecv, "Recv",
     NULL, recvbuff, count, datatype, ncclSum, peer, comm, stream, /* Args */

projects/rccl/src/debug.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -160,7 +159,7 @@ void ncclDebugLog(ncclDebugLogLevel level, unsigned long flags, const char *file
 
   int cudaDev;
   if (!(level == NCCL_LOG_TRACE && flags == NCCL_CALL)) {
-    hipGetDevice(&cudaDev);
+    cudaGetDevice(&cudaDev);
   }
 
   char buffer[1024];

projects/rccl/src/enqueue.cc

@@ -42,11 +42,11 @@ static ncclResult_t computeColl(struct ncclInfo* info /* input */, int* workFunc
 size_t ncclKernMaxLocalSize() {
   ncclResult_t res = ncclSuccess;
   int numNcclKerns = sizeof(ncclKerns)/sizeof(ncclKerns[0]);
-  hipFuncAttributes attr = {0};
+  cudaFuncAttributes attr = {0};
   size_t max = 0;
   for (int i = 0; i < numNcclKerns; i++) {
     if (ncclKerns[i].kernelFn != nullptr) {
-      CUDACHECKGOTO(hipFuncGetAttributes(&attr, reinterpret_cast<const void*>(ncclKerns[i].kernelFn)), res, error);
+      CUDACHECKGOTO(cudaFuncGetAttributes(&attr, reinterpret_cast<const void*>(ncclKerns[i].kernelFn)), res, error);
       if (attr.localSizeBytes > max) max = attr.localSizeBytes;
     }
   }
@@ -59,9 +59,9 @@ error:
 size_t ncclKernLocalSize(int i) {
   ncclResult_t res = ncclSuccess;
   int numNcclKerns = sizeof(ncclKerns)/sizeof(ncclKerns[0]);
-  hipFuncAttributes attr = {0};
+  cudaFuncAttributes attr = {0};
   if (i < numNcclKerns)
-    CUDACHECKGOTO(hipFuncGetAttributes(&attr, (const void*)(ncclKerns[i].kernelFn)), res, error);
+    CUDACHECKGOTO(cudaFuncGetAttributes(&attr, reinterpret_cast<const void*>(ncclKerns[i].kernelFn)), res, error);
 
 error:
   return (res != ncclSuccess) ? 0 : attr.localSizeBytes;
@@ -73,7 +73,7 @@ ncclResult_t ncclKernSetSharedMemoryCarveout(int carveOut) {
   ncclResult_t res = ncclSuccess;
   int numNcclKerns = sizeof(ncclKerns)/sizeof(ncclKerns[0]);
   for (int i = 0; i < numNcclKerns; i++) {
-    CUDACHECKGOTO(hipFuncSetAttribute((const void *)ncclKerns[i].kernelFn, hipFuncAttributePreferredSharedMemoryCarveout, carveOut), res, error);
+    CUDACHECKGOTO(cudaFuncSetAttribute((const void *)ncclKerns[i].kernelFn, cudaFuncAttributePreferredSharedMemoryCarveout, carveOut), res, error);
   }
 
 error:
@@ -305,7 +305,7 @@ static ncclResult_t addP2pToPlan(
   struct ncclInfo info = {
     isSendNotRecv ? ncclFuncSend : ncclFuncRecv,
     isSendNotRecv ? "Send" : "Recv",
-    nullptr, addr, bytes, ncclInt8, ncclSum, peer, comm, (hipStream_t)0,
+    nullptr, addr, bytes, ncclInt8, ncclSum, peer, comm, (cudaStream_t)0,
     /*Args*/1, 1
   };
 
@@ -366,7 +366,7 @@ static void finishPlan(struct ncclKernelPlan* plan) {
   plan->channelCount = channelCount;
   plan->channelMask = channelMask;
   plan->hasProxyOps = hasProxyOps;
-  plan->threadPerBlock = std::max(plan->threadPerBlock, 4*WARP_SIZE);
+  plan->threadPerBlock = std::max(plan->threadPerBlock, 3*plan->comm->WarpSize);
 }
 
 static ncclResult_t registerIntraNodeBuffers(
@@ -831,6 +831,7 @@ static ncclResult_t hostStreamPlanTask(struct ncclComm* comm, struct ncclKernelP
 }
 
 static void HIPRT_CB hostStreamPlanCallback(void *plan_) {
+  NVTX3_FUNC_RANGE_IN(nccl_domain);
   struct ncclKernelPlan* plan = (struct ncclKernelPlan*)plan_;
   ncclResult_t result = hostStreamPlanTask(plan->comm, plan);
   if (result != ncclSuccess) {
@@ -845,7 +846,7 @@ static ncclResult_t reclaimPlan(struct ncclComm* comm, struct ncclCommCallback*
     NCCLCHECK(ncclCudaFree(plan->workHead));
     while (!ncclIntruQueueEmpty(&plan->ipcMemQueue)) {
       struct ncclPointerList* q = ncclIntruQueueDequeue(&plan->ipcMemQueue);
-      CUDACHECKIGNORE(hipIpcCloseMemHandle(q->ptr));
+      CUDACHECKIGNORE(cudaIpcCloseMemHandle(q->ptr));
       ncclMemoryPoolFree(&comm->memPool_ncclPointerList, q);
     }
   }
@@ -918,19 +919,38 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
     struct ncclKernelPlan* planHead = ncclIntruQueueHead(&comm->planQueue);
     comm->unlaunchedPlansHead = planHead;
 
+    // Semantically we want these dependencies for the kernels launched:
+    //   1. Launch host task on hostStream.
+    //   2. Launch kernel, depends on all of {deviceStream, hostStream, userStream[i]...}
+    //   3. {deviceStream, userStream[i]...} depend on kernel.
+    // We achieve this by:
+    //   1. userStream[0] waits on deviceStream
+    //   2. deviceStream waits on each of userStream[1...]
+    //   3. host task launch on hostStream
+    //   4. userStream[0] waits on hostStream
+    //   5. kernel launch on userStream[0]
+    //   6. deviceStream waits on userStream[0]
+    //   7. userStream[1...] each waits on deviceStream
+    // The two-level fan-in fan-out is because ncclStrongStreamWaitStream() requires
+    // at least one of the two streams to be strong-stream.
+    cudaStream_t launchStream = tasks->streams->stream;
     NCCLCHECKGOTO(ncclStrongStreamAcquire(tasks->capturingGraph, &comm->deviceStream), result, failure);
 
-    // Create dependency for nccl device work on user streams.
-    for (struct ncclCudaStreamList* l=tasks->streams; l != nullptr && tasks->numStreams != 1; l = l->next) {
-      NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->deviceStream, l->stream), result, failure);
-    }
-
-    if (tasks->numStreams == 1 && tasks->streams->stream != comm->lastStream) {
+    if (tasks->numStreams != 1) {
+      // Create dependency for device stream on user streams. First from extra user
+      // streams to deviceStream. Then deviceStream to first user stream.
+      for (struct ncclCudaStreamList* l=tasks->streams->next; l != nullptr; l = l->next) {
+        NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->deviceStream, l->stream), result, failure);
+      }
+      NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, launchStream, &comm->deviceStream), result, failure);
+    } else if (tasks->streams->stream != comm->lastStream) {
       // Stream changed from last call, create dependency against last NCCL kernel launch
       CUDACHECK(hipStreamWaitEvent(tasks->streams->stream, comm->doneEvent, 0));
     }
 
     if (persistent || comm->persistentRefs != 0) {
+      // We have to launch host tasks to push proxy args. We are careful to only
+      // do this if necessary since host tasks impose a high performance cost in CUDA.
       bool acquired = false;
       for (struct ncclKernelPlan* plan=planHead; plan != nullptr; plan = plan->next) {
         if (plan->hasProxyOps) {
@@ -942,6 +962,8 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
         }
       }
       if (acquired) {
+        // Make to-be-launched kernels dependent on just-launched host stream tasks.
+        if (tasks->numStreams != 1) NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, launchStream, &comm->hostStream), result, failure);
         NCCLCHECKGOTO(ncclStrongStreamRelease(tasks->capturingGraph, &comm->hostStream), result, failure);
       }
     }
@@ -967,8 +989,15 @@ ncclResult_t ncclLaunchKernelBefore_NoUncapturedCuda(struct ncclComm* comm, stru
   return ncclSuccess;
 }
 
+#if CUDART_VERSION >= 11080
+#define NCCL_MAX_CGA_CLUSTER_SIZE 8
+NCCL_PARAM(CGAClusterSize, "CGA_CLUSTER_SIZE", 0);
+#endif
+
 ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan) {
   struct ncclTasks* tasks = &comm->tasks;
+  void *fn = plan->kernelFn;
+  hipStream_t launchStream = tasks->streams->stream;
   dim3 grid = {(unsigned)plan->channelCount, 1, 1};
   dim3 block = {(unsigned)plan->threadPerBlock, 1, 1};
   void *args[3] = {&comm->devComm, &plan->channelMask, &plan->workHead};
@@ -976,9 +1005,54 @@ ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan
     CUDACHECK(hipExtLaunchKernel(plan->kernelFn, grid, block, args, 0, tasks->streams->stream, NULL, comm->doneEvent, 0));
     comm->lastStream = tasks->streams->stream;
   } else {
-    NCCLCHECK(ncclStrongStreamLaunchKernel(
-      tasks->capturingGraph, &comm->deviceStream, plan->kernelFn, grid, block, args, 0
-    ));
+    #if CUDART_VERSION >= 11080
+    int driverVersion;
+    NCCLCHECK(ncclCudaDriverVersion(&driverVersion));
+
+    unsigned int clusterSize = 0;
+    clusterSize = ncclParamCGAClusterSize();
+    if (clusterSize > NCCL_MAX_CGA_CLUSTER_SIZE) {
+      static bool warned = false;
+      if (warned == false) {
+        WARN("NCCL_CGA_CLUSTER_SIZE value %d is too big. Limiting value to %d.",
+             clusterSize, NCCL_MAX_CGA_CLUSTER_SIZE);
+        warned = true;
+      }
+      clusterSize = NCCL_MAX_CGA_CLUSTER_SIZE;
+    }
+
+    if (clusterSize && driverVersion >= 11080) {
+      cudaLaunchConfig_t launchConfig = {0};
+      cudaLaunchAttribute launchAttrs[2];
+      /* Cooperative Group Array (CGA)
+       * On sm90 and later we have an extra level of hierarchy where we
+       * can group together several blocks within the Grid, called
+       * Thread Block Clusters.
+       * Clusters enable multiple thread blocks running concurrently
+       * across multiple SMs to synchronize and collaboratively fetch
+       * and exchange data. A cluster of blocks are guaranteed to be
+       * concurrently scheduled onto a group of SMs.
+       * The maximum value is 8 and it must be divisible into the grid dimensions
+       */
+      // Grid dimension must be divisible by clusterSize
+      if (grid.x % clusterSize) clusterSize = 1;
+      launchAttrs[0].id = cudaLaunchAttributeClusterDimension;
+      launchAttrs[0].val.clusterDim = {clusterSize, 1, 1};
+      launchAttrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference;
+      launchAttrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread;
+
+      launchConfig.gridDim = grid;
+      launchConfig.blockDim = block;
+      launchConfig.attrs = launchAttrs;
+      launchConfig.numAttrs = sizeof(launchAttrs)/sizeof(launchAttrs[0]);
+      launchConfig.stream = launchStream;
+
+      CUDACHECK(cudaLaunchKernelExC(&launchConfig, fn, args));
+      return ncclSuccess;
+    }
+    #endif
+    // Standard kernel launch
+    CUDACHECK(cudaLaunchKernel(fn, grid, block, args, 0, launchStream));
   }
   return ncclSuccess;
 }
@@ -1005,18 +1079,22 @@ ncclResult_t ncclLaunchFinish(struct ncclComm* comm) {
     // Reset queue to empty without destroying plans since those will be sent
     // back to us for reclaiming via callbackQueue.
     ncclIntruQueueConstruct(&comm->planQueue);
-    // Close strong stream "transaction" encompassing cuda launches
-    NCCLCHECKGOTO(ncclStrongStreamRelease(tasks->capturingGraph, &comm->deviceStream), result, resume1);
+    cudaStream_t launchStream = tasks->streams->stream; // First user stream gets launch
+    // Create dependency for deviceStream on launchStream.
+    if (tasks->numStreams != 1) NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, &comm->deviceStream, launchStream), result, resume1);
   resume1:
-    // Create dependency for user streams on nccl device work.
-    struct ncclCudaStreamList* sl = tasks->streams;
-    tasks->streams = nullptr; // reset streams to empty
+    // Create dependency for other user streams (skip launch stream).
+    struct ncclCudaStreamList* sl = tasks->streams->next;
+    tasks->streams = nullptr; // Reset comm->tasks.streams to empty.
     while (sl != nullptr && tasks->numStreams != 1) {
       NCCLCHECKGOTO(ncclStrongStreamWaitStream(tasks->capturingGraph, sl->stream, &comm->deviceStream), result, resume2);
     resume2:
       sl = sl->next;
     }
     tasks->numStreams = 0;
+    // Release device stream as acquired in ncclLaunchPrepare()
+    NCCLCHECKGOTO(ncclStrongStreamRelease(tasks->capturingGraph, &comm->deviceStream), result, resume3);
+  resume3:;
   }
   return result;
 }
@@ -1412,20 +1490,20 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo const* inf
           NCCLCHECK(ncclChannelComputeFromBase(comm, channelBaseId, c, &channelId));
           if (isSendNotRecv) {
             if (comm->channels[channelId].peers[peer].send[1].connected == 0) { // P2P uses only 1 connector
-              comm->connectSend[peer] |= (1<<channelId);
+              comm->connectSend[peer] |= (1UL<<channelId);
               ncclGroupCommPreconnect(comm);
             }
             if (comm->p2pNet && comm->channels[channelId].peers[peer].send[NCCL_CONN_IDX_P2P_NET].connected == 0) {
-              comm->connectSend[peer+comm->nRanks*NCCL_CONN_IDX_P2P_NET] |= (1<<channelId);
+              comm->connectSend[peer+comm->nRanks*NCCL_CONN_IDX_P2P_NET] |= (1UL<<channelId);
               ncclGroupCommPreconnect(comm);
             }
           } else {
             if (comm->channels[channelId].peers[peer].recv[1].connected == 0) { // P2P uses only 1 connector
-              comm->connectRecv[peer] |= (1<<channelId);
+              comm->connectRecv[peer] |= (1UL<<channelId);
               ncclGroupCommPreconnect(comm);
             }
             if (comm->p2pNet && comm->channels[channelId].peers[peer].recv[NCCL_CONN_IDX_P2P_NET].connected == 0) {
-              comm->connectRecv[peer+comm->nRanks*NCCL_CONN_IDX_P2P_NET] |= (1<<channelId);
+              comm->connectRecv[peer+comm->nRanks*NCCL_CONN_IDX_P2P_NET] |= (1UL<<channelId);
               ncclGroupCommPreconnect(comm);
             }
           }
@@ -1442,7 +1520,7 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo const* inf
     if (comm->nRanks == 1 && opFull.op < ncclDevPreMulSum) {
       if (info->sendbuff != info->recvbuff) {
         size_t bytes = info->count*ncclTypeSize(info->datatype);
-        CUDACHECK(hipMemcpyAsync(info->recvbuff, info->sendbuff, bytes, hipMemcpyDeviceToDevice, info->stream));
+        CUDACHECK(cudaMemcpyAsync(info->recvbuff, info->sendbuff, bytes, cudaMemcpyDeviceToDevice, info->stream));
       }
       return ncclSuccess;
     } else {
@@ -1486,6 +1564,7 @@ static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo const* inf
       }
       if (l->stream == info->stream)
         break; // Already seen stream.
+      l = l->next;
     }
   }
   return ncclSuccess;
@@ -1501,8 +1580,8 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
   NCCLCHECKGOTO(ncclCommEnsureReady(info->comm), ret, fail);
 
   if (info->comm->checkPointers) {
-    CUDACHECKGOTO(hipGetDevice(&devOld), ret, fail);
-    CUDACHECKGOTO(hipSetDevice(info->comm->cudaDev), ret, fail);
+    CUDACHECKGOTO(cudaGetDevice(&devOld), ret, fail);
+    CUDACHECKGOTO(cudaSetDevice(info->comm->cudaDev), ret, fail);
   }
   NCCLCHECKGOTO(ArgsCheck(info), ret, fail);
 
@@ -1514,7 +1593,7 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
   NCCLCHECKGOTO(taskAppend(info->comm, info), ret, fail);
 
 exit:
-  if (devOld != -1) CUDACHECK(hipSetDevice(devOld));
+  if (devOld != -1) CUDACHECK(cudaSetDevice(devOld));
   ncclGroupErrCheck(ret);
   NCCLCHECK(ncclGroupEndInternal());
   /* if depth is 1, ncclGroupEndInternal() will trigger group ops. The state can change

projects/rccl/src/graph/connect.cc

@@ -35,7 +35,7 @@
 /******************************************************************/
 
 ncclResult_t ncclTopoPreset(struct ncclComm* comm,
-    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph,
+    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph, struct ncclTopoGraph* collNetGraph,
     struct ncclTopoRanks* topoRanks) {
   int rank = comm->rank;
   int nChannels = comm->nChannels;
@@ -60,6 +60,7 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm,
 
     int* ringIntra = ringGraph->intra+c*localRanks;
     int* treeIntra = treeGraph->intra+c*localRanks;
+    int* collNetIntra = collNetGraph->intra+c*localRanks;
 
     for (int i=0; i<localRanks; i++) {
       if (ringIntra[i] == rank) {
@@ -78,8 +79,10 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm,
         topoRanks->treeToChild1[c] = treeIntra[child1Index];
         channel->tree.up         = i == 0 ? -1 : treeIntra[i-1];
         channel->tree.down[0]    = i == localRanks-1 ? -1 : treeIntra[i+1];
-        channel->collnetChain.up         = i == 0 ? comm->nRanks : treeIntra[i-1];
-        channel->collnetChain.down[0]    = i == localRanks-1 ? -1 : treeIntra[i+1];
+      }
+      if (collNetIntra[i] == rank) {
+        channel->collnetChain.up      = i == 0 ? comm->nRanks : collNetIntra[i-1];
+        channel->collnetChain.down[0] = i == localRanks-1 ? -1 : collNetIntra[i+1];
       }
     }
     topoRanks->ringPrev[c] = channel->ring.prev;

projects/rccl/src/graph/paths.cc

@@ -424,6 +424,19 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* system, int64_t busId, int
   return ncclSuccess;
 }
 
+// Set to 0 to disable the flush on Hopper when using GDR
+NCCL_PARAM(NetForceFlush, "NET_FORCE_FLUSH", 1);
+
+// Determine whether we need to flush the GDR recv buffers
+ncclResult_t ncclTopoNeedFlush(struct ncclTopoSystem* system, int64_t busId, int* flush) {
+  int g;
+  NCCLCHECK(ncclTopoIdToIndex(system, GPU, busId, &g));
+  struct ncclTopoNode* gpu = system->nodes[GPU].nodes+g;
+  // Flush is required on Ampere and earlier
+  *flush = gpu->gpu.cudaCompCap < 90 ? 1 : ncclParamNetForceFlush();
+  return ncclSuccess;
+}
+
 NCCL_PARAM(NetDisableIntra, "NET_DISABLE_INTRA", 1);
 
 // Check whether going through the network would be faster than going through P2P/SHM.

projects/rccl/src/graph/rome_models.cc

@@ -721,7 +721,7 @@ newchannel:
   } while (str[offset++] != 0);
 end:
   graph->nChannels = nChannels;
-  graph->bwIntra = graph->bwInter = system->maxBw;
+  graph->bwIntra = graph->bwInter = system->totalBw/nChannels;
   if (graph->id == 1) {
     for (int i=0; i<graph->nChannels; i++) {
       int net;

projects/rccl/src/graph/search.cc

@@ -797,7 +797,7 @@ float speedArrayIntra[] = { 24.0, 20.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0, 6.0, 5
 float speedArrayInter[] = { 24.0, 20.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.4, 1.2, 0.24, 0.12 };
 #else
 float speedArrayIntra[] = { 44.0, 30.0, 22.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0, 6.0, 5.0, 4.0, 3.0 };
-float speedArrayInter[] = { 48.0, 30.0, 24.0, 22.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.4, 1.2, 0.24, 0.12 };
+float speedArrayInter[] = { 48.0, 30.0, 28.0, 24.0, 22.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.4, 1.2, 0.24, 0.12 };
 #endif
 #define NSPEEDSINTRA (sizeof(speedArrayIntra)/sizeof(float))
 #define NSPEEDSINTER (sizeof(speedArrayInter)/sizeof(float))
@@ -918,6 +918,7 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
   while (speedArray[speedIndex] > system->maxBw && speedIndex < nspeeds-1) speedIndex++;
   tmpGraph.bwIntra = tmpGraph.bwInter = speedArray[speedIndex];
   int64_t globalTimeout = NCCL_SEARCH_GLOBAL_TIMEOUT;
+
 search:
   int time = tmpGraph.sameChannels ? NCCL_SEARCH_TIMEOUT_SAMECHANNELS :
     tmpGraph.pattern == NCCL_TOPO_PATTERN_TREE ? NCCL_SEARCH_TIMEOUT_TREE : NCCL_SEARCH_TIMEOUT;
@@ -953,6 +954,7 @@ search:
     if (time != -1) globalTimeout += time;
     else globalTimeout = NCCL_SEARCH_GLOBAL_TIMEOUT;
     if (globalTimeout < 0 && graph->nChannels) goto done;
+
     int maxTypeIntra = system->nodes[NET].count > 0 ? tmpGraph.typeInter : PATH_SYS;
     if (tmpGraph.typeIntra < maxTypeIntra && (graph->nChannels == 0 || tmpGraph.typeIntra < graph->typeIntra)) {
       tmpGraph.typeIntra += 1;

projects/rccl/src/graph/tuning.cc

@@ -13,11 +13,11 @@
 NCCL_PARAM(Nthreads, "NTHREADS", -2);
 NCCL_PARAM(Ll128Nthreads, "LL128_NTHREADS", -2);
 
-static int getNthreads(const char* name, int env, int min, int max, int def) {
+static int getNthreads(const char* name, int env, int min, int max, int def, int WarpSize) {
   int nt = env;
   if (nt > 0) {
-    if (nt % WARP_SIZE != 0) {
-      WARN("Invalid %s %d (must be a multiple of %d)", name, nt, WARP_SIZE);
+    if (nt % WarpSize != 0) {
+      WARN("Invalid %s %d (must be a multiple of %d)", name, nt, WarpSize);
       nt = max;
     } else if (nt > max) {
       WARN("Invalid %s %d (maximum %d).", name, nt, max);
@@ -71,30 +71,30 @@ struct tuningModel {
 static struct tuningModel tuning_model_0 {
   .hwLat = {
     /* NVLINK */
-    { /* Tree (LL/LL128/Simple)*/ { 1.5, 1.5, 4.5 }, /* Ring (LL/LL128/Simple)*/ { 1.5, 1.5, 4.5 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 4.5 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 4.5 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.8, 1.4, 2.5 }, /* Ring (LL/LL128/Simple)*/ { 0.8, 2.2, 3.6 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 0.8 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 1.4 } },
     /* PCI */
     { /* Tree (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* Ring (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 5.7 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 5.7 } },
     /* NET */
-    { /* Tree (LL/LL128/Simple)*/ { 28.3, 28.3, 45.4 }, /* Ring (LL/LL128/Simple)*/ { 2.0, 2.0, 24.1 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 45.4 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 45.4 } },
+    { /* Tree (LL/LL128/Simple)*/ { 11.8, 18.2, 20.8 }, /* Ring (LL/LL128/Simple)*/ { 9.5, 19.8, 15.1 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 11.8 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 18.2 } },
   },
 
   .bwRatio = {
     /* 2 nodes */
-    { /* Tree (LL/LL128/Simple)*/ { 0.06, 1.00, 1.30 }, /* Ring (LL/LL128/Simple)*/ { 0.07, 1.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.04, 0.22, 0.91 }, /* Ring (LL/LL128/Simple)*/ { 0.04, 0.34, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
     /* more than 2 nodes */
-    { /* Tree (LL/LL128/Simple)*/ { 0.06, 1.00, 0.30 }, /* Ring (LL/LL128/Simple)*/ { 0.07, 1.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.04, 0.22, 0.95 }, /* Ring (LL/LL128/Simple)*/ { 0.04, 0.34, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
   },
 
   .treeCorrectionFactor = {
-    { 0.3, 0.9, 0.8, 0.7, 0.6, 0.3, 0.1, 0.4, 0.8, 0.8, 0.5, 0.3, 0.4, 0.3, 0.2, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.3, 0.9, 0.8, 0.7, 0.6, 0.3, 0.1, 0.4, 0.8, 0.8, 0.5, 0.3, 0.4, 0.3, 0.2, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.4, 1.0, 1.0, 0.8, 1.0, 1.0, 0.3, 1.0, 0.9, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.7, 0.4, 0.3, 0.3, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, },
+    { 0.1, 0.2, 0.1, 0.1, 0.9, 0.3, 0.4, 0.1, 0.2, 0.4, 0.2, 0.1, 0.3, 0.3, 0.2, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
+    { 0.1, 0.3, 1.0, 0.1, 0.5, 1.0, 0.9, 1.0, 1.0, 1.0, 0.3, 0.1, 0.4, 0.5, 0.5, 0.4, 0.4, 0.3, 0.3, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, },
+    { 0.2, 1.0, 0.1, 0.1, 0.7, 0.2, 0.4, 0.1, 0.1, 0.3, 0.4, 0.3, 0.6, 0.8, 1.0, 1.0, 1.0, 1.0, 0.9, 0.8, 0.8, 0.8, 0.8, 0.8, 0.9, 0.9, 0.9, },
   },
 
   .ringCorrectionFactor = {
-    { 0.2, 0.7, 0.7, 0.6, 0.6, 0.3, 0.2, 0.5, 1.0, 1.0, 0.8, 0.6, 0.8, 0.6, 0.3, 0.3, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.2, 0.7, 0.7, 0.6, 0.6, 0.3, 0.2, 0.5, 1.0, 1.0, 0.8, 0.6, 0.8, 0.6, 0.3, 0.3, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 1.0, 0.6, 0.9, 1.0, 0.7, 0.7, 1.0, 0.6, 0.8, 0.2, 0.1, 0.1, 0.1, 0.1, 0.2, 0.5, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
+    { 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.4, 0.2, 0.3, 0.5, 0.3, 0.1, 0.5, 0.5, 0.3, 0.2, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, },
+    { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.8, 0.7, 0.5, 0.4, 0.4, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, },
+    { 1.0, 0.8, 0.2, 1.0, 1.0, 0.3, 1.0, 0.1, 0.1, 0.2, 0.2, 0.1, 0.5, 1.0, 0.8, 0.8, 1.0, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
 };
 
@@ -161,30 +161,30 @@ static struct tuningModel tuning_model_2 {
 static struct tuningModel tuning_model_3 {
   .hwLat = {
     /* NVLINK */
-    { /* Tree (LL/LL128/Simple)*/ { 1.5, 1.5, 4.5 }, /* Ring (LL/LL128/Simple)*/ { 1.5, 1.5, 4.5 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 4.5 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 4.5 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.8, 0.0, 2.5 }, /* Ring (LL/LL128/Simple)*/ { 0.8, 0.0, 3.6 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 0.8 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 } },
     /* PCI */
     { /* Tree (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* Ring (LL/LL128/Simple)*/ { 2.2, 2.2, 5.7 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 5.7 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 5.7 } },
     /* NET */
-    { /* Tree (LL/LL128/Simple)*/ { 17.4, 17.4, 40.3 }, /* Ring (LL/LL128/Simple)*/ { 4.1, 4.1, 40.6 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 40.3 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 40.3 } },
+    { /* Tree (LL/LL128/Simple)*/ { 12.5, 0.0, 22.4 }, /* Ring (LL/LL128/Simple)*/ { 9.5, 0.0, 19.8 }, /* CollNetDirect (Simple)*/ { 0.0, 0.0, 12.5 }, /* CollNetChain (Simple)*/ { 0.0, 0.0, 0.0 } },
   },
 
   .bwRatio = {
     /* 2 nodes */
-    { /* Tree (LL/LL128/Simple)*/ { 0.08, 1.00, 0.95 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 1.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.20, 0.00, 1.75 }, /* Ring (LL/LL128/Simple)*/ { 0.20, 0.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
     /* more than 2 nodes */
-    { /* Tree (LL/LL128/Simple)*/ { 0.08, 1.00, 0.41 }, /* Ring (LL/LL128/Simple)*/ { 0.08, 1.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
+    { /* Tree (LL/LL128/Simple)*/ { 0.20, 0.00, 0.96 }, /* Ring (LL/LL128/Simple)*/ { 0.20, 0.00, 1.00 }, /* CollNetDirect (Simple)*/ { 0.00, 0.00, 1.00 }, /* CollNetChain (Simple)*/ { 0.00, 0.00, 1.00 } },
   },
 
   .treeCorrectionFactor = {
-    { 0.6, 1.0, 1.0, 1.0, 0.1, 0.2, 0.1, 0.2, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.9, 0.9, 0.8, 0.6, 0.4, 0.5, 0.4, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.6, 1.0, 1.0, 1.0, 0.1, 0.2, 0.1, 0.2, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.9, 0.9, 0.8, 0.6, 0.4, 0.5, 0.4, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 1.0, 0.1, 0.1, 0.1, 1.0, 1.0, 1.0, 1.0, 0.7, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.5, 0.6, 0.6, 0.6, 0.6, },
+    { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0, 1.0, 0.2, 1.0, 0.9, 1.0, 0.6, 0.4, 0.6, 0.4, 0.3, 0.3, 0.3, 0.3, 0.3, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, },
+    { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, },
+    { 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.1, 0.1, 0.1, 0.2, 1.0, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.8, 0.7, 0.8, 0.9, 0.7, 0.7, },
   },
 
   .ringCorrectionFactor = {
-    { 0.5, 0.2, 0.1, 0.1, 0.3, 0.3, 0.1, 0.3, 0.7, 0.8, 0.5, 0.4, 0.2, 0.1, 0.1, 0.1, 0.3, 0.5, 0.4, 0.4, 0.3, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.5, 0.2, 0.1, 0.1, 0.3, 0.3, 0.1, 0.3, 0.7, 0.8, 0.5, 0.4, 0.2, 0.1, 0.1, 0.1, 0.3, 0.5, 0.4, 0.4, 0.3, 0.2, 0.1, 0.1, 0.1, 0.1, 0.1, },
-    { 0.2, 0.3, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 1.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.2, 0.6, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
+    { 0.1, 0.1, 0.1, 0.1, 0.1, 0.3, 0.1, 0.2, 0.1, 0.4, 0.4, 0.2, 0.2, 0.3, 0.7, 0.5, 0.4, 0.3, 0.3, 0.3, 0.3, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, },
+    { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, },
+    { 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.5, 1.0, 0.1, 0.3, 0.1, 0.1, 0.1, 0.2, 0.2, 0.2, 0.3, 0.4, 0.7, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, },
   },
 };
 
@@ -226,22 +226,36 @@ static struct tuningModel rcclTuningModel[] = {
   tuning_model_4,
 };
 
+/* Array indexes used below */
+#define VOLTA_COMPCAP_IDX 0
+#define AMPERE_COMPCAP_IDX 1
+#define HOPPER_COMPCAP_IDX 2
+
 // LL128 max BW per channel
 static const double ll128MaxBwPerCh = 20.0;
-static const double llMaxBws[2][3] = { /* Volta-N1/Intel-N2/Intel-N4) */ {39.0, 39.0, 20.4}, /* Ampere-N1/AMD-N2/AMD-N4) */ {87.7, 22.5 /*avg of ring & tree*/, 19.0} };
-static const double perChMaxTreeBws[2][3] = { /* Volta (N1/N2/N4) */ {26.5, 18.5, 10.0}, /* Ampere (N1/N2/N4) */ {24.0, 23.6, 17.8} };
+static const double llMaxBws[3][3] = {
+  /* Volta-N1/Intel-N2/Intel-N4) */ {39.0, 39.0, 20.4},
+  /* Ampere-N1/AMD-N2/AMD-N4) */ {87.7, 22.5 /*avg of ring & tree*/, 19.0},
+  /* Hopper-N1/AMD-N2/AMD-N4) */ {87.7, 22.5 /*avg of ring & tree*/, 19.0}
+};
+
+static const double perChMaxTreeBws[3][3] = {
+  /* Volta (N1/N2/N4) */ {26.5, 18.5, 10.0},
+  /* Ampere (N1/N2/N4) */ {24.0, 23.6, 17.8},
+  /* Hopper (N1/N2/N4) */ {24.0, 23.6, 17.8},
+};
 
 ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCompCap, struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph, struct ncclTopoGraph* collNetGraph) {
   int simpleDefaultThreads = (ringGraph->bwIntra*ringGraph->nChannels <= PCI_BW) ? 256 : NCCL_SIMPLE_MAX_NTHREADS;
   comm->maxThreads[NCCL_ALGO_RING][NCCL_PROTO_SIMPLE] =
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, simpleDefaultThreads);
+    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, simpleDefaultThreads, comm->WarpSize);
   comm->maxThreads[NCCL_ALGO_TREE][NCCL_PROTO_SIMPLE] = comm->maxThreads[NCCL_ALGO_COLLNET_DIRECT][NCCL_PROTO_SIMPLE] =
-    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, NCCL_MAX_NTHREADS);
+    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, NCCL_MAX_NTHREADS, comm->WarpSize);
   comm->maxThreads[NCCL_ALGO_RING][NCCL_PROTO_LL] = comm->maxThreads[NCCL_ALGO_TREE][NCCL_PROTO_LL] = comm->maxThreads[NCCL_ALGO_COLLNET_DIRECT][NCCL_PROTO_LL] =
-    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, NCCL_MAX_NTHREADS);
+    getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 4*comm->WarpSize, NCCL_MAX_NTHREADS, NCCL_MAX_NTHREADS, comm->WarpSize);
   comm->maxThreads[NCCL_ALGO_RING][NCCL_PROTO_LL128] = comm->maxThreads[NCCL_ALGO_TREE][NCCL_PROTO_LL128] =
-    getNthreads("NCCL_LL128_NTHREADS", ncclParamLl128Nthreads(), 4*comm->WarpSize, NCCL_LL128_MAX_NTHREADS, NCCL_LL128_MAX_NTHREADS);
+    getNthreads("NCCL_LL128_NTHREADS", ncclParamLl128Nthreads(), 4*comm->WarpSize, NCCL_LL128_MAX_NTHREADS, NCCL_LL128_MAX_NTHREADS, comm->WarpSize);
 #else
     getNthreads("NCCL_NTHREADS", ncclParamNthreads(), 2*WARP_SIZE, NCCL_SIMPLE_MAX_NTHREADS, simpleDefaultThreads);
   comm->maxThreads[NCCL_ALGO_TREE][NCCL_PROTO_SIMPLE] =
@@ -258,14 +272,14 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
   int nRanks = comm->nRanks;
   if (nRanks <= 1) return ncclSuccess;
 
-  int compCap80 = minCompCap == 80 && maxCompCap == 80 ? 1 : 0;
+  int compCapIndex = (minCompCap == 80 && maxCompCap == 80) ? AMPERE_COMPCAP_IDX : ((minCompCap == 90 && maxCompCap == 90) ? HOPPER_COMPCAP_IDX : VOLTA_COMPCAP_IDX);
   int cpuArch, cpuVendor, cpuModel;
   NCCLCHECK(ncclTopoCpuType(comm->topo, &cpuArch, &cpuVendor, &cpuModel));
   int index2 = nNodes <= 2 ? nNodes-1 : 2;
   // LL: for single node, we look at GPU type; for multi-node, we look at CPU type
-  int index1 = nNodes == 1 ? compCap80 : cpuVendor == NCCL_TOPO_CPU_VENDOR_AMD ? 1 : 0;
+  int index1 = nNodes == 1 ? compCapIndex : cpuVendor == NCCL_TOPO_CPU_VENDOR_AMD ? 1 : 0;
   double llMaxBw = llMaxBws[index1][index2];
-  double perChMaxTreeBw = perChMaxTreeBws[compCap80][index2];
+  double perChMaxTreeBw = perChMaxTreeBws[compCapIndex][index2];
   // De-penalize Tree/Simple latency on Power systems to favor Tree than Ring
   //if (cpuArch == NCCL_TOPO_CPU_ARCH_POWER) hwLat[NCCL_HW_PCI][NCCL_ALGO_TREE][NCCL_PROTO_SIMPLE] = hwLat[NCCL_HW_PCI][NCCL_ALGO_RING][NCCL_PROTO_SIMPLE];
   float ppn = (float)nRanks / nNodes; // if ppn < 2, then we are sending/receiving at the same GPU through the NIC, apply some bw discount
@@ -290,6 +304,7 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
         int collnet = (a == NCCL_ALGO_COLLNET_DIRECT || a == NCCL_ALGO_COLLNET_CHAIN) ? 1 : 0;
         float bw = nNodes <= 2 || collnet ? graphs[a]->bwIntra : graphs[a]->bwInter;
         float busBw = comm->topo->baseBw != 0.0 ? comm->topo->baseBw : graphs[a]->nChannels * bw;
+        //INFO(NCCL_INIT, "algo %s proto %s busBw %f baseBw %f bw %f nChannels %d bwIntra %f bwInter %f", ncclAlgoStr[a], ncclProtoStr[p], busBw, comm->topo->baseBw, bw, graphs[a]->nChannels, graphs[a]->bwIntra, graphs[a]->bwInter);
 
         // Various model refinements
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
@@ -298,7 +313,7 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
         else
           busBw *= rcclTuningModel[comm->topo->tuning].bwRatio[1][a][p];
 #else
-        if (compCap80) busBw = std::min(busBw, 235.0f);
+        if (compCapIndex == AMPERE_COMPCAP_IDX) busBw = std::min(busBw, 235.0f);
         if (a == NCCL_ALGO_RING && p == NCCL_PROTO_LL) { busBw = std::min(llMaxBw, busBw * ((nNodes > 1 || coll == ncclFuncAllReduce || coll == ncclFuncReduce) ? 1.0/4.0 : 1.0/3.0)); }
         if (a == NCCL_ALGO_RING && p == NCCL_PROTO_LL128) busBw = std::min(busBw * (ppn < 2 ? 0.7 : 0.92 /*120.0/128.0*/), ll128MaxBwPerCh*graphs[a]->nChannels);
         if (a == NCCL_ALGO_TREE) busBw = std::min(busBw*.92, graphs[a]->nChannels*perChMaxTreeBw);
@@ -306,14 +321,14 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
         if (a == NCCL_ALGO_TREE && p == NCCL_PROTO_LL128) busBw = std::min(busBw * (nNodes == 1 ? 7.0/9.0 : 120.0/128.0), ll128MaxBwPerCh*graphs[a]->nChannels);
         if (a == NCCL_ALGO_COLLNET_DIRECT && p != NCCL_PROTO_SIMPLE) busBw = 0;  // Not used
         if (a == NCCL_ALGO_COLLNET_CHAIN && p != NCCL_PROTO_SIMPLE) busBw = 0;  // Not used
-      	if (a == NCCL_ALGO_COLLNET_DIRECT && p == NCCL_PROTO_SIMPLE) {
-                // Collnet+Direct requires all GPUs to have a local NIC to work at full speed
-                float factor = ppn / (1.0*graphs[a]->nChannels); // GPU/NIC ratio
-      	  factor -= (factor-1)/2;
-                busBw /= factor;
-      	}
-      	if (a == NCCL_ALGO_COLLNET_CHAIN && p == NCCL_PROTO_SIMPLE) busBw *= .75;
+        if (a == NCCL_ALGO_COLLNET_DIRECT && p == NCCL_PROTO_SIMPLE) {
+          // Collnet+Direct requires all GPUs to have a local NIC to work at full speed
+          float factor = ppn / (1.0*graphs[a]->nChannels); // GPU/NIC ratio
+          factor -= (factor-1)/2;
+          busBw /= factor;
+        }
 #endif
+        if (a == NCCL_ALGO_COLLNET_CHAIN && p == NCCL_PROTO_SIMPLE) busBw *= .75;
 
         // Convert bus BW to algorithm BW
         float ratio = (a != NCCL_ALGO_RING) ? .5 : (1.0 * nRanks) / nsteps;
@@ -387,11 +402,18 @@ ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCom
       pEnable = (graphs[a]->typeInter <= PATH_PXB) && graphs[a]->typeIntra <= PATH_NVL &&
         (comm->topo->nodes[GPU].nodes[0].gpu.gcn == 910 && comm->topo->ll128Enabled) ? 1 : 0;
 #else
-      // Enable LL128 by default only on Volta/Ampere+NVLink. Other cases are not tested and may cause silent data corruption.
-      pEnable = (graphs[a]->typeInter <= PATH_PXB) && graphs[a]->typeIntra <= PATH_NVL &&
-        ((minCompCap == 70 && maxCompCap == 70) || (minCompCap == 80 && maxCompCap == 80)) ? 1 : 0;
+      // Enable LL128 by default only on Volta/Ampere/Hopper+NVLink. Other cases are not tested and may cause silent data corruption.
+      pEnable = 1;
+      pEnable &= (graphs[a]->typeInter <= PATH_PXB);
+      pEnable &= (graphs[a]->typeIntra <= PATH_NVL);
+      pEnable &= (minCompCap == maxCompCap);
+      switch (minCompCap) {
+      case 70: pEnable &= 1; break;
+      case 80: pEnable &= 1; break;
+      case 90: pEnable &= !(CUDART_VERSION == 11080 && c == ncclFuncAllReduce && a == NCCL_ALGO_RING && comm->nRanks == 2); break;
+      default: pEnable &= 0; break;
+      }
 #endif
-      if (comm->rank == 0 && c == 0 && a == 0) INFO(NCCL_INIT, "Using tuning table %d with LL128 %s", comm->topo->tuning, pEnable ? "enabled" : "disabled");
     }
     if (pEnable == 0) comm->bandwidths[c][a][p] = 0;
     // Only disable algo for Allreduce since others only have one

projects/rccl/src/graph/xml.cc

@@ -615,7 +615,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev
     if (rocmDev == -1) {
       const char* busId;
       NCCLCHECK(xmlGetAttr(pciNode, "busid", &busId));
-      if (busId == NULL || hipDeviceGetByPCIBusId(&dev, busId) != hipSuccess) dev = -1;
+      if (busId == NULL || cudaDeviceGetByPCIBusId(&dev, busId) != cudaSuccess) dev = -1;
     } else {
       dev = rocmDev;
     }
@@ -627,8 +627,8 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "sm", &index));
   if (index == -1) {
     int cudaMajor, cudaMinor;
-    hipDeviceProp_t devProp;
-    CUDACHECK(hipGetDeviceProperties(&devProp, 0));
+    cudaDeviceProp devProp;
+    CUDACHECK(cudaGetDeviceProperties(&devProp, dev));
     cudaMajor = devProp.major; cudaMinor = devProp.minor;
     NCCLCHECK(xmlSetAttrInt(gpuNode, "sm", cudaMajor*10+cudaMinor));
   }
@@ -688,7 +688,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev
     }
 #else
     // NVML NVLink detection
-    int maxNvLinks = (sm < 60) ? 0 : (sm < 70) ? 4 : (sm < 80) ? 6 : 12;
+    int maxNvLinks = (sm < 60) ? 0 : (sm < 70) ? 4 : (sm < 80) ? 6 : (sm < 90) ? 12 : 18;
 
     if (maxNvLinks > 0 && nvmlDev == NULL) {
       WARN("No NVML device handle. Skipping nvlink detection.");
@@ -701,8 +701,21 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev
       if ((ncclNvmlDeviceGetNvLinkCapability(nvmlDev, l, NVML_NVLINK_CAP_P2P_SUPPORTED, &canP2P) != ncclSuccess) || !canP2P) continue;
 
       // Make sure the Nvlink is up. The previous call should have trained the link.
-      nvmlEnableState_t isActive;
-      if ((ncclNvmlDeviceGetNvLinkState(nvmlDev, l, &isActive) != ncclSuccess) || (isActive != NVML_FEATURE_ENABLED)) continue;
+      nvmlEnableState_t isActive = NVML_FEATURE_DISABLED;
+#if CUDART_VERSION >= 11080
+      if (sm >= 90) {
+        nvmlFieldValue_t fv;
+        fv.fieldId = NVML_FI_DEV_NVLINK_GET_STATE;
+        fv.scopeId = l;
+        // fv.value will contain NV_FEATURE_ENABLED or NV_FEATURE_DISABLED
+        if ((ncclNvmlDeviceGetFieldValues(nvmlDev, 1, &fv) == ncclSuccess) && (fv.nvmlReturn == NVML_SUCCESS))
+          isActive = (nvmlEnableState_t) fv.value.uiVal;
+      } else /* FALLTHRU to GetNvLinkState if before SM90 */
+#endif
+      {
+        (void) ncclNvmlDeviceGetNvLinkState(nvmlDev, l, &isActive);
+      }
+      if (isActive != NVML_FEATURE_ENABLED) continue;
 
       // Try to figure out what's on the other side of the NVLink
       nvmlPciInfo_t remoteProc;

projects/rccl/src/group.cc

@@ -117,7 +117,7 @@ struct ncclPreconnectJob {
 ncclResult_t ncclPreconnectFunc(struct ncclAsyncJob* job_) {
   struct ncclPreconnectJob* job = (struct ncclPreconnectJob*)job_;
   struct ncclComm* comm = job->comm;
-  CUDACHECK(hipSetDevice(comm->cudaDev));
+  CUDACHECK(cudaSetDevice(comm->cudaDev));
   if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
   NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 1));
   if (comm->p2pNet) NCCLCHECK(ncclTransportP2pSetup(comm, NULL, NCCL_CONN_IDX_P2P_NET));
@@ -138,7 +138,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
     bool capturingYes = false, capturingNo = false;
     do {
       (ncclCudaGraphValid(comm->tasks.capturingGraph) ? capturingYes : capturingNo) = true;
-      CUDACHECKGOTO(hipSetDevice(comm->cudaDev), result, failure);
+      CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
       NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);
       if (useBarrier) ncclCommIntraBarrierIn(comm, 1);
       comm = comm->groupNext;
@@ -170,7 +170,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
           struct ncclKernelPlan* plan = comm->unlaunchedPlansHead;
           if (plan != nullptr) {
             comm->unlaunchedPlansHead = plan->next;
-            CUDACHECKGOTO(hipSetDevice(comm->cudaDev), result, failure);
+            CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
             NCCLCHECKGOTO(ncclLaunchKernelBefore_NoUncapturedCuda(comm, plan), result, failure);
             NCCLCHECKGOTO(ncclLaunchKernel(comm, plan), result, failure);
           }
@@ -180,7 +180,7 @@ static ncclResult_t doLaunches(struct ncclComm* head) {
             NCCLCHECKGOTO(ncclLaunchKernelAfter_NoCuda(comm, plan), result, failure);
           }
         } else { // Final round.
-          CUDACHECKGOTO(hipSetDevice(comm->cudaDev), result, failure);
+          CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
           NCCLCHECKGOTO(ncclLaunchFinish(comm), result, failure);
         }
         comm = next;
@@ -213,8 +213,8 @@ static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclComm** g
     for (int i = 0; i < comm->nRanks; i++) {
       comm->tasks.peers[i].sendSeen = false;
       comm->tasks.peers[i].recvSeen = false;
-      comm->connectSend[i] = 0;
-      comm->connectRecv[i] = 0;
+      comm->connectSend[i] = 0UL;
+      comm->connectRecv[i] = 0UL;
     }
     comm->unlaunchedPlansHead = nullptr;
     // Reclaim abandoned kernel plan memory. Note ncclWork structs were already
@@ -276,7 +276,7 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
   struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> *asyncJobsMain = gjob->asyncJobsPtr;
   volatile bool *groupAbortFlag = gjob->abortFlagPtr;
 
-  CUDACHECKGOTO(hipGetDevice(&savedDev), ret, fail);
+  CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
 
   if (groupCommPreconnectHeadMain != nullptr) {
     struct ncclComm* comm = groupCommPreconnectHeadMain;
@@ -333,6 +333,8 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
 
         job = job->next;
       } while (job != nullptr);
+      // Let preconnect threads progress.
+      if (jobsDone == false) usleep(1);
     } while (jobsDone == false);
 
     if (ret != ncclSuccess) goto fail;
@@ -366,7 +368,7 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_) {
   *gjob->groupCommHeadPtr = nullptr;
   *gjob->groupCommPreconnectHeadPtr = nullptr;
 
-  CUDACHECK(hipSetDevice(savedDev));
+  CUDACHECK(cudaSetDevice(savedDev));
 
 exit:
   return ret;

projects/rccl/src/include/alloc.h

@@ -23,13 +23,13 @@ uint64_t clockNano(); // from utils.h with which we have a circular dependency
 template <typename T>
 ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
   *ptr = nullptr;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), hipHostMallocMapped), result, finish);
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECKGOTO(hipHostMalloc(ptr, nelem*sizeof(T), cudaHostAllocMapped), result, finish);
   memset(*ptr, 0, nelem*sizeof(T));
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (*ptr == nullptr) WARN("Failed to CUDA host alloc %ld bytes", nelem*sizeof(T));
   INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return result;
@@ -37,7 +37,7 @@ finish:
 #define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 inline ncclResult_t ncclCudaHostFree(void* ptr) {
-  CUDACHECK(hipHostFree(ptr));
+  CUDACHECK(cudaFreeHost(ptr));
   return ncclSuccess;
 }
 
@@ -90,15 +90,15 @@ extern struct allocationTracker allocTracker[];
 template <typename T>
 ncclResult_t ncclCudaMallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
   *ptr = nullptr;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (isFineGrain)
     CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
   else
-    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
+    CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (*ptr == nullptr) WARN("Failed to CUDA malloc %ld bytes", nelem*sizeof(T));
   INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return result;
@@ -106,23 +106,23 @@ finish:
 #define ncclCudaMalloc(...) ncclCudaMallocDebug( __FILE__, __LINE__, __VA_ARGS__)
 
 template <typename T>
-ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t sideStream = nullptr, bool isFineGrain = false) {
+ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, cudaStream_t sideStream = nullptr, bool isFineGrain = false) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
   *ptr = nullptr;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   // Need a side stream so as not to interfere with graph capture.
-  hipStream_t stream = sideStream;
+  cudaStream_t stream = sideStream;
   if (stream == nullptr)
-    CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+    CUDACHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
   if (isFineGrain)
     CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
   else
-    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
-  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
-  CUDACHECKGOTO(hipStreamSynchronize(stream), result, finish);
+    CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
+  CUDACHECKGOTO(cudaMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
+  CUDACHECKGOTO(cudaStreamSynchronize(stream), result, finish);
   if (sideStream == nullptr)
-    CUDACHECKGOTO(hipStreamDestroy(stream), result, finish);
+    CUDACHECKGOTO(cudaStreamDestroy(stream), result, finish);
   int dev;
   CUDACHECK(hipGetDevice(&dev));
   if (dev < MAX_ALLOC_TRACK_NGPU) {
@@ -130,8 +130,8 @@ ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t
     __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
   }
 finish:
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (*ptr == nullptr) WARN("Failed to CUDA calloc %ld bytes", nelem*sizeof(T));
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
   INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return result;
 }
@@ -140,14 +140,14 @@ finish:
 template <typename T>
 ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t stream, bool isFineGrain = false) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
   *ptr = nullptr;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (isFineGrain)
     CUDACHECKGOTO(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained), result, finish);
   else
-    CUDACHECKGOTO(hipMalloc(ptr, nelem*sizeof(T)), result, finish);
-  CUDACHECKGOTO(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
+    CUDACHECKGOTO(cudaMalloc(ptr, nelem*sizeof(T)), result, finish);
+  CUDACHECKGOTO(cudaMemsetAsync(*ptr, 0, nelem*sizeof(T), stream), result, finish);
   int dev;
   CUDACHECK(hipGetDevice(&dev));
   if (dev < MAX_ALLOC_TRACK_NGPU) {
@@ -155,7 +155,7 @@ ncclResult_t ncclCudaCallocAsyncDebug(const char *filefunc, int line, T** ptr, s
     __atomic_fetch_add(&allocTracker[dev].totalAllocSize, nelem*sizeof(T), __ATOMIC_RELAXED);
   }
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   if (*ptr == nullptr) WARN("Failed to CUDA calloc async %ld bytes", nelem*sizeof(T));
   INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return result;
@@ -165,38 +165,38 @@ finish:
 template <typename T>
 ncclResult_t ncclCudaMemcpy(T* dst, T* src, size_t nelem) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   // Need a side stream so as not to interfere with graph capture.
-  hipStream_t stream;
-  CUDACHECKGOTO(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking), result, finish);
+  cudaStream_t stream;
+  CUDACHECKGOTO(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), result, finish);
   NCCLCHECKGOTO(ncclCudaMemcpyAsync(dst, src, nelem, stream), result, finish);
-  CUDACHECKGOTO(hipStreamSynchronize(stream), result, finish);
-  CUDACHECKGOTO(hipStreamDestroy(stream), result, finish);
+  CUDACHECKGOTO(cudaStreamSynchronize(stream), result, finish);
+  CUDACHECKGOTO(cudaStreamDestroy(stream), result, finish);
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   return result;
 }
 
 template <typename T>
-ncclResult_t ncclCudaMemcpyAsync(T* dst, T* src, size_t nelem, hipStream_t stream) {
+ncclResult_t ncclCudaMemcpyAsync(T* dst, T* src, size_t nelem, cudaStream_t stream) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  CUDACHECKGOTO(hipMemcpyAsync(dst, src, nelem*sizeof(T), hipMemcpyDefault, stream), result, finish);
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECKGOTO(cudaMemcpyAsync(dst, src, nelem*sizeof(T), cudaMemcpyDefault, stream), result, finish);
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   return result;
 }
 
 template <typename T>
 ncclResult_t ncclCudaFree(T* ptr) {
   ncclResult_t result = ncclSuccess;
-  hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
-  CUDACHECKGOTO(hipFree(ptr), result, finish);
+  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECKGOTO(cudaFree(ptr), result, finish);
 finish:
-  CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode));
+  CUDACHECK(cudaThreadExchangeStreamCaptureMode(&mode));
   return result;
 }
 

projects/rccl/src/include/checks.h

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -12,17 +11,17 @@
 
 // Check CUDA RT calls
 #define CUDACHECK(cmd) do {                                 \
-    hipError_t err = cmd;                                    \
-    if( err != hipSuccess ) {                                \
-        WARN("HIP failure '%s'", hipGetErrorString(err));   \
+    cudaError_t err = cmd;                                  \
+    if( err != cudaSuccess ) {                              \
+        WARN("Cuda failure '%s'", cudaGetErrorString(err)); \
         return ncclUnhandledCudaError;                      \
     }                                                       \
 } while(false)
 
 #define CUDACHECKGOTO(cmd, res, label) do {                 \
-    hipError_t err = cmd;                                    \
-    if( err != hipSuccess ) {                                \
-        WARN("HIP failure '%s'", hipGetErrorString(err));   \
+    cudaError_t err = cmd;                                  \
+    if( err != cudaSuccess ) {                              \
+        WARN("Cuda failure '%s'", cudaGetErrorString(err)); \
         res = ncclUnhandledCudaError;                       \
         goto label;                                         \
     }                                                       \
@@ -30,10 +29,10 @@
 
 // Report failure but clear error and continue
 #define CUDACHECKIGNORE(cmd) do {  \
-    hipError_t err = cmd;         \
-    if( err != hipSuccess ) {     \
-        INFO(NCCL_ALL,"%s:%d Cuda failure '%s'", __FILE__, __LINE__, hipGetErrorString(err)); \
-        (void) hipGetLastError(); \
+    cudaError_t err = cmd;         \
+    if( err != cudaSuccess ) {     \
+        INFO(NCCL_ALL,"%s:%d Cuda failure '%s'", __FILE__, __LINE__, cudaGetErrorString(err)); \
+        (void) cudaGetLastError(); \
     }                              \
 } while(false)
 
@@ -151,7 +150,7 @@
 } while(0)
 
 #define CUDACHECKTHREAD(a) do { \
-  if ((a) != hipSuccess) { \
+  if ((a) != cudaSuccess) { \
     INFO(NCCL_INIT,"%s:%d -> %d [Async thread]", __FILE__, __LINE__, args->ret); \
     args->ret = ncclUnhandledCudaError; \
     return args; \

projects/rccl/src/include/comm.h

@@ -174,8 +174,8 @@ struct ncclComm {
   ncclCollNet_t* ncclCollNet;
   void* bootstrap;
   // Bitmasks for ncclTransportP2pSetup
-  uint32_t* connectSend;
-  uint32_t* connectRecv;
+  uint64_t* connectSend;
+  uint64_t* connectRecv;
 
   int rank;    // my rank in the communicator
   int nRanks;  // number of GPUs in communicator

projects/rccl/src/include/cudawrap.h

@@ -8,6 +8,8 @@
 #define NCCL_CUDAWRAP_H_
 
 #include <cuda.h>
+#include <cuda_runtime.h>
+#include "checks.h"
 
 #if CUDART_VERSION >= 11030
 #include <cudaTypedefs.h>
@@ -83,6 +85,18 @@ DECLARE_CUDA_PFN_EXTERN(cuDriverGetVersion, 2020);
 DECLARE_CUDA_PFN_EXTERN(cuGetProcAddress, 11030);
 
 
-ncclResult_t cudaLibraryInit(void);
+ncclResult_t ncclCudaLibraryInit(void);
+
+extern int ncclCudaDriverVersionCache;
+
+inline ncclResult_t ncclCudaDriverVersion(int* driver) {
+  int version = __atomic_load_n(&ncclCudaDriverVersionCache, __ATOMIC_RELAXED);
+  if (version == -1) {
+    CUDACHECK(cudaDriverGetVersion(&version));
+    __atomic_store_n(&ncclCudaDriverVersionCache, version, __ATOMIC_RELAXED);
+  }
+  *driver = version;
+  return ncclSuccess;
+}
 
 #endif

projects/rccl/src/include/graph.h

@@ -37,6 +37,7 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 #define MAX_XGMI_INTER_GPUS 4
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int* dev);
 ncclResult_t ncclTopoGetLinkType(struct ncclTopoSystem* system, int cudaDev1, int cudaDev2, bool* isXGMI, int maxInter=MAX_XGMI_INTER_GPUS, int nInter=0, int *inter=nullptr);
+ncclResult_t ncclTopoNeedFlush(struct ncclTopoSystem* system, int64_t busId, int* flush);
 ncclResult_t ncclTopoCheckNet(struct ncclTopoSystem* system, int64_t id1, int64_t id2, int* net);
 int ncclPxnDisable(struct ncclComm* comm);
 ncclResult_t ncclTopoGetPxnRanks(struct ncclComm* comm, int** intermediateRanks, int* nranks);
@@ -109,7 +110,7 @@ struct ncclTopoRanks {
 };
 
 ncclResult_t ncclTopoPreset(struct ncclComm* comm,
-    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph,
+    struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph, struct ncclTopoGraph* collNetGraph,
     struct ncclTopoRanks* topoRanks);
 
 ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePatterns,

projects/rccl/src/include/info.h

@@ -41,7 +41,7 @@ struct ncclInfo {
   ncclRedOp_t op;
   int root; // peer for p2p operations
   ncclComm_t comm;
-  hipStream_t stream;
+  cudaStream_t stream;
   // Algorithm details
   int chunkSteps;
   int sliceSteps;
@@ -91,7 +91,7 @@ struct ncclTaskP2p {
 
 struct ncclCudaStreamList {
   struct ncclCudaStreamList *next;
-  hipStream_t stream;
+  cudaStream_t stream;
 };
 
 struct ncclTasks {
@@ -111,7 +111,7 @@ struct ncclTasks {
   // Keep track of the number of user streams
   int numStreams;
   // The most recent user stream. Ignored if streams==nullptr
-  hipStream_t streamRecent;
+  cudaStream_t streamRecent;
   // The graph capturing all user streams or invalid if none. Thus we restrict the
   // user that all streams must be captured in the same graph or not captured
   // at all. Technically we could probably relax this, but that would mean

projects/rccl/src/include/nvmlwrap.h

@@ -107,6 +107,75 @@ typedef enum nvmlGpuP2PCapsIndex_enum
     NVML_P2P_CAPS_INDEX_UNKNOWN
 } nvmlGpuP2PCapsIndex_t;
 
+/**
+ * Represents the type for sample value returned
+ */
+typedef enum nvmlValueType_enum
+{
+    NVML_VALUE_TYPE_DOUBLE = 0,
+    NVML_VALUE_TYPE_UNSIGNED_INT = 1,
+    NVML_VALUE_TYPE_UNSIGNED_LONG = 2,
+    NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3,
+    NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4,
+
+    // Keep this last
+    NVML_VALUE_TYPE_COUNT
+}nvmlValueType_t;
+
+
+/**
+ * Union to represent different types of Value
+ */
+typedef union nvmlValue_st
+{
+    double dVal;                    //!< If the value is double
+    unsigned int uiVal;             //!< If the value is unsigned int
+    unsigned long ulVal;            //!< If the value is unsigned long
+    unsigned long long ullVal;      //!< If the value is unsigned long long
+    signed long long sllVal;        //!< If the value is signed long long
+}nvmlValue_t;
+
+/**
+ * Field Identifiers.
+ *
+ * All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change.
+ */
+
+/* NVLink Speed */
+#define NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON 90  //!< Common NVLink Speed in MBps for active links
+#define NVML_FI_DEV_NVLINK_LINK_COUNT        91  //!< Number of NVLinks present on the device
+
+/**
+ * Remote device NVLink ID
+ *
+ * Link ID needs to be specified in the scopeId field in nvmlFieldValue_t.
+ */
+#define NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID     146 //!< Remote device NVLink ID
+
+/**
+ * NVSwitch: connected NVLink count
+ */
+#define NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT   147  //!< Number of NVLinks connected to NVSwitch
+
+#define NVML_FI_DEV_NVLINK_GET_SPEED                  164
+#define NVML_FI_DEV_NVLINK_GET_STATE                  165
+#define NVML_FI_DEV_NVLINK_GET_VERSION                166
+#define NVML_FI_MAX 167 //!< One greater than the largest field ID defined above
+
+/**
+ * Information for a Field Value Sample
+ */
+typedef struct nvmlFieldValue_st
+{
+    unsigned int fieldId;       //!< ID of the NVML field to retrieve. This must be set before any call that uses this struct. See the constants starting with NVML_FI_ above.
+    unsigned int scopeId;       //!< Scope ID can represent data used by NVML depending on fieldId's context. For example, for NVLink throughput counter data, scopeId can represent linkId.
+    long long timestamp;        //!< CPU Timestamp of this value in microseconds since 1970
+    long long latencyUsec;      //!< How long this field value took to update (in usec) within NVML. This may be averaged across several fields that are serviced by the same driver call.
+    nvmlValueType_t valueType;  //!< Type of the value stored in value
+    nvmlReturn_t nvmlReturn;    //!< Return code for retrieving this value. This must be checked before looking at value, as value is undefined if nvmlReturn != NVML_SUCCESS
+    nvmlValue_t value;          //!< Value for this field. This is only valid if nvmlReturn == NVML_SUCCESS
+} nvmlFieldValue_t;
+
 /* End of nvml.h */
 #endif // NCCL_NVML_DIRECT
 
@@ -135,4 +204,6 @@ ncclResult_t ncclNvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned
 ncclResult_t ncclNvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, nvmlNvLinkCapability_t capability, unsigned int *capResult);
 ncclResult_t ncclNvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor);
 ncclResult_t ncclNvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, nvmlGpuP2PCapsIndex_t p2pIndex, nvmlGpuP2PStatus_t* p2pStatus);
+ncclResult_t ncclNvmlDeviceGetFieldValues(nvmlDevice_t device, int valuesCount, nvmlFieldValue_t *values);
+
 #endif // End include guard

projects/rccl/src/include/proxy.h

@@ -126,7 +126,7 @@ struct ncclProxySharedP2p {
   int size;
   char* cudaBuff;
   char* hostBuff;
-  hipIpcMemHandle_t ipc;
+  cudaIpcMemHandle_t ipc;
   struct ncclProxyArgs* proxyAppend[MAXCHANNELS]; // Separate send and recv
 };
 
@@ -172,7 +172,7 @@ struct ncclProxyState {
   pthread_t thread;
   struct ncclSocket* listenSock;
   int stop;
-  hipCtx_t cudaCtx;
+  CUcontext cudaCtx;
   int safeAbortFlag;
 
   // Used by main thread

projects/rccl/src/include/strongstream.h

@@ -18,11 +18,11 @@
 struct ncclCudaGraph {
 #if CUDART_VERSION >= 11030
   cudaGraph_t graph;
-  uint64_t graphId;
+  unsigned long long graphId;
 #endif
 };
 
-inline struct ncclCudaGraph ncclCudaGraphNull() {
+inline struct ncclCudaGraph ncclCudaGraphNone() {
   struct ncclCudaGraph tmp;
   #if CUDART_VERSION >= 11030
     tmp.graph = nullptr;
@@ -47,9 +47,8 @@ inline bool ncclCudaGraphSame(struct ncclCudaGraph a, struct ncclCudaGraph b) {
   #endif
 }
 
-ncclResult_t ncclCudaGetCapturingGraph(struct ncclCudaGraph* graph, hipStream_t stream);
-ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t fn, void* arg);
-
+ncclResult_t ncclCudaGetCapturingGraph(struct ncclCudaGraph* graph, cudaStream_t stream);
+ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, cudaHostFn_t fn, void* arg);
 
 /* ncclStrongStream: An abstraction over CUDA streams that do not lose their
  * identity while being captured. Regular streams have the deficiency that the
@@ -58,37 +57,30 @@ ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t
  * streams unfit for the use of serializing access to a persistent resource.
  * Strong streams have been introduced to address this need.
  *
- * Constraints of using strong streams:
+ * - All updates to a strong stream must be enclosed by a Acquire/Release pair.
  *
- * - Operations that enqueue work to the strong stream need to be enclosed by
- *   ncclStrongStream[Acquire/Release] pairs. Acquire/release act like fences,
- *   the strong stream is not stateful so there is no harm in redundant acquire
- *   or releases.
+ * - The Acquire, Release, and all updates take a ncclCudaGraph parameter
+ *   indicating the currently capturing graph (or none). This parameter must be
+ *   the same for the entire sequence of {Acquire; ...; Release}.
  *
  * - An {Acquire; ...; Release} sequence must not be concurrent with any
  *   other operations against the strong stream including graph launches which
  *   reference this stream.
- *
- * - All strong stream functions take a "graph" parameter which must reference
- *   the currently capturing graph, or null if none.
  */
 struct ncclStrongStream;
 
 ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss);
 ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss);
 
-// Has this strong stream ever been captured in a graph.
-bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss);
-
 // Acquire-fence the strong stream.
 ncclResult_t ncclStrongStreamAcquire(
   struct ncclCudaGraph graph, struct ncclStrongStream* ss
 );
 
 // Acquire-fence the strong stream assuming no graph is capturing. This permits
-// the caller to enqueue directly to the `ss->stream` member using native CUDA
-// calls. Strong stream must be released via:
-//   ncclStrongStreamRelease(ncclCudaGraphNull(), graphRefs, ss);
+// the caller to enqueue directly to the `ss->cudaStream` member using native CUDA
+// calls. Strong stream still must be released via:
+//   ncclStrongStreamRelease(ncclCudaGraphNone(), ss);
 ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss);
 
 // Release-fence of the strong stream.
@@ -97,24 +89,25 @@ ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStro
 // Add a host launch to the stream.
 ncclResult_t ncclStrongStreamLaunchHost(
   struct ncclCudaGraph graph, struct ncclStrongStream* ss,
-  hipHostFn_t fn, void* arg
+  cudaHostFn_t fn, void* arg
 );
 // Add a kernel launch to the stream.
 ncclResult_t ncclStrongStreamLaunchKernel(
   struct ncclCudaGraph graph, struct ncclStrongStream* ss,
   void* fn, dim3 grid, dim3 block, void** args, size_t sharedMemBytes
 );
+
 // Cause `a` to wait for the current state `b`. Both `a` and `b` must be acquired.
 ncclResult_t ncclStrongStreamWaitStream(
   struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b
 );
 // `b` must be capturing within `graph`.
 ncclResult_t ncclStrongStreamWaitStream(
-  struct ncclCudaGraph graph, struct ncclStrongStream* a, hipStream_t b
+  struct ncclCudaGraph graph, struct ncclStrongStream* a, cudaStream_t b
 );
 // `a` must be capturing within `graph`.
 ncclResult_t ncclStrongStreamWaitStream(
-  struct ncclCudaGraph graph, hipStream_t a, struct ncclStrongStream* b
+  struct ncclCudaGraph graph, cudaStream_t a, struct ncclStrongStream* b
 );
 
 // Synchrnoization does not need the strong stream to be acquired.
@@ -122,21 +115,23 @@ ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct ncclStrongStreamGraph; // internal to ncclStrongStream
+
 struct ncclStrongStream {
-  hipStream_t stream;
-  hipEvent_t event;
-  #if CUDART_VERSION >= 11030
-  cudaGraphNode_t node; // null if never captured, otherwise never null again
-  uint64_t graphId:63, eventIsLagging:1;
-  #endif
+  // Used when not graph capturing.
+  cudaStream_t cudaStream;
+#if CUDART_VERSION >= 11030
+  // The event used to establish order between graphs and streams. During acquire
+  // this event is waited on, during release it is recorded to.
+  cudaEvent_t serialEvent;
+  // This stream ever appeared in a graph capture.
+  bool everCaptured;
+  // Tracks whether serialEvent needs to be recorded to upon Release().
+  bool serialEventNeedsRecord;
+  struct ncclStrongStreamGraph* graphHead;
+#else
+  cudaEvent_t scratchEvent;
+#endif
 };
 
-inline bool ncclStrongStreamEverCaptured(struct ncclStrongStream* ss) {
-  #if CUDART_VERSION >= 11030
-    return ss->node != nullptr;
-  #else
-    return false;
-  #endif
-}
-
 #endif

projects/rccl/src/init.cc

@@ -259,7 +259,7 @@ void ncclCommPushFree(struct ncclComm* comm, void* obj) {
 }
 
 static ncclResult_t ncclDestructorFnCudaFree(struct ncclDestructor* dtor) {
-  CUDACHECK(hipFree(dtor->obj));
+  CUDACHECK(cudaFree(dtor->obj));
   return ncclSuccess;
 }
 void ncclCommPushCudaFree(struct ncclComm* comm, void* obj) {
@@ -271,7 +271,7 @@ void ncclCommPushCudaFree(struct ncclComm* comm, void* obj) {
 }
 
 static ncclResult_t ncclDestructorFnCudaHostFree(struct ncclDestructor* dtor) {
-  CUDACHECK(hipHostFree(dtor->obj));
+  CUDACHECK(cudaFreeHost(dtor->obj));
   return ncclSuccess;
 }
 void ncclCommPushCudaHostFree(struct ncclComm* comm, void* obj) {
@@ -410,13 +410,13 @@ static ncclResult_t dmaBufSupported(struct ncclComm* comm) {
   if (ncclParamDmaBufEnable() == 0 || comm->ncclNet->regMrDmaBuf == NULL) return ncclInternalError;
 #if CUDA_VERSION >= 11070
   int flag = 0;
-  hipDevice_t dev;
+  CUdevice dev;
   int cudaDriverVersion;
-  CUCHECK(hipDriverGetVersion(&cudaDriverVersion));
+  CUCHECK(cuDriverGetVersion(&cudaDriverVersion));
   if (cudaDriverVersion < 11070) return ncclInternalError;
-  CUCHECK(hipDeviceGet(&dev, comm->cudaDev));
+  CUCHECK(cuDeviceGet(&dev, comm->cudaDev));
   // Query device to see if DMA-BUF support is available
-  (void) CUPFN(hipDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, dev));
+  (void) CUPFN(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, dev));
   if (flag == 0) return ncclInternalError;
   INFO(NCCL_INIT, "DMA-BUF is available on GPU device %d", comm->cudaDev);
   return ncclSuccess;
@@ -490,7 +490,7 @@ static ncclResult_t commAlloc(ncclComm_t* comret, int ndev, int rank, int virtua
   comm->doneEvent = doneEvent;
   comm->lastStream = nullptr;
   comm->virtualId = virtualId;
-  hipGetDevice(&comm->cudaDev);
+  cudaGetDevice(&comm->cudaDev);
   NCCLCHECK(getBusId(comm->cudaDev, &comm->busId));
   TRACE(NCCL_INIT,"comm %p rank %d nranks %d cudaDev %d busId %lx", comm, rank, ndev, comm->cudaDev, comm->busId);
 
@@ -540,7 +540,7 @@ static ncclResult_t devCommSetup(ncclComm_t comm) {
 
   int nRanks = comm->nRanks;
   struct ncclDevCommAndChannels *devCommAndChans, tmpCommAndChans;
-  NCCLCHECK(ncclCudaCallocAsync(&devCommAndChans, 1, comm->deviceStream.stream));
+  NCCLCHECK(ncclCudaCallocAsync(&devCommAndChans, 1, comm->deviceStream.cudaStream));
   ncclCommPushCudaFree(comm, devCommAndChans);
   comm->devComm = &devCommAndChans->comm;
   tmpCommAndChans.comm.rank = comm->rank;
@@ -587,7 +587,7 @@ static ncclResult_t devCommSetup(ncclComm_t comm) {
     tmpCommAndChans.channels[c].workFifoDone = &comm->workFifoDone[c];
 
     if (comm->channels[c].ring.userRanks != nullptr) {
-      NCCLCHECK(ncclCudaMemcpyAsync(tmpCommAndChans.channels[c].ring.userRanks, comm->channels[c].ring.userRanks, nRanks, comm->deviceStream.stream));
+      NCCLCHECK(ncclCudaMemcpyAsync(tmpCommAndChans.channels[c].ring.userRanks, comm->channels[c].ring.userRanks, nRanks, comm->deviceStream.cudaStream));
     }
   }
 
@@ -608,10 +608,9 @@ static ncclResult_t devCommSetup(ncclComm_t comm) {
   NCCLCHECK(ncclCudaCalloc(&tmpCommAndChans.comm.devProf, MAXCHANNELS*PROFILE_NUM_LAUNCHES), comm->sideStream);
 #endif
 
-  NCCLCHECK(ncclCudaMemcpyAsync(devCommAndChans, &tmpCommAndChans, 1, comm->deviceStream.stream));
-  CUDACHECK(hipStreamSynchronize(comm->deviceStream.stream));
-  NCCLCHECK(ncclStrongStreamRelease(ncclCudaGraphNull(), &comm->deviceStream));
-
+  NCCLCHECK(ncclCudaMemcpyAsync(devCommAndChans, &tmpCommAndChans, 1, comm->deviceStream.cudaStream));
+  CUDACHECK(cudaStreamSynchronize(comm->deviceStream.cudaStream));
+  NCCLCHECK(ncclStrongStreamRelease(ncclCudaGraphNone(), &comm->deviceStream));
   return ncclSuccess;
 }
 
@@ -635,7 +634,7 @@ static void showVersion() {
 static ncclResult_t fillInfo(struct ncclComm* comm, struct ncclPeerInfo* info, uint64_t commHash) {
   info->rank = comm->rank;
   info->virtualId = comm->virtualId;
-  CUDACHECK(hipGetDevice(&info->cudaDev));
+  CUDACHECK(cudaGetDevice(&info->cudaDev));
   info->hostHash=getHostHash()+commHash;
   info->pidHash=getPidHash()+commHash;
 
@@ -942,7 +941,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
 
   comm->nChannels = (comm->topo->nodes[GPU].count != comm->topo->nRanks && comm->topo->nodes[NET].count)
     ? std::min(treeGraph.nChannels, ringGraph.nChannels) : ringGraph.nChannels;
-  NCCLCHECK(ncclTopoPreset(comm, &treeGraph, &ringGraph, &allGather3Data[rank].topoRanks));
+  NCCLCHECK(ncclTopoPreset(comm, &treeGraph, &ringGraph, &collNetGraph, &allGather3Data[rank].topoRanks));
 
   NCCLCHECK(bootstrapAllGather(comm->bootstrap, allGather3Data, sizeof(*allGather3Data)));
 
@@ -1288,13 +1287,13 @@ collnet_cleanup:
       for (int c=0; c<comm->p2pnChannelsPerPeer; c++) {
         NCCLCHECK(ncclChannelCompute(comm, peer, c, ncclFuncSend, &channelId));
         if (comm->channels[channelId].peers[peer].send[1].connected == 0) {
-          comm->connectSend[peer] |= (1<<channelId);
+          comm->connectSend[peer] |= (1UL<<channelId);
         }
       }
       for (int c=0; c<comm->p2pnChannelsPerPeer; c++) {
         NCCLCHECK(ncclChannelCompute(comm, peer, c, ncclFuncRecv, &channelId));
         if (comm->channels[channelId].peers[peer].recv[1].connected == 0) {
-          comm->connectRecv[peer] |= (1<<channelId);
+          comm->connectRecv[peer] |= (1UL<<channelId);
         }
       }
     }
@@ -1415,7 +1414,7 @@ static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) {
   int virtualId = job->virtualId;
   ncclResult_t res = ncclSuccess;
 
-  CUDACHECK(hipSetDevice(cudaDev));
+  CUDACHECK(cudaSetDevice(cudaDev));
   // Set the maximum kernel stack size of all kernels to avoid
   // a CUDA memory reconfig on load (c.f. NVSHMEM issue)
   if (maxLocalSizeBytes > 0 && ncclParamSetStackSize() == 1) {
@@ -1471,7 +1470,7 @@ static ncclResult_t ncclCommInitRankDev(ncclComm_t* newcomm, int nranks, ncclUni
 
   memset(allocTracker+cudaDev, 0, sizeof(struct allocationTracker));
   // Make sure the CUDA runtime is initialized.
-  CUDACHECKGOTO(hipFree(NULL), res, fail);
+  CUDACHECKGOTO(cudaFree(NULL), res, fail);
 
   NCCLCHECKGOTO(PtrCheck(newcomm, "CommInitRank", "newcomm"), res, fail);
   if (nranks < 1 || myrank < 0 || myrank >= nranks) {
@@ -1512,7 +1511,7 @@ ncclResult_t ncclCommInitRank(ncclComm_t* newcomm, int nranks, ncclUniqueId comm
   if (ncclParamDmaBufEnable()) rocmLibraryInit();
 
   int cudaDev;
-  CUDACHECK(hipGetDevice(&cudaDev));
+  CUDACHECK(cudaGetDevice(&cudaDev));
   NCCLCHECK(ncclCommInitRankDev(newcomm, nranks, commId, myrank, cudaDev, NULL, -1));
   return ncclSuccess;
 }
@@ -1543,7 +1542,7 @@ ncclResult_t ncclCommInitAll(ncclComm_t* comms, int ndev, const int* devlist) {
     goto fail;
   }
 
-  CUDACHECKGOTO(hipGetDeviceCount(&totalnDev), ret, fail);
+  CUDACHECKGOTO(cudaGetDeviceCount(&totalnDev), ret, fail);
   if (devlist) {
     NCCLCHECKGOTO(ncclCalloc(&gpuFlags, totalnDev), ret, fail);
     for (int i = 0; i < ndev; ++i) {
@@ -1562,6 +1561,7 @@ ncclResult_t ncclCommInitAll(ncclComm_t* comms, int ndev, const int* devlist) {
       gpuFlags[devlist[i]] = 1;
     }
     free(gpuFlags);
+    gpuFlags = nullptr;
   }
 
   ncclUniqueId uniqueId;
@@ -1573,11 +1573,9 @@ ncclResult_t ncclCommInitAll(ncclComm_t* comms, int ndev, const int* devlist) {
   }
   NCCLCHECKGOTO(ncclGroupEnd(), ret, fail);
 
-exit:
-  return ret;
 fail:
-  if (gpuFlags) free(gpuFlags);
-  goto exit;
+  free(gpuFlags);
+  return ret;
 }
 
 ncclResult_t ncclCommSetAsyncError(ncclComm_t comm, ncclResult_t nextState) {
@@ -1628,7 +1626,7 @@ ncclResult_t ncclCommInitRankConfig(ncclComm_t *newcomm, int nranks, ncclUniqueI
   if (blockingEnv == 1) internalConfigPtr->blocking = blockingEnv;
 
   if (ncclParamDmaBufEnable()) (void) rocmLibraryInit();
-  CUDACHECKGOTO(hipGetDevice(&cudaDev), ret, exit);
+  CUDACHECKGOTO(cudaGetDevice(&cudaDev), ret, exit);
   NCCLCHECKGOTO(ncclCommInitRankDev(newcomm, nranks, commId, myrank, cudaDev, internalConfigPtr, -1), ret, fail);
 
 exit:
@@ -1648,13 +1646,13 @@ static ncclResult_t commDestroySync(struct ncclAsyncJob* job_) {
 #ifdef ENABLE_TRACE
   int rank = comm->rank;
 #endif
-  CUDACHECK(hipGetDevice(&savedDevice));
+  CUDACHECK(cudaGetDevice(&savedDevice));
   int commDevice = comm->cudaDev;
   ncclResult_t ret;
 
-  CUDACHECKGOTO(hipGetDevice(&savedDevice), ret, fail);
+  CUDACHECKGOTO(cudaGetDevice(&savedDevice), ret, fail);
   if (savedDevice != commDevice) {
-    CUDACHECKGOTO(hipSetDevice(commDevice), ret, fail);
+    CUDACHECKGOTO(cudaSetDevice(commDevice), ret, fail);
   }
 
   TRACE(NCCL_INIT, "Destroying comm %p rank %d abortFlag %d asyncResult %d", comm, comm->rank, *comm->abortFlag, comm->asyncResult);
@@ -1670,7 +1668,7 @@ static ncclResult_t commDestroySync(struct ncclAsyncJob* job_) {
   }
 
   if (savedDevice != commDevice) {
-    CUDACHECKGOTO(hipSetDevice(savedDevice), ret, fail);
+    CUDACHECKGOTO(cudaSetDevice(savedDevice), ret, fail);
   }
 
 exit:
@@ -1683,15 +1681,16 @@ static ncclResult_t commCleanup(ncclComm_t comm) {
   int savedDevice;
   int commDevice = comm->cudaDev;
 
-  CUDACHECK(hipGetDevice(&savedDevice));
+  CUDACHECK(cudaGetDevice(&savedDevice));
   if (savedDevice != commDevice) {
-    CUDACHECK(hipSetDevice(commDevice));
+    CUDACHECK(cudaSetDevice(commDevice));
   }
 
   NCCLCHECK(commFree(comm));
 
-  if (savedDevice != commDevice)
-    CUDACHECK(hipSetDevice(savedDevice));
+  if (savedDevice != commDevice) {
+    CUDACHECK(cudaSetDevice(savedDevice));
+  }
 
 #if defined(ENABLE_NPKIT)
   // Dump NPKit events and shutdown

projects/rccl/src/misc/argcheck.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2019-2022, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,16 +8,16 @@
 #include "comm.h"
 
 static ncclResult_t CudaPtrCheck(const void* pointer, struct ncclComm* comm, const char* ptrname, const char* opname) {
-  hipPointerAttribute_t attr;
-  hipError_t err = hipPointerGetAttributes(&attr, pointer);
-  if (err != hipSuccess || attr.devicePointer == NULL) {
+  cudaPointerAttributes attr;
+  cudaError_t err = cudaPointerGetAttributes(&attr, pointer);
+  if (err != cudaSuccess || attr.devicePointer == NULL) {
     WARN("%s : %s %p is not a valid pointer", opname, ptrname, pointer);
     return ncclInvalidArgument;
   }
 #if CUDART_VERSION >= 10000
-  if (attr.type == hipMemoryTypeDevice && attr.device != comm->cudaDev) {
+  if (attr.type == cudaMemoryTypeDevice && attr.device != comm->cudaDev) {
 #else
-  if (attr.memoryType == hipMemoryTypeDevice && attr.device != comm->cudaDev) {
+  if (attr.memoryType == cudaMemoryTypeDevice && attr.device != comm->cudaDev) {
 #endif
     WARN("%s : %s allocated on device %d mismatchs with NCCL device %d", opname, ptrname, attr.device, comm->cudaDev);
     return ncclInvalidArgument;
@@ -44,7 +43,7 @@ ncclResult_t ArgsCheck(struct ncclInfo* info) {
     WARN("%s : invalid type %d", info->opName, info->datatype);
     return ncclInvalidArgument;
   }
-  // Type is OK, compute nbytes. Convert Allgather/Broadcast/P2P/AllToAllPivot calls to chars.
+  // Type is OK, compute nbytes. Convert Allgather/Broadcast/P2P calls to chars.
   NCCLCHECK(ncclInfoSetDerived(info, info->comm->nRanks));
 
   if (info->op < 0 || ncclMaxRedOp < info->op) {

projects/rccl/src/misc/cudawrap.cc

@@ -38,7 +38,7 @@ DECLARE_CUDA_PFN(cuGetProcAddress, 11030);
 #define CUDA_DRIVER_MIN_VERSION 11030
 
 static void *cudaLib;
-static int cudaDriverVersion;
+int ncclCudaDriverVersionCache = -1;
 
 #if CUDART_VERSION >= 11030
 /*
@@ -107,16 +107,17 @@ static void initOnceFunc() {
     goto error;
   }
 
-  res = pfn_cuDriverGetVersion(&cudaDriverVersion);
+  int driverVersion;
+  res = pfn_cuDriverGetVersion(&driverVersion);
   if (res != 0) {
     WARN("cuDriverGetVersion failed with %d", res);
     goto error;
   }
 
-  INFO(NCCL_INIT, "cudaDriverVersion %d", cudaDriverVersion);
+  INFO(NCCL_INIT, "cudaDriverVersion %d", driverVersion);
 
-  if (cudaDriverVersion < CUDA_DRIVER_MIN_VERSION) {
-    // WARN("CUDA Driver version found is %d. Minimum requirement is %d", cudaDriverVersion, CUDA_DRIVER_MIN_VERSION);
+  if (driverVersion < CUDA_DRIVER_MIN_VERSION) {
+    // WARN("CUDA Driver version found is %d. Minimum requirement is %d", driverVersion, CUDA_DRIVER_MIN_VERSION);
     // Silently ignore version check mismatch for backwards compatibility
     goto error;
   }
@@ -148,7 +149,7 @@ error:
   return;
 }
 
-ncclResult_t cudaLibraryInit() {
+ncclResult_t ncclCudaLibraryInit() {
   pthread_once(&initOnceControl, initOnceFunc);
   return initResult;
 }

projects/rccl/src/misc/nvmlwrap.cc

@@ -38,6 +38,7 @@ namespace {
   NCCL_NVML_FN(nvmlDeviceGetNvLinkCapability, nvmlReturn_t, (nvmlDevice_t device, unsigned int link, nvmlNvLinkCapability_t capability, unsigned int *capResult))
   NCCL_NVML_FN(nvmlDeviceGetCudaComputeCapability, nvmlReturn_t, (nvmlDevice_t device, int* major, int* minor))
   NCCL_NVML_FN(nvmlDeviceGetP2PStatus, nvmlReturn_t, (nvmlDevice_t device1, nvmlDevice_t device2, nvmlGpuP2PCapsIndex_t p2pIndex, nvmlGpuP2PStatus_t* p2pStatus))
+  NCCL_NVML_FN(nvmlDeviceGetFieldValues, nvmlReturn_t, (nvmlDevice_t device, int valuesCount, nvmlFieldValue_t *values))
 
   std::mutex lock; // NVML has had some thread safety bugs
   bool initialized = false;
@@ -80,7 +81,8 @@ ncclResult_t ncclNvmlEnsureInitialized() {
       {(void**)&pfn_nvmlDeviceGetNvLinkRemotePciInfo, "nvmlDeviceGetNvLinkRemotePciInfo"},
       {(void**)&pfn_nvmlDeviceGetNvLinkCapability, "nvmlDeviceGetNvLinkCapability"},
       {(void**)&pfn_nvmlDeviceGetCudaComputeCapability, "nvmlDeviceGetCudaComputeCapability"},
-      {(void**)&pfn_nvmlDeviceGetP2PStatus, "nvmlDeviceGetP2PStatus"}
+      {(void**)&pfn_nvmlDeviceGetP2PStatus, "nvmlDeviceGetP2PStatus"},
+      {(void**)&pfn_nvmlDeviceGetFieldValues, "nvmlDeviceGetFieldValues"}
     };
     for(Symbol sym: symbols) {
       *sym.ppfn = dlsym(libhandle, sym.name);
@@ -260,3 +262,10 @@ ncclResult_t ncclNvmlDeviceGetP2PStatus(
   }
   return ncclSuccess;
 }
+
+ncclResult_t ncclNvmlDeviceGetFieldValues(nvmlDevice_t device, int valuesCount, nvmlFieldValue_t *values) {
+  NCCLCHECK(ncclNvmlEnsureInitialized());
+  std::lock_guard<std::mutex> locked(lock);
+  NVMLTRY(nvmlDeviceGetFieldValues, device, valuesCount, values);
+  return ncclSuccess;
+}

projects/rccl/src/misc/shmutils.cc

@@ -60,15 +60,15 @@ ncclResult_t ncclShmOpen(char* shmPath, const int shmSize, void** shmPtr, void**
 
   NCCLCHECKGOTO(ncclShmSetup(shmPath, shmSize, &fd, &ptr, create), res, sysError);
   if (devShmPtr) {
-    CUDACHECKGOTO(hipHostRegister(ptr, shmSize, hipHostRegisterMapped), res, hipError_t);
-    CUDACHECKGOTO(hipHostGetDevicePointer(devShmPtr, ptr, 0), res, hipError_t);
+    CUDACHECKGOTO(cudaHostRegister(ptr, shmSize, cudaHostRegisterMapped), res, cudaError);
+    CUDACHECKGOTO(cudaHostGetDevicePointer(devShmPtr, ptr, 0), res, cudaError);
   }
 
   *shmPtr = ptr;
   return ncclSuccess;
 sysError:
   WARN("Error while %s shared memory segment %s (size %d)", create ? "creating" : "attaching to", shmPath, shmSize);
-hipError_t:
+cudaError:
   if (fd != -1) close(fd);
   if (create) shm_unlink(shmPath);
   if (ptr != MAP_FAILED) munmap(ptr, shmSize);
@@ -83,7 +83,7 @@ ncclResult_t ncclShmUnlink(const char* shmPath) {
 
 ncclResult_t ncclShmClose(void* shmPtr, void* devShmPtr, const int shmSize) {
   if (shmPtr) {
-    if (devShmPtr) CUDACHECK(hipHostUnregister(shmPtr));
+    if (devShmPtr) CUDACHECK(cudaHostUnregister(shmPtr));
     if (munmap(shmPtr, shmSize) != 0) {
       WARN("munmap of shared memory failed");
       return ncclSystemError;

projects/rccl/src/misc/strongstream.cc

@@ -5,43 +5,71 @@
  ************************************************************************/
 
 #include "strongstream.h"
+#include "rocmwrap.h"
 #include "checks.h"
 #include "param.h"
 
+// Tracks the chain of graph nodes for a given graph captured identified by
+// its graph id. This state has to live for as long as captured work is being
+// submitted. CUDA doesn't have mechanism to inform us when the user ends capture
+// so the best we can do is get notified when the graph is destroyed.
+struct ncclStrongStreamGraph {
+  struct ncclStrongStreamGraph* next;
+  // Atomically exchanged to false by both the main thread or the graph destructor
+  // callback. The last to arrive deletes the node.
+  bool alive;
+  unsigned long long graphId;
+  // For each graph we track the "tip" of the chain of graph nodes. A linear
+  // chain would always have just one node at its tip, but since we have to merge
+  // in chains from other streams (via ncclStrongStreamWaitStream) some spots
+  // in the chain can be wider than a single node and thus need a list, so we
+  // maintain a dynamically sized array of tip nodes.
+  int tipCount, tipCapacity;
+  cudaGraphNode_t* tipNodes;
+};
+
+static void ncclStrongStreamGraphDelete(struct ncclStrongStreamGraph* g) {
+  free(g->tipNodes);
+  free(g);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 ncclResult_t ncclCudaGetCapturingGraph(
-    struct ncclCudaGraph* graph, hipStream_t stream
+    struct ncclCudaGraph* graph, cudaStream_t stream
   ) {
-  #if CUDART_VERSION >= 11030
-    thread_local int driver = -1;
-    if (driver == -1) {
-      CUDACHECK(cudaDriverGetVersion(&driver));
-    }
-    if (driver < 11030) {
+  #if CUDART_VERSION >= 10000 // cudaStreamGetCaptureInfo
+    int driver;
+    NCCLCHECK(ncclCudaDriverVersion(&driver));
+    if (CUDART_VERSION < 11030 || driver < 11030) {
       cudaStreamCaptureStatus status;
       unsigned long long gid;
-      graph->graph = nullptr;
       CUDACHECK(cudaStreamGetCaptureInfo(stream, &status, &gid));
+      #if CUDART_VERSION >= 11030
+        graph->graph = nullptr;
+        graph->graphId = ULLONG_MAX;
+      #endif
       if (status != cudaStreamCaptureStatusNone) {
-        WARN("The installed CUDA driver is older than the minimum version (R465) required for NCCL's CUDA Graphs support");
+        WARN("NCCL cannot be captured in a graph if either it wasn't built with CUDA runtime >= 11.3 or if the installed CUDA driver < R465.");
         return ncclInvalidUsage;
       }
     } else {
-      cudaStreamCaptureStatus status;
-      unsigned long long gid;
-      CUDACHECK(cudaStreamGetCaptureInfo_v2(stream, &status, &gid, &graph->graph, nullptr, nullptr));
-      if (status != cudaStreamCaptureStatusActive) {
-        graph->graph = nullptr;
-        gid = ULLONG_MAX;
-      }
-      graph->graphId = gid;
+      #if CUDART_VERSION >= 11030
+        cudaStreamCaptureStatus status;
+        unsigned long long gid;
+        CUDACHECK(cudaStreamGetCaptureInfo_v2(stream, &status, &gid, &graph->graph, nullptr, nullptr));
+        if (status != cudaStreamCaptureStatusActive) {
+          graph->graph = nullptr;
+          gid = ULLONG_MAX;
+        }
+        graph->graphId = gid;
+      #endif
     }
   #endif
   return ncclSuccess;
 }
 
-ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t fn, void* arg) {
+ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, cudaHostFn_t fn, void* arg) {
   #if CUDART_VERSION >= 11030
     cudaUserObject_t object;
     CUDACHECK(cudaUserObjectCreate(
@@ -58,52 +86,114 @@ ncclResult_t ncclCudaGraphAddDestructor(struct ncclCudaGraph graph, hipHostFn_t
 ////////////////////////////////////////////////////////////////////////////////
 
 ncclResult_t ncclStrongStreamConstruct(struct ncclStrongStream* ss) {
-  CUDACHECK(hipStreamCreateWithFlags(&ss->stream, hipStreamNonBlocking));
-  CUDACHECK(hipEventCreateWithFlags(&ss->event, hipEventDisableTiming));
+  CUDACHECK(cudaStreamCreateWithFlags(&ss->cudaStream, cudaStreamNonBlocking));
   #if CUDART_VERSION >= 11030
-    ss->node = nullptr;
-    ss->graphId = (1ull<<(8*sizeof(long long)-1))-1;
-    ss->eventIsLagging = 0;
+    CUDACHECK(cudaEventCreateWithFlags(&ss->serialEvent, cudaEventDisableTiming));
+    ss->everCaptured = false;
+    ss->serialEventNeedsRecord = false;
+    ss->graphHead = nullptr;
+  #else
+    CUDACHECK(cudaEventCreateWithFlags(&ss->scratchEvent, cudaEventDisableTiming));
   #endif
   return ncclSuccess;
 }
 
+static void graphDestructor(void* arg) {
+  struct ncclStrongStreamGraph* g = (struct ncclStrongStreamGraph*)arg;
+  if (false == __atomic_exchange_n(&g->alive, false, __ATOMIC_ACQ_REL)) {
+    // Last to arrive deletes list node.
+    ncclStrongStreamGraphDelete(g);
+  }
+}
+
 ncclResult_t ncclStrongStreamDestruct(struct ncclStrongStream* ss) {
+  CUDACHECK(cudaStreamDestroy(ss->cudaStream));
   #if CUDART_VERSION >= 11030
-    CUDACHECK(cudaEventDestroy(ss->event));
+    CUDACHECK(cudaEventDestroy(ss->serialEvent));
+    // Delete list of per-graph chains.
+    struct ncclStrongStreamGraph* g = ss->graphHead;
+    while (g != nullptr) {
+      struct ncclStrongStreamGraph* next = g->next;
+      if (false == __atomic_exchange_n(&g->alive, false, __ATOMIC_ACQ_REL)) {
+        // Last to arrive deletes list node.
+        ncclStrongStreamGraphDelete(g);
+      }
+      g = next;
+    }
+  #else
+    CUDACHECK(cudaEventDestroy(ss->scratchEvent));
   #endif
-  CUDACHECK(hipStreamDestroy(ss->stream));
   return ncclSuccess;
 }
 
 NCCL_PARAM(GraphMixingSupport, "GRAPH_MIXING_SUPPORT", 1)
 
+static void ensureTips(struct ncclStrongStreamGraph* g, int n) {
+  if (g->tipCapacity < n) {
+    g->tipNodes = (cudaGraphNode_t*)realloc(g->tipNodes, n*sizeof(cudaGraphNode_t));
+    g->tipCapacity = n;
+  }
+}
+
 ncclResult_t ncclStrongStreamAcquire(
     struct ncclCudaGraph graph, struct ncclStrongStream* ss
   ) {
   #if CUDART_VERSION >= 11030
     bool mixing = ncclParamGraphMixingSupport();
     if (graph.graph == nullptr) {
-      if (mixing && ncclStrongStreamEverCaptured(ss)) {
-        CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
-        ss->eventIsLagging = 0;
+      if (mixing && ss->everCaptured) {
+        CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
+        ss->serialEventNeedsRecord = false;
       }
     } else {
-      if (ss->graphId != graph.graphId) {
-        if (mixing && ss->eventIsLagging) {
-          // Can only be here if previous release was for uncaptured work that
-          // elided updating the event because no capture had yet occurred.
-          CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
-          CUDACHECK(cudaEventRecord(ss->event, ss->stream));
-        }
-        ss->graphId = graph.graphId;
-        ss->eventIsLagging = 0;
-        if (mixing) {
-          CUDACHECK(cudaGraphAddEventWaitNode(&ss->node, graph.graph, nullptr, 0, ss->event));
+      ss->everCaptured = true;
+      // Find the current graph in our list of graphs if it exists.
+      struct ncclStrongStreamGraph** pg = &ss->graphHead;
+      struct ncclStrongStreamGraph* g;
+      while (*pg != nullptr) {
+        g = *pg;
+        if (g->graphId == graph.graphId) {
+          // Move to front of list so that operations after acquire don't have to search the list.
+          *pg = g->next;
+          g->next = ss->graphHead;
+          ss->graphHead = g;
+          return ncclSuccess;
+        } else if (false == __atomic_load_n(&g->alive, __ATOMIC_ACQUIRE)) {
+          // Unrelated graph that has been destroyed. Remove and delete.
+          *pg = g->next;
+          ncclStrongStreamGraphDelete(g);
         } else {
-          CUDACHECK(cudaGraphAddEmptyNode(&ss->node, graph.graph, nullptr, 0));
+          pg = &g->next;
         }
       }
+
+      // This is a new graph so add to the list.
+      g = (struct ncclStrongStreamGraph*)malloc(sizeof(struct ncclStrongStreamGraph));
+      g->graphId = graph.graphId;
+      g->tipNodes = nullptr;
+      g->tipCapacity = 0;
+      g->tipCount = 0;
+      g->next = ss->graphHead;
+      ss->graphHead = g;
+      g->alive = true;
+      NCCLCHECK(ncclCudaGraphAddDestructor(graph, graphDestructor, (void*)g));
+
+      if (mixing && ss->serialEventNeedsRecord) {
+        // Can only be here if previous release was for uncaptured work that
+        // elided updating the event because no capture had yet occurred.
+        CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
+        CUDACHECK(cudaEventRecord(ss->serialEvent, ss->cudaStream));
+      }
+      ss->serialEventNeedsRecord = false;
+
+      // First node in the chain must be a wait on the serialEvent.
+      if (mixing) {
+        ensureTips(g, 1);
+        CUDACHECK(cudaGraphAddEventWaitNode(&g->tipNodes[0], graph.graph, nullptr, 0, ss->serialEvent));
+        g->tipCount = 1;
+      } else {
+        g->tipCount = 0;
+      }
     }
   #endif
   return ncclSuccess;
@@ -112,26 +202,38 @@ ncclResult_t ncclStrongStreamAcquire(
 ncclResult_t ncclStrongStreamAcquireUncaptured(struct ncclStrongStream* ss) {
   #if CUDART_VERSION >= 11030
     bool mixing = ncclParamGraphMixingSupport();
-    if (mixing && ncclStrongStreamEverCaptured(ss)) {
-      CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+    if (mixing && ss->everCaptured) {
+      CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
     }
-    ss->eventIsLagging = 1; // Assume the caller is going to add work to stream.
+    ss->serialEventNeedsRecord = true; // Assume the caller is going to add work to stream.
   #endif
   return ncclSuccess;
 }
 
+static ncclResult_t checkGraphId(struct ncclStrongStreamGraph* g, unsigned long long id) {
+  if (g == nullptr || g->graphId != id) {
+    WARN("Expected graph id=%llu was not at head of strong stream's internal list.", id);
+    return ncclInternalError;
+  }
+  return ncclSuccess;
+}
+
 ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStrongStream* ss) {
   #if CUDART_VERSION >= 11030
     bool mixing = ncclParamGraphMixingSupport();
-    if (mixing && ss->eventIsLagging) {
+    if (mixing && ss->serialEventNeedsRecord) {
       if (graph.graph == nullptr) {
-        if (ncclStrongStreamEverCaptured(ss)) {
-          CUDACHECK(cudaEventRecord(ss->event, ss->stream));
-          ss->eventIsLagging = 0;
+        if (ss->everCaptured) {
+          CUDACHECK(cudaEventRecord(ss->serialEvent, ss->cudaStream));
+          ss->serialEventNeedsRecord = false;
         }
       } else {
-        CUDACHECK(cudaGraphAddEventRecordNode(&ss->node, graph.graph, &ss->node, 1, ss->event));
-        ss->eventIsLagging = 0;
+        struct ncclStrongStreamGraph* g = ss->graphHead;
+        NCCLCHECK(checkGraphId(g, graph.graphId));
+        ensureTips(g, 1);
+        CUDACHECK(cudaGraphAddEventRecordNode(&g->tipNodes[0], graph.graph, g->tipNodes, g->tipCount, ss->serialEvent));
+        g->tipCount = 1;
+        ss->serialEventNeedsRecord = false;
       }
     }
   #endif
@@ -139,21 +241,24 @@ ncclResult_t ncclStrongStreamRelease(struct ncclCudaGraph graph, struct ncclStro
 }
 
 ncclResult_t ncclStrongStreamLaunchHost(
-    struct ncclCudaGraph graph, struct ncclStrongStream* ss, hipHostFn_t fn, void* arg
+    struct ncclCudaGraph graph, struct ncclStrongStream* ss, cudaHostFn_t fn, void* arg
   ) {
   #if CUDART_VERSION >= 11030
     if (graph.graph == nullptr) {
-      CUDACHECK(cudaLaunchHostFunc(ss->stream, fn, arg));
+      CUDACHECK(cudaLaunchHostFunc(ss->cudaStream, fn, arg));
     } else {
       cudaHostNodeParams p;
       p.fn = fn;
       p.userData = arg;
-      CUDACHECK(cudaGraphAddHostNode(&ss->node, graph.graph, &ss->node, 1, &p));
+      struct ncclStrongStreamGraph* g = ss->graphHead;
+      NCCLCHECK(checkGraphId(g, graph.graphId));
+      ensureTips(g, 1);
+      CUDACHECK(cudaGraphAddHostNode(&g->tipNodes[0], graph.graph, g->tipNodes, g->tipCount, &p));
+      g->tipCount = 1;
     }
-    ss->eventIsLagging = 1;
+    ss->serialEventNeedsRecord = true;
   #else
-    //CUDACHECK(hipLaunchHostFunc(ss->stream, fn, arg));
-    CUDACHECK(hipStreamAddCallback(ss->stream, (hipStreamCallback_t)fn, arg, 0));
+    CUDACHECK(hipStreamAddCallback(ss->cudaStream, (hipStreamCallback_t)fn, arg, 0));
   #endif
   return ncclSuccess;
 }
@@ -164,9 +269,8 @@ ncclResult_t ncclStrongStreamLaunchKernel(
   ) {
   #if CUDART_VERSION >= 11030
     if (graph.graph == nullptr) {
-      CUDACHECK(cudaLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->stream));
+      CUDACHECK(cudaLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->cudaStream));
     } else {
-      cudaGraphNode_t tip = ss->node;
       cudaKernelNodeParams p;
       p.func = fn;
       p.gridDim = grid;
@@ -174,100 +278,116 @@ ncclResult_t ncclStrongStreamLaunchKernel(
       p.kernelParams = args;
       p.sharedMemBytes = sharedMemBytes;
       p.extra = nullptr;
-      CUDACHECK(cudaGraphAddKernelNode(&ss->node, graph.graph, &tip, 1, &p));
+      struct ncclStrongStreamGraph* g = ss->graphHead;
+      NCCLCHECK(checkGraphId(g, graph.graphId));
+      ensureTips(g, 1);
+      CUDACHECK(cudaGraphAddKernelNode(&g->tipNodes[0], graph.graph, g->tipNodes, g->tipCount, &p));
+      g->tipCount = 1;
     }
-    ss->eventIsLagging = 1;
+    ss->serialEventNeedsRecord = true;
   #else
-    CUDACHECK(hipLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->stream));
+    CUDACHECK(cudaLaunchKernel(fn, grid, block, args, sharedMemBytes, ss->cudaStream));
   #endif
   return ncclSuccess;
 }
 
+// Merge node list `b` into list `a` but don't add duplicates.
+static void mergeTips(struct ncclStrongStreamGraph* a, cudaGraphNode_t const* bNodes, int bn) {
+  int an = a->tipCount;
+  ensureTips(a, an + bn);
+  for (int bi=0; bi < bn; bi++) {
+    for (int ai=0; ai < an; ai++) {
+      if (a->tipNodes[ai] == bNodes[bi]) goto next_b;
+    }
+    a->tipNodes[a->tipCount++] = bNodes[bi];
+  next_b:;
+  }
+}
+
 ncclResult_t ncclStrongStreamWaitStream(
     struct ncclCudaGraph graph, struct ncclStrongStream* a, struct ncclStrongStream* b
   ) {
   #if CUDART_VERSION >= 11030
     if (graph.graph == nullptr) {
-      if (b->eventIsLagging) {
-        b->eventIsLagging = 0;
-        CUDACHECK(cudaEventRecord(b->event, b->stream));
+      if (b->serialEventNeedsRecord) {
+        b->serialEventNeedsRecord = false;
+        CUDACHECK(cudaEventRecord(b->serialEvent, b->cudaStream));
       }
-      CUDACHECK(cudaStreamWaitEvent(a->stream, b->event, 0));
-      a->eventIsLagging = 1;
+      CUDACHECK(cudaStreamWaitEvent(a->cudaStream, b->serialEvent, 0));
     } else {
-      cudaGraphNode_t pair[2] = {a->node, b->node};
-      CUDACHECK(cudaGraphAddEmptyNode(&a->node, graph.graph, pair, 2));
+      struct ncclStrongStreamGraph* ag = a->graphHead;
+      NCCLCHECK(checkGraphId(ag, graph.graphId));
+      struct ncclStrongStreamGraph* bg = b->graphHead;
+      NCCLCHECK(checkGraphId(bg, graph.graphId));
+      mergeTips(ag, bg->tipNodes, bg->tipCount);
     }
+    a->serialEventNeedsRecord = true;
   #else
-    CUDACHECK(hipEventRecord(b->event, b->stream));
-    CUDACHECK(hipStreamWaitEvent(a->stream, b->event, 0));
+    CUDACHECK(cudaEventRecord(b->scratchEvent, b->cudaStream));
+    CUDACHECK(cudaStreamWaitEvent(a->cudaStream, b->scratchEvent, 0));
   #endif
   return ncclSuccess;
 }
 
 ncclResult_t ncclStrongStreamWaitStream(
-    struct ncclCudaGraph graph, struct ncclStrongStream* a, hipStream_t b
+    struct ncclCudaGraph graph, struct ncclStrongStream* a, cudaStream_t b
   ) {
   #if CUDART_VERSION >= 11030
     if (graph.graph == nullptr) {
-      CUDACHECK(cudaEventRecord(a->event, b));
-      CUDACHECK(cudaStreamWaitEvent(a->stream, a->event, 0));
-      // We used a->event to record b so it no longer reflects anything about a.
-      a->eventIsLagging = 1;
+      // It is ok to use a->serialEvent to record b since we'll be setting
+      // a->serialEventNeedsRecord so the event won't be considered accurate
+      // until re-recorded.
+      CUDACHECK(cudaEventRecord(a->serialEvent, b));
+      CUDACHECK(cudaStreamWaitEvent(a->cudaStream, a->serialEvent, 0));
     } else {
       cudaStreamCaptureStatus status;
-      unsigned long long gid1;
-      cudaGraphNode_t const* deps;
-      size_t depN = 0;
-      CUDACHECK(cudaStreamGetCaptureInfo_v2(b, &status, &gid1, nullptr, &deps, &depN));
-      if (status != cudaStreamCaptureStatusActive || graph.graphId != gid1) {
+      unsigned long long bGraphId;
+      cudaGraphNode_t const* bNodes;
+      size_t bCount = 0;
+      CUDACHECK(cudaStreamGetCaptureInfo_v2(b, &status, &bGraphId, nullptr, &bNodes, &bCount));
+      if (status != cudaStreamCaptureStatusActive || graph.graphId != bGraphId) {
         WARN("Stream is not being captured by the expected graph.");
         return ncclInvalidUsage;
       }
-      if (depN > 0 && (depN > 1 || deps[0] != a->node)) {
-        cudaGraphNode_t tie;
-        if (depN == 1) {
-          tie = deps[0];
-        } else {
-          CUDACHECK(cudaGraphAddEmptyNode(&tie, graph.graph, deps, depN));
-        }
-        cudaGraphNode_t pair[2] = {a->node, tie};
-        CUDACHECK(cudaGraphAddEmptyNode(&a->node, graph.graph, pair, 2));
-      }
-      // a->eventIsLagging doesn't change since we are just updating the
-      // dependencies of a->node.
+      struct ncclStrongStreamGraph* ag = a->graphHead;
+      NCCLCHECK(checkGraphId(ag, graph.graphId));
+      mergeTips(ag, bNodes, bCount);
     }
+    a->serialEventNeedsRecord = true;
   #else
-    CUDACHECK(hipEventRecord(a->event, b));
-    CUDACHECK(hipStreamWaitEvent(a->stream, a->event, 0));
+    CUDACHECK(cudaEventRecord(a->scratchEvent, b));
+    CUDACHECK(cudaStreamWaitEvent(a->cudaStream, a->scratchEvent, 0));
   #endif
   return ncclSuccess;
 }
 
 ncclResult_t ncclStrongStreamWaitStream(
-    struct ncclCudaGraph graph, hipStream_t a, struct ncclStrongStream* b
+    struct ncclCudaGraph graph, cudaStream_t a, struct ncclStrongStream* b
   ) {
   #if CUDART_VERSION >= 11030
     if (graph.graph == nullptr) {
-      if (b->eventIsLagging) {
-        b->eventIsLagging = 0;
-        CUDACHECK(cudaEventRecord(b->event, b->stream));
+      if (b->serialEventNeedsRecord) {
+        b->serialEventNeedsRecord = false;
+        CUDACHECK(cudaEventRecord(b->serialEvent, b->cudaStream));
       }
-      CUDACHECK(cudaStreamWaitEvent(a, b->event, 0));
+      CUDACHECK(cudaStreamWaitEvent(a, b->serialEvent, 0));
     } else {
-      CUDACHECK(cudaStreamUpdateCaptureDependencies(a, &b->node, 1, cudaStreamAddCaptureDependencies));
+      struct ncclStrongStreamGraph* bg = b->graphHead;
+      NCCLCHECK(checkGraphId(bg, graph.graphId));
+      CUDACHECK(cudaStreamUpdateCaptureDependencies(a, bg->tipNodes, bg->tipCount, cudaStreamAddCaptureDependencies));
     }
   #else
-    CUDACHECK(hipEventRecord(b->event, b->stream));
-    CUDACHECK(hipStreamWaitEvent(a, b->event, 0));
+    CUDACHECK(cudaEventRecord(b->scratchEvent, b->cudaStream));
+    CUDACHECK(cudaStreamWaitEvent(a, b->scratchEvent, 0));
   #endif
   return ncclSuccess;
 }
 
 ncclResult_t ncclStrongStreamSynchronize(struct ncclStrongStream* ss) {
   #if CUDART_VERSION >= 11030
-    CUDACHECK(cudaStreamWaitEvent(ss->stream, ss->event, 0));
+    CUDACHECK(cudaStreamWaitEvent(ss->cudaStream, ss->serialEvent, 0));
+    ss->serialEventNeedsRecord = false;
   #endif
-  CUDACHECK(hipStreamSynchronize(ss->stream));
+  CUDACHECK(cudaStreamSynchronize(ss->cudaStream));
   return ncclSuccess;
 }

projects/rccl/src/misc/utils.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
  * Copyright (c) 2016-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,17 +8,16 @@
 #include "core.h"
 
 #include "nvmlwrap.h"
-#include <hip/hip_runtime.h>
 
 #include <stdlib.h>
 
 // Get current Compute Capability
 int ncclCudaCompCap() {
   int cudaDev;
-  if (hipGetDevice(&cudaDev) != hipSuccess) return 0;
+  if (cudaGetDevice(&cudaDev) != cudaSuccess) return 0;
   int ccMajor, ccMinor;
-  if (hipDeviceGetAttribute(&ccMajor, hipDeviceAttributeComputeCapabilityMajor, cudaDev) != hipSuccess) return 0;
-  if (hipDeviceGetAttribute(&ccMinor, hipDeviceAttributeComputeCapabilityMinor, cudaDev) != hipSuccess) return 0;
+  if (cudaDeviceGetAttribute(&ccMajor, cudaDevAttrComputeCapabilityMajor, cudaDev) != cudaSuccess) return 0;
+  if (cudaDeviceGetAttribute(&ccMinor, cudaDevAttrComputeCapabilityMinor, cudaDev) != cudaSuccess) return 0;
   return ccMajor*10+ccMinor;
 }
 
@@ -51,7 +49,7 @@ ncclResult_t getBusId(int cudaDev, int64_t *busId) {
   // format. Still need to allocate proper space in case PCI domain goes
   // higher.
   char busIdStr[] = "00000000:00:00.0";
-  CUDACHECK(hipDeviceGetPCIBusId(busIdStr, sizeof(busIdStr), cudaDev));
+  CUDACHECK(cudaDeviceGetPCIBusId(busIdStr, sizeof(busIdStr), cudaDev));
   NCCLCHECK(busIdToInt64(busIdStr, busId));
   return ncclSuccess;
 }

projects/rccl/src/net.cc

@@ -337,7 +337,7 @@ ncclResult_t ncclGpuGdrSupport(struct ncclComm* comm, int* gdrSupport) {
     NCCLCHECKGOTO(ncclNetListen(comm, dev, &handle, &lComm), ret, cleanup1);
     NCCLWAITGOTO(ncclNetConnect(comm, dev, &handle, &sComm), sComm != NULL, comm->abortFlag, ret, cleanup2);
     NCCLWAITGOTO(ncclNetAccept(comm, lComm, &rComm), rComm != NULL, comm->abortFlag, ret, cleanup3);
-    CUDACHECKGOTO(hipMalloc(&gpuPtr, GPU_BUF_SIZE), ret, cleanup4);
+    CUDACHECKGOTO(cudaMalloc(&gpuPtr, GPU_BUF_SIZE), ret, cleanup4);
     if (ncclNetRegMr(comm, sComm, gpuPtr, GPU_BUF_SIZE, NCCL_PTR_CUDA, &mHandle) == ncclSuccess) {
       NCCLCHECK(ncclNetDeregMr(comm, sComm, mHandle));
       NCCLCHECK(ncclNetRegMr(comm, rComm, gpuPtr, GPU_BUF_SIZE, NCCL_PTR_CUDA, &mHandle));
@@ -345,7 +345,7 @@ ncclResult_t ncclGpuGdrSupport(struct ncclComm* comm, int* gdrSupport) {
       *gdrSupport = 1;
     }
     ncclDebugNoWarn = 0;
-    CUDACHECK(hipFree(gpuPtr));
+    CUDACHECK(cudaFree(gpuPtr));
 cleanup4:
     NCCLCHECK(ncclNetCloseRecv(comm, rComm));
 cleanup3:

projects/rccl/src/proxy.cc

@@ -661,7 +661,7 @@ void* ncclProxyProgress(void *comm_) {
   struct ncclComm* comm = (struct ncclComm*)comm_;
   if (ncclSetThreadContext(comm) != ncclSuccess) {
     WARN("[Proxy Progress] Failed to set CUDA context on device %d", comm->cudaDev);
-  } else if (hipSetDevice(comm->cudaDev) != hipSuccess) {
+  } else if (cudaSetDevice(comm->cudaDev) != cudaSuccess) {
     WARN("[Proxy Progress] Failed to set CUDA device %d", comm->cudaDev);
   }
   if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
@@ -1028,7 +1028,7 @@ void* ncclProxyService(void* _args) {
   if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
   if (ncclSetThreadContext(comm) != ncclSuccess) {
     WARN("[Proxy Service] Failed to set CUDA context on device %d", comm->cudaDev);
-  } else if (hipSetDevice(comm->cudaDev) != hipSuccess) {
+  } else if (cudaSetDevice(comm->cudaDev) != cudaSuccess) {
     WARN("[Proxy Service] Failed to set CUDA device %d", comm->cudaDev);
   }
   if (CPU_COUNT(&comm->cpuAffinity)) sched_setaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);
@@ -1056,8 +1056,8 @@ void* ncclProxyService(void* _args) {
   int asyncOpCount = 0;
   while ((stop == 0 || (stop == 1 && npeers > 0)) && *comm->abortFlag == 0) {
     /* never let proxy service thread blocks in poll, or it cannot receive abortFlag. */
-    if (int error = poll(pollfds, NCCL_MAX_LOCAL_RANKS+1, asyncOpCount ? 0 : 500) < 0) {
-      WARN("[Proxy Service] Poll failed with error %d", error);
+    if (poll(pollfds, NCCL_MAX_LOCAL_RANKS+1, asyncOpCount ? 0 : 500) < 0) {
+      WARN("[Proxy Service] Poll failed: %s\n", strerror(errno));
       return NULL;
     }
     if (pollfds[NCCL_MAX_LOCAL_RANKS].revents) {
@@ -1186,7 +1186,7 @@ ncclResult_t ncclProxyDestroy(struct ncclComm* comm) {
           NCCLCHECK(ncclShmClose(state->proxyOps[i].pool, NULL, sizeof(struct ncclProxyOpsPool)));
         }
         if (state->sharedDevMems[i]) {
-          CUDACHECK(hipIpcCloseMemHandle(state->sharedDevMems[i]));
+          CUDACHECK(cudaIpcCloseMemHandle(state->sharedDevMems[i]));
         }
         int type = ncclProxyMsgClose;
         if (*comm->abortFlag == 0) NCCLCHECK(ncclSocketSend(state->peerSocks+i, &type, sizeof(int)));

projects/rccl/src/transport.cc

@@ -54,7 +54,7 @@ static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph*
 ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex) {
   TRACE(NCCL_INIT, "nsend %d nrecv %d", nsend, nrecv);
   struct ncclChannel* channel = &comm->channels[channelId];
-  uint32_t mask = 1 << channelId;
+  uint64_t mask = 1UL << channel->id;
   for (int i=0; i<nrecv; i++) {
     int peer = peerRecv[i];
     if (peer == -1 || peer >= comm->nRanks || peer == comm->rank || channel->peers[peer].recv[connIndex].connected) continue;
@@ -85,15 +85,15 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
     int bootstrapTag = (i<<8) + (graph ? graph->id+1 : 0);
     int recvPeer = (comm->rank - i + comm->nRanks) % comm->nRanks;
     int sendPeer = (comm->rank + i) % comm->nRanks;
-    uint32_t recvMask = comm->connectRecv[recvPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)];
-    uint32_t sendMask = comm->connectSend[sendPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)];
+    uint64_t recvMask = comm->connectRecv[recvPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)];
+    uint64_t sendMask = comm->connectSend[sendPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)];
 
     struct ncclConnect* recvData = data;
     int sendChannels = 0, recvChannels = 0;
     int type;
     TIME_START(0);
     for (int c=0; c<MAXCHANNELS; c++) {
-      if (recvMask & (1<<c)) {
+      if (recvMask & (1UL<<c)) {
         NCCLCHECK(selectTransport<0>(comm, graph, recvData+recvChannels++, c, recvPeer, connIndex, &type));
         if (type > highestType) highestType = type;
       }
@@ -102,7 +102,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
     TIME_START(1);
     struct ncclConnect* sendData = recvData+recvChannels;
     for (int c=0; c<MAXCHANNELS; c++) {
-      if (sendMask & (1<<c)) {
+      if (sendMask & (1UL<<c)) {
         NCCLCHECK(selectTransport<1>(comm, graph, sendData+sendChannels++, c, sendPeer, connIndex, &type));
         if (type > highestType) highestType = type;
       }
@@ -127,30 +127,28 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
 
     TIME_START(3);
     for (int c=0; c<MAXCHANNELS; c++) {
-      if (sendMask & (1<<c)) {
+      if (sendMask & (1UL<<c)) {
         struct ncclConnector* conn = comm->channels[c].peers[sendPeer].send + connIndex;
         NCCLCHECK(conn->transportComm->connect(comm, sendData++, 1, comm->rank, conn));
         conn->connected = 1;
-
-        CUDACHECK(hipMemcpyAsync(&comm->channels[c].devPeers[sendPeer].send[connIndex], &conn->conn, sizeof(struct ncclConnInfo), hipMemcpyHostToDevice, comm->sideStream));
-        CUDACHECK(hipMemcpyAsync(&comm->channels[c].devPeers[sendPeer].send[connIndex], &conn->conn, sizeof(struct ncclConnInfo), hipMemcpyHostToDevice, comm->sideStream));
+        CUDACHECK(cudaMemcpyAsync(&comm->channels[c].devPeers[sendPeer].send[connIndex], &conn->conn, sizeof(struct ncclConnInfo), cudaMemcpyHostToDevice, comm->sideStream));
+        CUDACHECK(cudaMemcpyAsync(&comm->channels[c].devPeers[sendPeer].send[connIndex], &conn->conn, sizeof(struct ncclConnInfo), cudaMemcpyHostToDevice, comm->sideStream));
       }
     }
     TIME_STOP(3);
     TIME_START(4);
     for (int c=0; c<MAXCHANNELS; c++) {
-      if (recvMask & (1<<c)) {
+      if (recvMask & (1UL<<c)) {
         struct ncclConnector* conn = comm->channels[c].peers[recvPeer].recv + connIndex;
         NCCLCHECK(conn->transportComm->connect(comm, recvData++, 1, comm->rank, conn));
         conn->connected = 1;
-
-        CUDACHECK(hipMemcpyAsync(&comm->channels[c].devPeers[recvPeer].recv[connIndex], &conn->conn, sizeof(struct ncclConnInfo), hipMemcpyHostToDevice, comm->sideStream));
+        CUDACHECK(cudaMemcpyAsync(&comm->channels[c].devPeers[recvPeer].recv[connIndex], &conn->conn, sizeof(struct ncclConnInfo), cudaMemcpyHostToDevice, comm->sideStream));
       }
     }
     TIME_STOP(4);
-    comm->connectRecv[recvPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)] = comm->connectSend[sendPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)] = 0;
+    comm->connectRecv[recvPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)] = comm->connectSend[sendPeer+comm->nRanks*(connIndex == NCCL_CONN_IDX_P2P_NET ? NCCL_CONN_IDX_P2P_NET : 0)] = 0UL;
   }
-  CUDACHECK(hipStreamSynchronize(comm->sideStream));
+  CUDACHECK(cudaStreamSynchronize(comm->sideStream));
   if (highestTransportType != NULL) *highestTransportType = highestType;
   TIME_PRINT("P2P Setup/Connect");
   return ncclSuccess;
@@ -225,7 +223,7 @@ int ncclTransportCollNetSetup(struct ncclComm* comm, struct ncclTopoGraph* collN
     NCCLCHECKGOTO(transportComm->connect(comm, masterConnects, nMasters, rankInCollNet, conn), res, cleanup);
     struct ncclDevChannelPeer* devRoot = channel->devPeers+nranks;
     struct ncclConnInfo* devConnInfo = (type == collNetRecv) ? devRoot->recv+type : devRoot->send+type;
-    CUDACHECKGOTO(hipMemcpy(devConnInfo, &conn->conn, sizeof(struct ncclConnInfo), hipMemcpyHostToDevice), res, cleanup);
+    CUDACHECKGOTO(cudaMemcpy(devConnInfo, &conn->conn, sizeof(struct ncclConnInfo), cudaMemcpyHostToDevice), res, cleanup);
   }
   // recv side sends connect info to send side
   if (isMaster && type == collNetRecv) {

projects/rccl/src/transport/coll_net.cc

@@ -123,6 +123,7 @@ struct recvResources {
   int netDev;
   int useGdr;
   int useDmaBuf;
+  int needFlush;
   uint64_t* gdcSync;
   uint64_t* gdcFlush;
   void* gdrDesc;
@@ -142,6 +143,7 @@ static ncclResult_t canConnect(int* ret, struct ncclTopoSystem* topo, struct ncc
 struct setupReq {
   int netDev;
   int useGdr;
+  int needFlush;
 };
 
 
@@ -154,6 +156,8 @@ static ncclResult_t sendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph
   NCCLCHECK(ncclTopoGetNetDev(comm, myInfo->rank, graph, channelId, -1, &req.netDev, &proxyRank));
   NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, req.netDev, 1, &req.useGdr));
   send->conn.direct |= req.useGdr ? NCCL_DIRECT_NIC : 0;
+  // Determine whether we need to flush the GDR buffer on recv or not
+  if (req.useGdr) NCCLCHECK(ncclTopoNeedFlush(comm->topo, myInfo->busId, &req.needFlush));
 
   NCCLCHECK(ncclTopoGetLocalRank(comm->topo, myInfo->rank, &send->proxyConn.localRank));
   NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_COLLNET, 1, myInfo->rank, &send->proxyConn));
@@ -377,7 +381,7 @@ static ncclResult_t sharedBuffersGet(struct ncclComm* comm, int type, int slot,
 static ncclResult_t sharedBuffersDestroy(struct ncclComm* comm) {
   struct ncclProxySharedCollNet* state = &comm->proxyState.progressState.collNet;
   if (state->size == 0) return ncclSuccess;
-  CUDACHECK(hipFree(state->cudaBuff));
+  CUDACHECK(cudaFree(state->cudaBuff));
   NCCLCHECK(ncclCudaHostFree(state->hostBuff));
   // This will be called multiple times, with multiple channels and send/recv. Make sure we only do it once.
   state->size = 0;
@@ -395,6 +399,7 @@ static ncclResult_t recvProxySetup(struct ncclProxyConnection* connection, struc
 
   resources->netDev = req->netDev;
   resources->useGdr = req->useGdr;
+  resources->needFlush = req->needFlush;
   ncclNetProperties_t props;
   NCCLCHECK(collNetGetProperties(comm, req->netDev, &props));
   /* DMA-BUF support */
@@ -567,7 +572,7 @@ static ncclResult_t sendProxyFree(struct ncclProxyConnection* connection, struct
     }
     struct connectMapMem* mems = resources->map.mems;
     NCCLCHECK(ncclCudaHostFree(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr));
-    CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+    CUDACHECK(cudaFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
     if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
     NCCLCHECK(sharedBuffersDestroy(comm));
     NCCLCHECK(sharedFree(comm, resources->netDev));
@@ -587,7 +592,7 @@ static ncclResult_t recvProxyFree(struct ncclProxyConnection* connection, struct
     }
     struct connectMapMem* mems = resources->map.mems;
     NCCLCHECK(ncclCudaHostFree(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr));
-    CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+    CUDACHECK(cudaFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
     if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
     NCCLCHECK(sharedBuffersDestroy(comm));
     NCCLCHECK(sharedFree(comm, resources->netDev));
@@ -763,7 +768,7 @@ static ncclResult_t recvProxyProgress(struct ncclComm* comm, struct ncclProxyArg
           TRACE(NCCL_NET, "recvProxy [%lu/%d/%d] received, size %d", sub->received, group, buffSlot, totalSize);
           sub->received += args->sliceSteps;
           sub->requests[buffSlot] = NULL;
-          if (1 && reqFifo[group][buffSlot].size > 0 && resources->useGdr) {
+          if (reqFifo[group][buffSlot].size > 0 && resources->useGdr && resources->needFlush) {
             // GDRCOPY support
             if (resources->gdcFlush) {
 #if defined (__x86_64__)

projects/rccl/src/transport/net.cc

@@ -67,7 +67,7 @@ struct connectMapMem{
   int size;
   union {
     char shmPath[PATH_MAX];
-    hipIpcMemHandle_t ipc;
+    cudaIpcMemHandle_t ipc;
   };
 };
 
@@ -125,6 +125,7 @@ struct recvResources {
   int netDev;
   int useGdr;
   int useDmaBuf;
+  int needFlush;
   int maxRecvs;
   uint64_t* gdcSync;
   uint64_t* gdcFlush;
@@ -163,6 +164,7 @@ struct setupReq {
   int shared;
   int netDev;
   int useGdr;
+  int needFlush;
   int channelId;
   int connIndex;
   uint32_t* curr_hdp_reg;
@@ -226,6 +228,9 @@ static ncclResult_t recvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph
   if (req.netDev < 0) NCCLCHECK(ncclTopoGetNetDev(comm, myInfo->rank, graph, channelId, myInfo->rank, &req.netDev, &proxyRank));
   NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, req.netDev, 0, &req.useGdr));
 
+  // Determine whether we need to flush the GDR buffer on recv or not
+  if (req.useGdr) NCCLCHECK(ncclTopoNeedFlush(comm->topo, myInfo->busId, &req.needFlush));
+
   // We don't support PXN on receive yet
   NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_NET, 0, myInfo->rank, &recv->proxyConn));
 
@@ -288,24 +293,24 @@ static ncclResult_t sendConnect(struct ncclComm* comm, struct ncclConnect* conne
   if (map->sameProcess) {
     if (map->cudaDev != comm->cudaDev) {
       // Enable P2P access
-      hipError_t err = hipDeviceEnablePeerAccess(map->cudaDev, 0);
-      if (err == hipErrorPeerAccessAlreadyEnabled) {
-        hipGetLastError();
-      } else if (err != hipSuccess) {
-        WARN("failed to peer with device %d: %d %s", map->cudaDev, err, hipGetErrorString(err));
+      cudaError_t err = cudaDeviceEnablePeerAccess(map->cudaDev, 0);
+      if (err == cudaErrorPeerAccessAlreadyEnabled) {
+        cudaGetLastError();
+      } else if (err != cudaSuccess) {
+        WARN("failed to peer with device %d: %d %s", map->cudaDev, err, cudaGetErrorString(err));
         return ncclInternalError;
       }
     }
   } else {
     NCCLCHECK(netMapShm(map->mems+NCCL_NET_MAP_HOSTMEM));
     if (map->mems[NCCL_NET_MAP_DEVMEM].size) {
-      CUDACHECK(hipIpcOpenMemHandle((void**)&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].ipc, hipIpcMemLazyEnablePeerAccess));
+      CUDACHECK(cudaIpcOpenMemHandle((void**)&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].ipc, cudaIpcMemLazyEnablePeerAccess));
       map->mems[NCCL_NET_MAP_DEVMEM].cpuPtr = NULL;
     }
     if (map->mems[NCCL_NET_MAP_SHARED_DEVMEM].size) {
       void** sharedDevMemPtr = comm->proxyState.sharedDevMems+send->proxyConn.localRank;
       if (*sharedDevMemPtr == NULL) {
-        CUDACHECK(hipIpcOpenMemHandle(sharedDevMemPtr, map->mems[NCCL_NET_MAP_SHARED_DEVMEM].ipc, hipIpcMemLazyEnablePeerAccess));
+        CUDACHECK(cudaIpcOpenMemHandle(sharedDevMemPtr, map->mems[NCCL_NET_MAP_SHARED_DEVMEM].ipc, cudaIpcMemLazyEnablePeerAccess));
       }
       map->mems[NCCL_NET_MAP_SHARED_DEVMEM].gpuPtr = (char*)(*sharedDevMemPtr);
       map->mems[NCCL_NET_MAP_SHARED_DEVMEM].cpuPtr = NULL;
@@ -357,7 +362,7 @@ static ncclResult_t sendFree(struct ncclConnector* send) {
     if (map->sameProcess == 0) {
       NCCLCHECK(ncclShmClose(map->mems[NCCL_NET_MAP_HOSTMEM].cpuPtr, map->mems[NCCL_NET_MAP_HOSTMEM].gpuPtr, map->mems[NCCL_NET_MAP_HOSTMEM].size));
       if (map->mems[NCCL_NET_MAP_DEVMEM].size) {
-        CUDACHECK(hipIpcCloseMemHandle(map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr));
+        CUDACHECK(cudaIpcCloseMemHandle(map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr));
       }
     }
   }
@@ -371,7 +376,7 @@ static ncclResult_t recvFree(struct ncclConnector* recv) {
 
 #define NCCL_SHARED_STEPS 16
 static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, int localRank, int type, int sameProcess,
-    int nChannels, char** gpuPtr, char** cpuPtr, int* size, hipIpcMemHandle_t* ipc) {
+    int nChannels, char** gpuPtr, char** cpuPtr, int* size, cudaIpcMemHandle_t* ipc) {
   if (cuda == 0 && sameProcess == 0) {
       WARN("PXN should not use host buffers for data");
       return ncclInternalError;
@@ -396,7 +401,7 @@ static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, int local
   if (cuda && state->cudaBuff == NULL) {
     NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, state->size, comm->sideStream, cuda));
     if (sameProcess == 0) {
-      CUDACHECK(hipIpcGetMemHandle(&state->ipc, state->cudaBuff));
+      CUDACHECK(cudaIpcGetMemHandle(&state->ipc, state->cudaBuff));
     }
   }
   if (!cuda && state->hostBuff == NULL) {
@@ -407,7 +412,7 @@ static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, int local
     if (gpuPtr) *gpuPtr = *cpuPtr;
   } else {
     if (gpuPtr) *gpuPtr = NULL;
-    if (ipc) memcpy(ipc, &state->ipc, sizeof(hipIpcMemHandle_t));
+    if (ipc) memcpy(ipc, &state->ipc, sizeof(cudaIpcMemHandle_t));
   }
   return ncclSuccess;
 }
@@ -427,7 +432,7 @@ static ncclResult_t sharedBuffersDestroy(struct ncclComm* comm, int localRank, i
   if (state->size == 0) NCCLCHECK(ncclInternalError);
   state->refcount--;
   if (state->refcount == 0) {
-    if (state->cudaBuff) CUDACHECK(hipFree(state->cudaBuff));
+    if (state->cudaBuff) CUDACHECK(cudaFree(state->cudaBuff));
     if (state->hostBuff) NCCLCHECK(ncclCudaHostFree(state->hostBuff));
   }
   if (peer->send.refcount || peer->recv.refcount) return ncclSuccess;
@@ -492,6 +497,7 @@ static ncclResult_t recvProxySetup(struct ncclProxyConnection* connection, struc
   resources->netDev = req->netDev;
   resources->shared = connection->shared = req->shared;
   resources->useGdr = req->useGdr;
+  resources->needFlush = req->needFlush;
   resources->channelId = req->channelId;
   resources->connIndex = req->connIndex;
   ncclNetProperties_t props;
@@ -550,7 +556,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
   map->sameProcess =
     comm->peerInfo[resources->rank].pidHash == comm->peerInfo[comm->rank].pidHash ? 1 : 0;
   map->shared = resources->shared;
-  CUDACHECK(hipGetDevice(&map->cudaDev));
+  CUDACHECK(cudaGetDevice(&map->cudaDev));
 
   if (resources->shared == 0) { // Only allocate dedicated buffers for ring/tree, not for p2p
     for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
@@ -586,7 +592,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
       map->mems[NCCL_NET_MAP_DEVMEM].cpuPtr = map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr;
     }
     if (!map->sameProcess) {
-      CUDACHECK(hipIpcGetMemHandle(&map->mems[NCCL_NET_MAP_DEVMEM].ipc, map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr));
+      CUDACHECK(cudaIpcGetMemHandle(&map->mems[NCCL_NET_MAP_DEVMEM].ipc, map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr));
     }
   }
   if (map->sameProcess) {
@@ -801,7 +807,7 @@ static ncclResult_t sendProxyFree(struct ncclProxyConnection* connection, struct
   } else {
     NCCLCHECK(ncclShmClose(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr, NULL, mems[NCCL_NET_MAP_HOSTMEM].size));
   }
-  CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+  CUDACHECK(cudaFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
   if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
   if (resources->shared) {
     NCCLCHECK(sharedBuffersDestroy(comm, resources->localRank, 0));
@@ -832,7 +838,7 @@ static ncclResult_t recvProxyFree(struct ncclProxyConnection* connection, struct
   }
   struct connectMapMem* mems = resources->map.mems;
   NCCLCHECK(ncclCudaHostFree(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr));
-  CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+  CUDACHECK(cudaFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
   if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
   if (resources->shared) {
     NCCLCHECK(sharedBuffersDestroy(comm, resources->localRank, 1));
@@ -1150,7 +1156,7 @@ static ncclResult_t recvProxyProgress(struct ncclComm* comm, struct ncclProxyArg
         for (int i=0; i<NCCL_PROXY_MAX_SUBS; i++) sizes[i] = 0;
         NCCLCHECK(ncclNetTest(comm, subGroup->requests[step%NCCL_STEPS], &done, sizes));
         if (done) {
-          int useGdr = 0;
+          int needFlush = 0;
           int totalSize = 0;
           for (int i=0; i<NCCL_PROXY_MAX_SUBS; i++) totalSize += sizes[i];
           for (int i=0; i<subGroup->groupSize; i++) {
@@ -1175,11 +1181,11 @@ static ncclResult_t recvProxyProgress(struct ncclComm* comm, struct ncclProxyArg
             for (uint64_t step=sub->received-args->sliceSteps; step<sub->received; step++) ncclProfilingRecord(args, s+i, step, ncclProxyProfileRecvFlushWait);
             if (step < sub->nsteps) {
               struct recvResources* resources = (struct recvResources*) (sub->connection->transportResources);
-              if (resources->useGdr) useGdr = 1;
+              if (resources->useGdr) needFlush |= resources->needFlush;
             }
           }
           subGroup->requests[step%NCCL_STEPS] = NULL;
-          if (totalSize > 0 && p == NCCL_PROTO_SIMPLE && useGdr) {
+          if (totalSize > 0 && p == NCCL_PROTO_SIMPLE && needFlush) {
             // GDRCOPY support
             struct recvResources* resources = (struct recvResources*) (subGroup->connection->transportResources);
             if (resources->gdcFlush) {

projects/rccl/src/transport/net_socket.cc

@@ -343,7 +343,7 @@ ncclResult_t ncclSocketConnect(int dev, void* opaqueHandle, void** sendComm) {
   comm->nSocks = handle->nSocks;
   comm->nThreads = handle->nThreads;
   comm->dev = dev;
-  CUDACHECK(hipGetDevice(&comm->cudaDev));
+  CUDACHECK(cudaGetDevice(&comm->cudaDev));
   for (; i<comm->nSocks+1; i++) {
     sock = i == comm->nSocks ? &comm->ctrlSock : comm->socks+i;
     NCCLCHECK(ncclSocketInit(sock, &handle->connectAddr, NULL, 1));
@@ -388,7 +388,7 @@ ncclResult_t ncclSocketAccept(void* listenComm, void** recvComm) {
   rComm->nSocks = lComm->nSocks;
   rComm->nThreads = lComm->nThreads;
   rComm->dev = lComm->dev;
-  CUDACHECK(hipGetDevice(&rComm->cudaDev));
+  CUDACHECK(cudaGetDevice(&rComm->cudaDev));
   lComm->sock.asyncFlag = 1;
   for (; i<rComm->nSocks+1; i++) {
     uint8_t sendSockIdx;

projects/rccl/src/transport/p2p.cc

@@ -14,7 +14,7 @@
 
 struct ncclP2pBuff {
   void* directPtr;
-  hipIpcMemHandle_t devIpc;
+  cudaIpcMemHandle_t devIpc;
 };
 
 struct p2pConnectInfo {
@@ -47,8 +47,8 @@ struct p2pProxyInfo {
 
   // Used by progress only
   uint64_t step;
-  hipStream_t stream;
-  hipEvent_t events[NCCL_STEPS];
+  cudaStream_t stream;
+  cudaEvent_t events[NCCL_STEPS];
 };
 static_assert(sizeof(p2pConnectInfo) <= CONNECT_SIZE, "P2P Connect info is too large");
 
@@ -74,11 +74,11 @@ struct p2pRecvResources {
 /* Convert a PCI busId string into a local cudaDev device index (cf. CUDA_VISIBLE_DEVICES) */
 static int busIdToCudaDev(int64_t busId) {
   int ndev;
-  if (hipGetDeviceCount(&ndev) != hipSuccess)
+  if (cudaGetDeviceCount(&ndev) != cudaSuccess)
     return -1;
   for (int i = 0; i < ndev; i++) {
     char devBusIdStr[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
-    if (hipDeviceGetPCIBusId(devBusIdStr, NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE, i) != hipSuccess)
+    if (cudaDeviceGetPCIBusId(devBusIdStr, NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE, i) != cudaSuccess)
       return -1;
     int64_t devBusId;
     NCCLCHECK(busIdToInt64(devBusIdStr, &devBusId));
@@ -141,7 +141,7 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
 
   // Check that CUDA can do P2P
   int p2p;
-  if (hipDeviceCanAccessPeer(&p2p, cudaDev1, cudaDev2) != hipSuccess) {
+  if (cudaDeviceCanAccessPeer(&p2p, cudaDev1, cudaDev2) != cudaSuccess) {
     INFO(NCCL_INIT|NCCL_P2P,"peer query failed between dev %d(=%lx) and dev %d(=%lx)",
          cudaDev1, info1->busId, cudaDev2, info2->busId);
     *ret = 0;
@@ -163,13 +163,13 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
     }
     // Check that legacy IPC support is available (WSL WAR)
     char *dummy;
-    hipIpcMemHandle_t ipc;
-    CUDACHECK(hipMalloc(&dummy, CUDA_IPC_MIN));
-    if (hipIpcGetMemHandle(&ipc, dummy) != hipSuccess) {
+    cudaIpcMemHandle_t ipc;
+    NCCLCHECK(ncclCudaCalloc(&dummy, CUDA_IPC_MIN));
+    if (cudaIpcGetMemHandle(&ipc, dummy) != cudaSuccess) {
       INFO(NCCL_INIT|NCCL_P2P,"Legacy IPC not supported");
       *ret = 0;
     }
-    CUDACHECK(hipFree(dummy));
+    CUDACHECK(cudaFree(dummy));
     legacyIPC = *ret;
     return ncclSuccess;
   }
@@ -211,19 +211,19 @@ static ncclResult_t p2pMap(struct ncclPeerInfo* myInfo, struct ncclPeerInfo* pee
   if (myInfo->pidHash == peerInfo->pidHash) {
     if (peerInfo->cudaDev != myInfo->cudaDev) {
       // Enable P2P access
-      hipError_t err = hipDeviceEnablePeerAccess(peerInfo->cudaDev, 0);
-      if (err == hipErrorPeerAccessAlreadyEnabled) {
-        hipGetLastError();
-      } else if (err != hipSuccess) {
+      cudaError_t err = cudaDeviceEnablePeerAccess(peerInfo->cudaDev, 0);
+      if (err == cudaErrorPeerAccessAlreadyEnabled) {
+        cudaGetLastError();
+      } else if (err != cudaSuccess) {
         WARN("failed to peer with device %d(=%lx): %d %s",
-            peerInfo->cudaDev, peerInfo->busId, err, hipGetErrorString(err));
+            peerInfo->cudaDev, peerInfo->busId, err, cudaGetErrorString(err));
         return ncclInternalError;
       }
     }
     *devMem = p2pBuff->directPtr;
     *ipcPtr = NULL;
   } else {
-    CUDACHECK(hipIpcOpenMemHandle(devMem, p2pBuff->devIpc, hipIpcMemLazyEnablePeerAccess));
+    CUDACHECK(cudaIpcOpenMemHandle(devMem, p2pBuff->devIpc, cudaIpcMemLazyEnablePeerAccess));
     *ipcPtr = *devMem;
   }
   return ncclSuccess;
@@ -411,8 +411,8 @@ ncclResult_t p2pRecvConnect(struct ncclComm* comm, struct ncclConnect* connectIn
 ncclResult_t p2pSendFree(struct ncclConnector* send) {
   struct p2pSendResources* resources = (struct p2pSendResources*)send->transportResources;
   if (resources) {
-    if (resources->sendMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->sendMemIpc));
-    if (resources->recvMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->recvMemIpc));
+    if (resources->sendMemIpc) CUDACHECK(cudaIpcCloseMemHandle(resources->sendMemIpc));
+    if (resources->recvMemIpc) CUDACHECK(cudaIpcCloseMemHandle(resources->recvMemIpc));
     free(resources);
   }
   return ncclSuccess;
@@ -421,8 +421,8 @@ ncclResult_t p2pSendFree(struct ncclConnector* send) {
 ncclResult_t p2pRecvFree(struct ncclConnector* recv) {
   struct p2pRecvResources* resources = (struct p2pRecvResources*)recv->transportResources;
   if (resources) {
-    if (resources->sendMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->sendMemIpc));
-    if (resources->recvMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->recvMemIpc));
+    if (resources->sendMemIpc) CUDACHECK(cudaIpcCloseMemHandle(resources->sendMemIpc));
+    if (resources->recvMemIpc) CUDACHECK(cudaIpcCloseMemHandle(resources->recvMemIpc));
     if (useMemcpy) {
       NCCLCHECK(ncclShmClose(resources->shm, resources->devShm, resources->shmSize));
     }
@@ -457,10 +457,10 @@ static ncclResult_t p2pSendProxySetup(struct ncclProxyConnection* connection, st
     struct ncclP2pBuff* p2pBuff = (struct ncclP2pBuff*)respBuff;
     NCCLCHECK(ncclCudaCalloc((char**)&p2pBuff->directPtr, size, comm->sideStream, true));
     connection->transportResources = p2pBuff->directPtr;
-    hipError_t res = hipIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
-    if (res != hipSuccess) {
-      WARN("hipIpcGetMemHandle failed : %s", hipGetErrorString(res));
-      hipFree(p2pBuff->directPtr);
+    cudaError_t res = cudaIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
+    if (res != cudaSuccess) {
+      WARN("cudaIpcGetMemHandle failed : %s", cudaGetErrorString(res));
+      cudaFree(p2pBuff->directPtr);
       free(p2pBuff);
       CUDACHECK(res);
     }
@@ -476,10 +476,10 @@ static ncclResult_t p2pRecvProxySetup(struct ncclProxyConnection* connection, st
   struct ncclP2pBuff* p2pBuff = (struct ncclP2pBuff*)respBuff;
   NCCLCHECK(ncclCudaCalloc((char**)&p2pBuff->directPtr, size, comm->sideStream, true));
   connection->transportResources = p2pBuff->directPtr;
-  hipError_t res = hipIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
-  if (res != hipSuccess) {
-    WARN("hipIpcGetMemHandle failed : %s", hipGetErrorString(res));
-    hipFree(p2pBuff->directPtr);
+  cudaError_t res = cudaIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
+  if (res != cudaSuccess) {
+    WARN("cudaIpcGetMemHandle failed : %s", cudaGetErrorString(res));
+    cudaFree(p2pBuff->directPtr);
     free(p2pBuff);
     CUDACHECK(res);
   }
@@ -493,9 +493,9 @@ static ncclResult_t p2pSendProxyConnect(struct ncclProxyConnection* connection,
   if (reqSize != sizeof(void*)) return ncclInternalError;
   proxyInfo->recvFifo = *((char**)reqBuff);
 
-  CUDACHECK(hipStreamCreateWithFlags(&proxyInfo->stream, hipStreamNonBlocking));
+  CUDACHECK(cudaStreamCreateWithFlags(&proxyInfo->stream, cudaStreamNonBlocking));
   for (int i=0; i<NCCL_STEPS; i++) {
-    CUDACHECK(hipEventCreate(proxyInfo->events+i));
+    CUDACHECK(cudaEventCreate(proxyInfo->events+i));
   }
   connection->proxyAppendPtr = &connection->proxyAppend;
   return ncclSuccess;
@@ -507,23 +507,23 @@ static ncclResult_t p2pSendProxyFree(struct ncclProxyConnection* connection, str
     if (proxyInfo) {
       NCCLCHECK(ncclShmClose(proxyInfo->shm, proxyInfo->devShm, proxyInfo->shmSize));
       NCCLCHECK(ncclCudaHostFree(proxyInfo->ceRecvMem));
-      CUDACHECK(hipFree(proxyInfo->ceDevBuff));
-      CUDACHECK(hipStreamDestroy(proxyInfo->stream));
+      CUDACHECK(cudaFree(proxyInfo->ceDevBuff));
+      CUDACHECK(cudaStreamDestroy(proxyInfo->stream));
       for (int i=0; i<NCCL_STEPS; i++) {
-        CUDACHECK(hipEventDestroy(proxyInfo->events[i]));
+        CUDACHECK(cudaEventDestroy(proxyInfo->events[i]));
       }
       free(proxyInfo);
     }
   } else {
     // Do not check return code as CUDA may have already shut down
-    hipFree(connection->transportResources);
+    cudaFree(connection->transportResources);
   }
   return ncclSuccess;
 }
 
 static ncclResult_t p2pRecvProxyFree(struct ncclProxyConnection* connection, struct ncclComm* comm) {
   // Do not check return code as CUDA may have already shut down
-  hipFree(connection->transportResources);
+  cudaFree(connection->transportResources);
   return ncclSuccess;
 }
 
@@ -545,7 +545,7 @@ static ncclResult_t p2pSendProxyProgress(struct ncclComm* comm, struct ncclProxy
     for (int s=0; s<args->nsubs; s++) {
       struct ncclProxySubArgs* sub = args->subs+s;
       struct p2pProxyInfo* resources = (struct p2pProxyInfo*) (sub->connection->transportResources);
-      if (p != NCCL_PROTO_SIMPLE) { // Only Simple uses hipMemcpy
+      if (p != NCCL_PROTO_SIMPLE) { // Only Simple uses cudaMemcpy
           resources->step = sub->base + sub->nsteps;
           args->done++;
           continue;
@@ -557,16 +557,16 @@ static ncclResult_t p2pSendProxyProgress(struct ncclComm* comm, struct ncclProxy
         // Check GPU has sent everything
         if ((*recvTail > sub->base+sub->transmitted)) {
           int size = sizesFifo[buffSlot];
-          CUDACHECK(hipMemcpyAsync(resources->recvFifo+buffSlot*stepSize, resources->ceDevBuff+buffSlot*stepSize, size, hipMemcpyDeviceToDevice, resources->stream));
-          CUDACHECK(hipEventRecord(resources->events[buffSlot], resources->stream));
+          CUDACHECK(cudaMemcpyAsync(resources->recvFifo+buffSlot*stepSize, resources->ceDevBuff+buffSlot*stepSize, size, cudaMemcpyDeviceToDevice, resources->stream));
+          CUDACHECK(cudaEventRecord(resources->events[buffSlot], resources->stream));
           sub->transmitted += args->sliceSteps;
         }
       }
       if (sub->done < sub->transmitted) {
         int buffSlot = (sub->base+sub->done)%NCCL_STEPS;
-        hipError_t res = hipEventQuery(resources->events[buffSlot]);
-        if (res != hipErrorNotReady) CUDACHECK(res);
-        if (res == hipSuccess) {
+        cudaError_t res = cudaEventQuery(resources->events[buffSlot]);
+        if (res != cudaErrorNotReady) CUDACHECK(res);
+        if (res == cudaSuccess) {
           sub->done += args->sliceSteps;
           // Notify SHM
           resources->shm->recvMem.tail = sub->base + sub->done;

projects/rccl/src/transport/shm.cc

@@ -123,8 +123,8 @@ struct shmProxyInfo {
 
   // used by progress only
   uint64_t step;
-  hipStream_t stream;
-  hipEvent_t events[NCCL_STEPS];
+  cudaStream_t stream;
+  cudaEvent_t events[NCCL_STEPS];
 };
 
 /* Connect to this peer */
@@ -220,9 +220,9 @@ static ncclResult_t shmSendProxyConnect(struct ncclProxyConnection* connection,
   memcpy(proxyInfo, reqBuff, reqSize);
   NCCLCHECK(ncclCudaCalloc(&proxyInfo->devFifo, comm->buffSizes[NCCL_PROTO_SIMPLE], comm->sideStream));
   NCCLCHECK(ncclCudaHostCalloc(&proxyInfo->ceRecvMem, 1));
-  CUDACHECK(hipStreamCreateWithFlags(&proxyInfo->stream, hipStreamNonBlocking));
+  CUDACHECK(cudaStreamCreateWithFlags(&proxyInfo->stream, cudaStreamNonBlocking));
   for (int i=0; i<NCCL_STEPS; i++) {
-    CUDACHECK(hipEventCreate(proxyInfo->events+i));
+    CUDACHECK(cudaEventCreate(proxyInfo->events+i));
   }
   connection->proxyAppendPtr = &connection->proxyAppend;
   connection->transportResources = proxyInfo;
@@ -238,9 +238,9 @@ static ncclResult_t shmRecvProxyConnect(struct ncclProxyConnection* connection,
   memcpy(proxyInfo, reqBuff, reqSize);
   NCCLCHECK(ncclCudaCalloc(&proxyInfo->devFifo, comm->buffSizes[NCCL_PROTO_SIMPLE], comm->sideStream));
   NCCLCHECK(ncclCudaHostCalloc(&proxyInfo->ceRecvMem, 1));
-  CUDACHECK(hipStreamCreateWithFlags(&proxyInfo->stream, hipStreamNonBlocking));
+  CUDACHECK(cudaStreamCreateWithFlags(&proxyInfo->stream, cudaStreamNonBlocking));
   for (int i=0; i<NCCL_STEPS; i++) {
-    CUDACHECK(hipEventCreate(proxyInfo->events+i));
+    CUDACHECK(cudaEventCreate(proxyInfo->events+i));
   }
   connection->proxyAppendPtr = &connection->proxyAppend;
   connection->transportResources = proxyInfo;
@@ -253,11 +253,11 @@ static ncclResult_t shmSendProxyFree(struct ncclProxyConnection* connection, str
   struct shmProxyInfo* resources = (struct shmProxyInfo*)connection->transportResources;
 
   if (resources) {
-    CUDACHECK(hipStreamDestroy(resources->stream));
-    CUDACHECK(hipFree(resources->devFifo));
+    CUDACHECK(cudaStreamDestroy(resources->stream));
+    CUDACHECK(cudaFree(resources->devFifo));
     NCCLCHECK(ncclCudaHostFree(resources->ceRecvMem));
     for (int i=0; i<NCCL_STEPS; i++) {
-      CUDACHECK(hipEventDestroy(resources->events[i]));
+      CUDACHECK(cudaEventDestroy(resources->events[i]));
     }
     free(connection->transportResources);
   }
@@ -268,11 +268,11 @@ static ncclResult_t shmRecvProxyFree(struct ncclProxyConnection* connection, str
   struct shmProxyInfo* resources = (struct shmProxyInfo*)connection->transportResources;
 
   if (resources) {
-    CUDACHECK(hipStreamDestroy(resources->stream));
-    CUDACHECK(hipFree(resources->devFifo));
+    CUDACHECK(cudaStreamDestroy(resources->stream));
+    CUDACHECK(cudaFree(resources->devFifo));
     NCCLCHECK(ncclCudaHostFree(resources->ceRecvMem));
     for (int i=0; i<NCCL_STEPS; i++) {
-      CUDACHECK(hipEventDestroy(resources->events[i]));
+      CUDACHECK(cudaEventDestroy(resources->events[i]));
     }
     free(connection->transportResources);
   }
@@ -309,8 +309,8 @@ static ncclResult_t shmSendProxyProgress(struct ncclComm* comm, struct ncclProxy
         // Check GPU has sent everything
         if ((*recvTail > sub->base+sub->transmitted)) {
           int size = sizesFifo[buffSlot];
-          CUDACHECK(hipMemcpyAsync(resources->shmFifo+buffSlot*stepSize, resources->devFifo+buffSlot*stepSize, size, hipMemcpyDeviceToHost, resources->stream));
-          CUDACHECK(hipEventRecord(resources->events[buffSlot], resources->stream));
+          CUDACHECK(cudaMemcpyAsync(resources->shmFifo+buffSlot*stepSize, resources->devFifo+buffSlot*stepSize, size, cudaMemcpyDeviceToHost, resources->stream));
+          CUDACHECK(cudaEventRecord(resources->events[buffSlot], resources->stream));
           resources->recvMem->sizesFifo[buffSlot] = size;
           __sync_synchronize(); // make sure sizesFifo is visible
           sub->transmitted += args->sliceSteps;
@@ -318,9 +318,9 @@ static ncclResult_t shmSendProxyProgress(struct ncclComm* comm, struct ncclProxy
       }
       if (sub->done < sub->transmitted) {
         int buffSlot = (sub->base+sub->done)%NCCL_STEPS;
-        hipError_t res = hipEventQuery(resources->events[buffSlot]);
-        if (res != hipErrorNotReady) CUDACHECK(res);
-        if (res == hipSuccess) {
+        cudaError_t res = cudaEventQuery(resources->events[buffSlot]);
+        if (res != cudaErrorNotReady) CUDACHECK(res);
+        if (res == cudaSuccess) {
           sub->done += args->sliceSteps;
           // Notify SHM
           resources->recvMem->tail = sub->base + sub->done;
@@ -368,16 +368,16 @@ static ncclResult_t shmRecvProxyProgress(struct ncclComm* comm, struct ncclProxy
         // Check data is ready in SHM
         if ((*recvTail > sub->base+sub->transmitted)) {
           int size = sizesFifo[buffSlot];
-          CUDACHECK(hipMemcpyAsync(resources->devFifo+buffSlot*stepSize, resources->shmFifo+buffSlot*stepSize, size, hipMemcpyHostToDevice, resources->stream));
-          CUDACHECK(hipEventRecord(resources->events[buffSlot], resources->stream));
+          CUDACHECK(cudaMemcpyAsync(resources->devFifo+buffSlot*stepSize, resources->shmFifo+buffSlot*stepSize, size, cudaMemcpyHostToDevice, resources->stream));
+          CUDACHECK(cudaEventRecord(resources->events[buffSlot], resources->stream));
           sub->transmitted += args->sliceSteps;
         }
       }
       if (sub->done < sub->transmitted) {
         int buffSlot = (sub->base+sub->done)%NCCL_STEPS;
-        hipError_t res = hipEventQuery(resources->events[buffSlot]);
-        if (res != hipErrorNotReady) CUDACHECK(res);
-        if (res == hipSuccess) {
+        cudaError_t res = cudaEventQuery(resources->events[buffSlot]);
+        if (res != cudaErrorNotReady) CUDACHECK(res);
+        if (res == cudaSuccess) {
           sub->done += args->sliceSteps;
           // Notify GPU
           resources->ceRecvMem->tail = sub->base + sub->done;

projects/rccl/tools/topo_expl/utils.cpp

@@ -32,7 +32,7 @@
 #include "rocm_smi/rocm_smi.h"
 
 const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+2] = { "Broadcast", "Reduce", "AllGather", "ReduceScatter", "AllReduce", "SendRecv", "AllToAllPivot" };
-const char* ncclAlgoStr[NCCL_NUM_ALGORITHMS] = { "Tree", "Ring", "CollNet" };
+const char* ncclAlgoStr[NCCL_NUM_ALGORITHMS] = { "Tree", "Ring", "CollNetDirect", "CollNetChain" };
 const char* ncclProtoStr[NCCL_NUM_PROTOCOLS] = { "LL", "LL128", "Simple" };
 
 extern NodeModel *node_model;
@@ -698,7 +698,7 @@ ncclResult_t initTransportsRank_1(struct ncclComm* comm, struct allGather3Data_t
 
   comm->nChannels = (comm->topo->nodes[GPU].count != comm->topo->nRanks && comm->topo->nodes[NET].count)
     ? std::min(treeGraph.nChannels, ringGraph.nChannels) : ringGraph.nChannels;
-  NCCLCHECK(ncclTopoPreset(comm, &treeGraph, &ringGraph, &allGather3Data[rank].topoRanks));
+  NCCLCHECK(ncclTopoPreset(comm, &treeGraph, &ringGraph, &collNetGraph, &allGather3Data[rank].topoRanks));
   return ncclSuccess;
 }