enable hip graph on multi-node (#884)

* initial checkin

* enable msccl when hip graph is on

* remove the commented out code of msccl enable check

* clean up the code

* remove the msccl HighestTransportType check logic

[ROCm/rccl commit: e1dc4d5e42]

projects/rccl/src/include/msccl/msccl_setup.h

@@ -11,6 +11,8 @@
 #include "comm.h"
 #include "msccl/msccl_struct.h"
 
+ncclResult_t mscclGetCaptureStatus(hipStream_t stream);
+
 ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream);
 
 ncclResult_t mscclSetupSyncFlags(hipStream_t stream);
@@ -19,7 +21,7 @@ ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm);
 
 ncclResult_t mscclSetupCount(struct mscclAlgo* hostAlgo, ncclComm_t comm, size_t count, ncclDataType_t dataType);
 
-ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm);
+ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm, hipStream_t stream);
 
 ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count,
     ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,

projects/rccl/src/include/msccl/msccl_status.h

@@ -10,6 +10,8 @@
 
 mscclStatus& mscclGetStatus();
 
+mscclSavedProxyArgs& mscclGetSavedProxyArgs();
+
 mscclThreadLocalStatus& mscclGetThreadLocalStatus();
 
 #endif

projects/rccl/src/include/msccl/msccl_struct.h

@@ -160,11 +160,29 @@ struct mscclSavedSchedulerParam {
   hipStream_t stream;
 };
 
+enum mscclCaptureStatus {
+  mscclNoCapture,
+  mscclNewCapture,
+  mscclExistingCapture
+};
+
+struct mscclProxyArg {
+  struct mscclAlgo* hostAlgo;
+  ncclComm_t comm;
+  mscclProxyArg(struct mscclAlgo* hostAlgo, ncclComm_t comm) 
+    : hostAlgo(hostAlgo), comm(comm) {}
+};
+
+typedef std::map<unsigned long long, std::vector<struct mscclProxyArg>> mscclSavedProxyArgs;
+
 struct mscclThreadLocalStatus {
   bool mscclIsCallerFlag;
   mscclGroupStatus groupStatus;
   int groupDepth;
   std::vector<struct mscclSavedSchedulerParam> savedSchedulerParams;
+  unsigned long long captureId;
+  mscclCaptureStatus captureStatus;
+  hipGraph_t graph;
 };
 
 struct mscclStatus {
@@ -189,6 +207,9 @@ struct mscclStatus {
   mscclSchedulerInterface* mscclSchedulerPtr;
   std::vector<mscclAlgoMeta> algoMetas;
   std::vector<std::map<int, mscclAlgoHandle_t>> rankToAlgoHandles;
+  bool graphEnabled;
+  bool graphFirstKernel;
+  bool needsProxy;
 };
 
 struct alignas(16) mscclWork {

projects/rccl/src/include/transport.h

@@ -101,7 +101,7 @@ struct ncclTransport {
 };
 
 ncclResult_t ncclTransportP2pConnect(struct ncclComm* comm, int channelId, int nrecv, int* peerRecv, int nsend, int* peerSend, int connIndex);
-ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL);
+ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType=NULL, bool* needsProxy=NULL);
 
 ncclResult_t ncclNvlsInit(struct ncclComm* comm);
 ncclResult_t ncclNvlsSetup(struct ncclComm* comm, struct ncclComm* parent);