From 123de1621f78f8cc07bff3dae4d04b3a980d320a Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Mon, 10 Jan 2022 08:26:01 -0800
Subject: [PATCH] Use hipGraph instead of cudaGraph (#487)

[ROCm/rccl commit: 3669e124323325ed11a397e0615eb00ad482957c]
---
 projects/rccl/src/enqueue.cc        | 4 ++--
 projects/rccl/src/group.cc          | 2 +-
 projects/rccl/src/include/comm.h    | 4 +---
 projects/rccl/src/include/enqueue.h | 4 ++--
 4 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/projects/rccl/src/enqueue.cc b/projects/rccl/src/enqueue.cc
index 15915fe19e..95035a0483 100644
--- a/projects/rccl/src/enqueue.cc
+++ b/projects/rccl/src/enqueue.cc
@@ -1101,7 +1101,7 @@ void* graphHelperFunc(void *args) {
   }
 }
 
-ncclResult_t ncclGetCudaGraph(ncclComm_t comm, cudaGraph_t* graph) {
+ncclResult_t ncclGetCudaGraph(ncclComm_t comm, hipGraph_t* graph) {
   comm->usingCudaGraph = 0;
 #if CUDART_VERSION >= 11030
   hipStreamCaptureStatus captureStatus;
@@ -1285,7 +1285,7 @@ ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
         info->datatype, info->op, info->root, info->comm, info->comm->nRanks, info->stream);
 
     // Check whether we are in cuda graph mode
-    cudaGraph_t graph;
+    hipGraph_t graph;
     ncclComm_t comm = info->comm;
     NCCLCHECKGOTO(ncclGetCudaGraph(comm, &graph), ret, end);
 
diff --git a/projects/rccl/src/group.cc b/projects/rccl/src/group.cc
index 8a8d212a61..dee0ccc288 100644
--- a/projects/rccl/src/group.cc
+++ b/projects/rccl/src/group.cc
@@ -174,7 +174,7 @@ ncclResult_t ncclGroupEnd() {
   for (int i=0; i<ncclGroupIndex; i++) doneArray[i] = 1;
   ncclResult_t ret = ncclGroupError;
   int usingCudaGraphAll = -1;
-  cudaGraph_t* graphs = NULL;
+  hipGraph_t* graphs = NULL;
   if (ret != ncclSuccess) goto group_cleanup;
 
   /* Launch async ncclCommInitRank */
diff --git a/projects/rccl/src/include/comm.h b/projects/rccl/src/include/comm.h
index 5ccd0eba97..8c7ab14ff1 100644
--- a/projects/rccl/src/include/comm.h
+++ b/projects/rccl/src/include/comm.h
@@ -15,8 +15,6 @@
 // [/RCCL]
 
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  typedef void *cudaGraph_t;
-  typedef void *cudaGraphNode_t;
   #define HIPRT_CB
 #else
 #include "collectives.h"
@@ -216,7 +214,7 @@ struct ncclComm {
   struct ncclQueueInfo* enqueueInfo;
   int nQueueInfoCreated;
   int nQueueInfoDestroyed;
-  cudaGraphNode_t lastSetupNode;
+  hipGraphNode_t lastSetupNode;
   unsigned long long lastCudaGraphId;
   int driverVersion;
   pfn_cuMemGetAddressRange_t pfnCuMemGetAddressRange;
diff --git a/projects/rccl/src/include/enqueue.h b/projects/rccl/src/include/enqueue.h
index b059879929..c6674a83ef 100644
--- a/projects/rccl/src/include/enqueue.h
+++ b/projects/rccl/src/include/enqueue.h
@@ -28,8 +28,8 @@ ncclResult_t ncclSetupP2pKernel(struct ncclInfo* info);
 ncclResult_t ncclSetupAsyncKernels(struct ncclComm* comm);
 template<int USING_CUDA_GRAPH>
 void HIPRT_CB ncclEnqueueHostSetup(void* arg);
-ncclResult_t ncclGetCudaGraph(ncclComm_t comm, cudaGraph_t* graph);
-ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, cudaGraph_t graph);
+ncclResult_t ncclGetCudaGraph(ncclComm_t comm, hipGraph_t* graph);
+ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, hipGraph_t graph);
 
 struct ncclBuffRegInfo {
   void* sendbuffsBase[NCCL_MAX_INTRA_RANKS];