From e1dc4d5e422e143333e9cec285169c935c79ce4a Mon Sep 17 00:00:00 2001 From: Andy li Date: Tue, 12 Sep 2023 06:30:04 +0800 Subject: [PATCH] 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 --- src/collectives/device/msccl_kernel_impl.h | 8 ++- src/collectives/msccl.cc | 7 +- src/include/msccl/msccl_setup.h | 4 +- src/include/msccl/msccl_status.h | 2 + src/include/msccl/msccl_struct.h | 21 ++++++ src/include/transport.h | 2 +- src/init.cc | 20 ++++-- src/misc/msccl/msccl_lifecycle.cc | 20 ++---- src/misc/msccl/msccl_setup.cc | 74 ++++++++++++++++++++-- src/misc/msccl/msccl_status.cc | 6 ++ src/transport.cc | 13 ++-- 11 files changed, 142 insertions(+), 35 deletions(-) diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h index fe4ff520bf..40cf056f64 100644 --- a/src/collectives/device/msccl_kernel_impl.h +++ b/src/collectives/device/msccl_kernel_impl.h @@ -23,6 +23,9 @@ extern __shared__ struct mscclShmemData mscclShmem; #define COMPUTE_FLAG(__WORKINDEX__,__GRIDOFFSET_ITER__,__STEP__) \ MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS*(uint64_t)__WORKINDEX__ + ((uint64_t)__GRIDOFFSET_ITER__ * MSCCL_MAX_NUM_STEPS + (uint64_t)__STEP__) +#define GET_WORKINDEX_FROM_FLAG(__FLAG__) \ + (__FLAG__) / (MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS) + // a copy of the volatile load/store from prims_ll template __device__ static U load(U *src) { @@ -293,7 +296,10 @@ __device__ __forceinline__ void mscclRunInterpreter( int8_t dependentBid = mscclShmem.mscclTB.dependentBid[dependentPointer+tid]; int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid]; uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep); - while ((mscclFlags + dependentBid)->flag < goalFlag); + while (true){ + uint64_t curFlag = (mscclFlags + dependentBid)->flag; + if (curFlag >= goalFlag && GET_WORKINDEX_FROM_FLAG(curFlag) == workIndex) break; + } } step += numDependencies-1; barrier(nthreads, mscclBarrierNext, mscclBarriers); diff --git a/src/collectives/msccl.cc b/src/collectives/msccl.cc index d8e7cf036c..9c93945cae 100644 --- a/src/collectives/msccl.cc +++ b/src/collectives/msccl.cc @@ -48,6 +48,8 @@ ncclResult_t mscclRunAlgo( struct mscclAlgo* hostAlgo = status.hostAlgos[mscclAlgoHandle]; struct mscclAlgo* devAlgo = status.devAlgos[mscclAlgoHandle]; + NCCLCHECK(mscclGetCaptureStatus(stream)); + NCCLCHECK(mscclSetupCount(hostAlgo, comm, count, dataType)); NCCLCHECK(mscclSetupScratch(hostAlgo, stream)); @@ -55,11 +57,14 @@ ncclResult_t mscclRunAlgo( NCCLCHECK(mscclSetupSyncFlags(stream)); if (status.connectedAlgos[comm].find(mscclAlgoHandle) == status.connectedAlgos[comm].end()) { + hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed; + CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode)); NCCLCHECK(mscclSetupConnections(hostAlgo, comm)); + CUDACHECK(hipThreadExchangeStreamCaptureMode(&mode)); status.connectedAlgos[comm].insert(mscclAlgoHandle); } - NCCLCHECK(mscclSetupProxy(hostAlgo, comm)); + NCCLCHECK(mscclSetupProxy(hostAlgo, comm, stream)); NCCLCHECK(mscclSetupKernel(sendBuff, recvBuff, count, dataType, op, hostAlgo, devAlgo, comm, stream)); diff --git a/src/include/msccl/msccl_setup.h b/src/include/msccl/msccl_setup.h index 75822052cf..893a2f5b7f 100644 --- a/src/include/msccl/msccl_setup.h +++ b/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, diff --git a/src/include/msccl/msccl_status.h b/src/include/msccl/msccl_status.h index ec38c31e1a..a8bfb09ed3 100644 --- a/src/include/msccl/msccl_status.h +++ b/src/include/msccl/msccl_status.h @@ -10,6 +10,8 @@ mscclStatus& mscclGetStatus(); +mscclSavedProxyArgs& mscclGetSavedProxyArgs(); + mscclThreadLocalStatus& mscclGetThreadLocalStatus(); #endif diff --git a/src/include/msccl/msccl_struct.h b/src/include/msccl/msccl_struct.h index 2e32049219..ec1ddca432 100644 --- a/src/include/msccl/msccl_struct.h +++ b/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> mscclSavedProxyArgs; + struct mscclThreadLocalStatus { bool mscclIsCallerFlag; mscclGroupStatus groupStatus; int groupDepth; std::vector savedSchedulerParams; + unsigned long long captureId; + mscclCaptureStatus captureStatus; + hipGraph_t graph; }; struct mscclStatus { @@ -189,6 +207,9 @@ struct mscclStatus { mscclSchedulerInterface* mscclSchedulerPtr; std::vector algoMetas; std::vector> rankToAlgoHandles; + bool graphEnabled; + bool graphFirstKernel; + bool needsProxy; }; struct alignas(16) mscclWork { diff --git a/src/include/transport.h b/src/include/transport.h index f3f47065a2..db1764e35f 100644 --- a/src/include/transport.h +++ b/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); diff --git a/src/init.cc b/src/init.cc index dc69dde031..e228120202 100644 --- a/src/init.cc +++ b/src/init.cc @@ -42,6 +42,7 @@ // [/RCCL] #include "msccl/msccl_lifecycle.h" +#include "msccl/msccl_status.h" #define STR2(v) #v #define STR(v) STR2(v) @@ -669,10 +670,6 @@ static ncclResult_t devCommSetup(ncclComm_t comm) { NCCLCHECK(ncclCudaCalloc(&tmpCommAndChans.comm.devProf, MAXCHANNELS*PROFILE_NUM_LAUNCHES), comm->sideStream); #endif - if (mscclEnabled()) { - NCCLCHECK(mscclInit(comm)); - } - NCCLCHECKGOTO(ncclCudaMemcpyAsync(devCommAndChans, &tmpCommAndChans, 1, comm->sharedRes->deviceStream.cudaStream), ret, fail); exit: CUDACHECK(cudaStreamSynchronize(comm->sharedRes->deviceStream.cudaStream)); @@ -1046,6 +1043,10 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p int *topParentLocalRanks = NULL; int tpProxyRank; + int highestTransportType = TRANSPORT_P2P; + bool needsProxy = false; + bool mscclNeedsProxy = needsProxy; + // AllGather1 - begin NCCLCHECKGOTO(ncclCalloc(&comm->peerInfo, nranks+1), ret, fail); // Extra rank to represent CollNet root NCCLCHECKGOTO(fillInfo(comm, comm->peerInfo+rank, comm->commHash), ret, fail); @@ -1414,7 +1415,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p if (comm->nRanks == 1) continue; NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &channel->ring.prev, 1, &channel->ring.next, 0), ret, fail); } - NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &ringGraph, 0), ret, fail); + NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &ringGraph, 0, &highestTransportType, &needsProxy), ret, fail); + mscclNeedsProxy |= needsProxy; if (ringGraph.nIntraChannels && rcclParamP2pNetDisable() == 0) { comm->useIntraNet = 1; // Connect NET for intranode use @@ -1434,7 +1436,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, NCCL_MAX_TREE_ARITY, channel->tree.down, 1, &channel->tree.up, 0), ret, fail); NCCLCHECKGOTO(ncclTransportP2pConnect(comm, c, 1, &channel->tree.up, NCCL_MAX_TREE_ARITY, channel->tree.down, 0), ret, fail); } - NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &treeGraph, 0), ret, fail); + NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &treeGraph, 0, &highestTransportType, &needsProxy), ret, fail); + mscclNeedsProxy |= needsProxy; INFO(NCCL_INIT, "Connected all trees"); // Setup NVLS @@ -1565,6 +1568,11 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p // Call devCommSetup before the last barrier, making sure we don't have a thread running in front and starting to // launch NCCL kernels before all cuda mem allocation is complete. That could cause a deadlock. NCCLCHECKGOTO(devCommSetup(comm), ret, fail); + if (mscclEnabled()) { + NCCLCHECK(mscclInit(comm)); + mscclStatus& status = mscclGetStatus(); + status.needsProxy |= mscclNeedsProxy; + } /* Local intra-node barrier */ NCCLCHECKGOTO(bootstrapBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks, comm->localRankToRank[0]), ret, fail); diff --git a/src/misc/msccl/msccl_lifecycle.cc b/src/misc/msccl/msccl_lifecycle.cc index 16559b73d0..540e7bf293 100644 --- a/src/misc/msccl/msccl_lifecycle.cc +++ b/src/misc/msccl/msccl_lifecycle.cc @@ -168,6 +168,8 @@ ncclResult_t mscclInit(ncclComm_t comm) { mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus(); threadLocalStatus.groupStatus = mscclNoGroup; threadLocalStatus.groupDepth = 0; + threadLocalStatus.captureId = ULLONG_MAX; + threadLocalStatus.captureStatus = mscclNoCapture; comm->mscclCompatible = mscclCommCompatible(comm); { @@ -187,6 +189,7 @@ ncclResult_t mscclInit(ncclComm_t comm) { } NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS)); status.lastStream = nullptr; + status.needsProxy = false; mscclSchedulerTriedLoadAlgo = false; NCCLCHECK(mscclSchedulerInit()); @@ -390,8 +393,6 @@ ncclResult_t mscclEnqueueCheck( size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op, mscclFunc_t func, ncclComm_t comm, hipStream_t stream) { mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); - hipStreamCaptureStatus captureStatus; - unsigned long long pid; threadLocalStatus.savedSchedulerParams.push_back({}); NCCLCHECK(mscclSetSavedSchedulerParam( @@ -402,29 +403,16 @@ ncclResult_t mscclEnqueueCheck( switch (threadLocalStatus.groupStatus) { case mscclNoGroup: if (comm->mscclCompatible) { - if (stream == (hipStream_t)0) { - captureStatus = hipStreamCaptureStatusNone; - } else { - CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid)); - } - if (captureStatus == hipStreamCaptureStatusNone) { NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back())); if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) { NCCLCHECK(mscclRunSavedParams()); break; } } - } NCCLCHECK(mscclFallBackSavedParams()); break; case mscclGroupSupportedOp: if (comm->mscclCompatible) { - if (stream == (hipStream_t)0) { - captureStatus = hipStreamCaptureStatusNone; - } else { - CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid)); - } - if (captureStatus == hipStreamCaptureStatusNone) { NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back())); if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) { // Only save counts and displs when there is suitable MSCCL algorithm for this @@ -432,8 +420,8 @@ ncclResult_t mscclEnqueueCheck( break; } } - } threadLocalStatus.groupStatus = mscclGroupUnsupportedOp; + NCCLCHECK(mscclFallBackSavedParams()); case mscclGroupUnsupportedOp: NCCLCHECK(mscclFallBackSavedParams()); break; diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc index e4254b631f..bf5aaa535c 100644 --- a/src/misc/msccl/msccl_setup.cc +++ b/src/misc/msccl/msccl_setup.cc @@ -19,6 +19,28 @@ RCCL_PARAM(MscclEnableDoneEvent, "MSCCL_ENABLE_DONE_EVENT", 1); #endif +ncclResult_t mscclGetCaptureStatus(hipStream_t stream) { + mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); + mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs(); + cudaStreamCaptureStatus captureStatus; + unsigned long long captureId; + CUDACHECK(hipStreamGetCaptureInfo_v2(stream, &captureStatus, &captureId, &threadLocalStatus.graph, nullptr, nullptr)); + if (captureStatus == cudaStreamCaptureStatusActive) { + if (savedProxyArgs.count(captureId) == 0) { + threadLocalStatus.captureStatus = mscclNewCapture; + savedProxyArgs[captureId] = std::vector(); + } else { + INFO(NCCL_INIT|NCCL_NET,"mscclGetCaptureStatus: captureId %llu is same with the previous one\n", captureId); + threadLocalStatus.captureStatus = mscclExistingCapture; + } + threadLocalStatus.captureId = captureId; + } else { + threadLocalStatus.captureStatus = mscclNoCapture; + } + INFO(NCCL_INIT|NCCL_NET,"mscclGetCaptureStatus: %d, captureId: %llu, size: %lu\n", threadLocalStatus.captureStatus, threadLocalStatus.captureId, mscclGetSavedProxyArgs()[captureId].size()); + return ncclSuccess; +} + ncclResult_t mscclSetupCount(struct mscclAlgo* hostAlgo, ncclComm_t comm, size_t count, ncclDataType_t dataType) { mscclStatus& status = mscclGetStatus(); status.stepSize = comm->buffSizes[hostAlgo->protocol] / NCCL_STEPS; @@ -55,14 +77,19 @@ ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream) { ncclResult_t mscclSetupSyncFlags(hipStream_t stream) { mscclStatus& status = mscclGetStatus(); - if (status.workIndex > (1ULL << (8*sizeof(status.workIndex))) - 2 * NCCL_MAX_OPS - 1) { + mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); + if (threadLocalStatus.captureStatus == mscclNewCapture || + status.workIndex > (1ULL << (8*sizeof(status.workIndex))) - 2 * NCCL_MAX_OPS - 1) { CUDACHECK(hipMemsetAsync(status.syncFlags, 0, sizeof(struct mscclFlag) * MSCCL_MAX_NUM_THREAD_BLOCKS, stream)); status.workIndex = 1; // setting the workIndex back to 1 for next iterations + status.graphFirstKernel = false; } return ncclSuccess; } ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm) { + mscclStatus& status = mscclGetStatus(); + // Check whether there is enough channels if (hostAlgo->nChannels > comm->nChannels) { WARN("MSCCL: number of channels available (%d) less than required (%d)", comm->nChannels, hostAlgo->nChannels); @@ -88,15 +115,19 @@ ncclResult_t mscclSetupConnections(struct mscclAlgo* hostAlgo, ncclComm_t comm) // Connect MSCCL connections mscclSetIsCallerFlag(); - NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 0)); + int highestTransportType = TRANSPORT_P2P; + bool needsProxy = false; + NCCLCHECK(ncclTransportP2pSetup(comm, NULL, 0, &highestTransportType, &needsProxy)); + status.needsProxy |= needsProxy; mscclClearIsCallerFlag(); INFO(NCCL_INIT, "MSCCL: Setup connections finished, used %ld", allocTracker[comm->cudaDev].totalAllocSize); return ncclSuccess; } -ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm) { +static ncclResult_t mscclSetupProxyImpl(struct mscclAlgo* hostAlgo, ncclComm_t comm) { mscclStatus& status = mscclGetStatus(); + mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); struct ncclProxyOp proxyOp = {}; proxyOp.connIndex = 0; proxyOp.sliceSteps = status.sliceSteps; @@ -147,6 +178,38 @@ ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm) { return ncclSuccess; } +static void HIPRT_CB mscclSetupProxyCallback(void *args) { + std::vector* params = (std::vector*)args; + INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxyCallback: proxy args size: %ld\n", params->size()); + for (auto &p : *params) { + mscclSetupProxyImpl(p.hostAlgo, p.comm); + } +} + +ncclResult_t mscclSetupProxy(struct mscclAlgo* hostAlgo, ncclComm_t comm, hipStream_t stream) { + mscclStatus& status = mscclGetStatus(); + mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); + mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs(); + if (threadLocalStatus.captureStatus == mscclNoCapture) { + INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: no capture\n"); + NCCLCHECK(mscclSetupProxyImpl(hostAlgo, comm)); + } else if (status.needsProxy) { + INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: capture\n"); + if (savedProxyArgs[threadLocalStatus.captureId].size() == 0) { + INFO(NCCL_INIT|NCCL_NET,"mscclSetupProxy: adding callback\n"); + + hipGraphNode_t callbackNode; + hipHostNodeParams p; + p.fn = mscclSetupProxyCallback; + auto params = &savedProxyArgs[threadLocalStatus.captureId]; + p.userData = params; + CUDACHECK(hipGraphAddHostNode(&callbackNode, threadLocalStatus.graph, nullptr, 0, &p)); + } + mscclGetSavedProxyArgs()[threadLocalStatus.captureId].emplace_back(hostAlgo, comm); + } + return ncclSuccess; +} + static ncclResult_t hostToDevRedOp( ncclDevRedOpFull *opFull, ncclRedOp_t op, ncclDataType_t datatype, ncclComm *comm ) { @@ -268,7 +331,7 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count dim3 grid = {(uint32_t)hostAlgo->nBlocks, 1, 1}; dim3 block = {NCCL_MAX_NTHREADS, 1, 1}; - ncclDevRedOpFull opFull; + ncclDevRedOpFull opFull = {}; NCCLCHECK(hostToDevRedOp(&opFull, op, dataType, comm)); mscclWork work; @@ -283,7 +346,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count work.maxAllowedCount = status.maxAllowedCount; work.hasReduce = hostAlgo->hasReduce; work.redOpArgIsPtr = opFull.scalarArgIsPtr; - + INFO(NCCL_INIT, "MSCCL: Setup Kernel finished"); + void *args[3] = {&comm->devComm, &devAlgo, &work}; void *func = mscclKernelEntries[(opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol]; if (enableDoneEvent) { diff --git a/src/misc/msccl/msccl_status.cc b/src/misc/msccl/msccl_status.cc index d9c2de4ec7..37a1641efd 100644 --- a/src/misc/msccl/msccl_status.cc +++ b/src/misc/msccl/msccl_status.cc @@ -4,6 +4,7 @@ ************************************************************************/ #include "msccl/msccl_status.h" +#include "msccl/msccl_struct.h" mscclStatus& mscclGetStatus() { static mscclStatus status; @@ -14,3 +15,8 @@ mscclThreadLocalStatus& mscclGetThreadLocalStatus() { static thread_local mscclThreadLocalStatus threadLocalStatus; return threadLocalStatus; } + +mscclSavedProxyArgs& mscclGetSavedProxyArgs() { + static mscclSavedProxyArgs savedProxyArgs; + return savedProxyArgs; +} diff --git a/src/transport.cc b/src/transport.cc index 896082059b..76e6119fb6 100644 --- a/src/transport.cc +++ b/src/transport.cc @@ -19,7 +19,7 @@ struct ncclTransport* ncclTransports[NTRANSPORTS] = { }; template -static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType) { +static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType, bool* needsProxy) { struct ncclPeerInfo* myInfo = comm->peerInfo+comm->rank; struct ncclPeerInfo* peerInfo = comm->peerInfo+peer; struct ncclConnector* connector = (type == 1) ? comm->channels[channelId].peers[peer]->send + connIndex : @@ -44,6 +44,7 @@ static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* connector->transportComm = transportComm; NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex)); if (transportType) *transportType = t; + if (needsProxy) *needsProxy = (transportComm->proxyProgress != NULL); return ncclSuccess; } } @@ -77,10 +78,11 @@ void dumpData(struct ncclConnect* data, int ndata) { } } -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*/) { // Stream used during transport setup; need for P2P pre-connect + CUDA Graph ncclResult_t ret = ncclSuccess; int highestType = TRANSPORT_P2P; // track highest transport type + bool needsProxyResult = false; struct ncclConnect** data = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Store intermediate send/recvData structs for connect struct ncclConnect** recvData = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Points to entries inside data for given recv connection within a channel struct ncclConnect** sendData = (ncclConnect**) malloc(sizeof(ncclConnect*) * comm->nRanks); // Points to entries inside data for given send connection within a channel @@ -103,10 +105,11 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* recvData[i] = data[i]; int sendChannels = 0, recvChannels = 0; int type; + bool proxy; TIME_START(0); for (int c=0; c(comm, graph, recvData[i]+recvChannels++, c, recvPeer, connIndex, &type), ret, fail); + NCCLCHECKGOTO(selectTransport<0>(comm, graph, recvData[i]+recvChannels++, c, recvPeer, connIndex, &type, &proxy), ret, fail); if (type > highestType) highestType = type; } } @@ -115,8 +118,9 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* sendData[i] = recvData[i]+recvChannels; for (int c=0; c(comm, graph, sendData[i]+sendChannels++, c, sendPeer, connIndex, &type), ret, fail); + NCCLCHECKGOTO(selectTransport<1>(comm, graph, sendData[i]+sendChannels++, c, sendPeer, connIndex, &type, &proxy), ret, fail); if (type > highestType) highestType = type; + needsProxyResult |= proxy; } } TIME_STOP(1); @@ -207,6 +211,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* free(recvData); if (highestTransportType != NULL) *highestTransportType = highestType; + if (needsProxy != NULL) *needsProxy = needsProxyResult; TIME_PRINT("P2P Setup/Connect"); exit: NCCLCHECK(ncclStrongStreamWaitStream(ncclCudaGraphNone(), &comm->sharedRes->deviceStream, &comm->sharedRes->hostStream));