Wenkai Du
87 строки
3.9 KiB
C++
87 строки
3.9 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2016-2020, NVIDIA CORPORATION. All rights reserved.
|
|
* Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
#include "comm.h"
|
|
#include "info.h"
|
|
#include "bootstrap.h"
|
|
|
|
extern struct ncclTransport p2pTransport;
|
|
extern struct ncclTransport shmTransport;
|
|
extern struct ncclTransport netTransport;
|
|
|
|
struct ncclTransport ncclTransports[NTRANSPORTS] = {
|
|
p2pTransport,
|
|
shmTransport,
|
|
netTransport,
|
|
};
|
|
|
|
template <int type>
|
|
static ncclResult_t selectTransport(struct ncclTopoSystem* topo, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connect, struct ncclConnector* connector, int channelId) {
|
|
for (int t=0; t<NTRANSPORTS; t++) {
|
|
struct ncclTransport *transport = ncclTransports+t;
|
|
struct ncclTransportComm* transportComm = type == 1 ? &transport->send : &transport->recv;
|
|
int ret = 0;
|
|
NCCLCHECK(transport->canConnect(&ret, topo, graph, myInfo, peerInfo));
|
|
if (ret) {
|
|
connector->transportComm = transportComm;
|
|
NCCLCHECK(transportComm->setup(topo, graph, myInfo, peerInfo, connect, connector, channelId));
|
|
return ncclSuccess;
|
|
}
|
|
}
|
|
WARN("No transport found !");
|
|
return ncclInternalError;
|
|
}
|
|
|
|
ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclChannel* channel, int nrecv, int* peerRecv, int nsend, int* peerSend) {
|
|
TRACE(NCCL_INIT, "nsend %d nrecv %d", nsend, nrecv);
|
|
uint32_t nSkippedSend = 0, nSkippedRecv = 0; /* for tracing */
|
|
struct ncclConnect connect;
|
|
struct ncclConnector* conn;
|
|
for (int i=0; i<nrecv; i++) {
|
|
int peer = peerRecv[i];
|
|
if (peer == -1 || peer >= comm->nRanks) continue;
|
|
conn = &channel->peers[peer].recv;
|
|
if (conn->connected) { ++nSkippedRecv; continue; }
|
|
memset(&connect, 0, sizeof(connect));
|
|
NCCLCHECK(selectTransport<0>(comm->topo, graph, comm->peerInfo+comm->rank, comm->peerInfo+peer, &connect, conn, channel->id));
|
|
NCCLCHECK(bootstrapSend(comm->bootstrap, peer, &connect, sizeof(struct ncclConnect)));
|
|
}
|
|
for (int i=0; i<nsend; i++) {
|
|
int peer = peerSend[i];
|
|
if (peer == -1 || peer >= comm->nRanks) continue;
|
|
conn = &channel->peers[peer].send;
|
|
if (conn->connected) { ++nSkippedSend; continue; }
|
|
memset(&connect, 0, sizeof(connect));
|
|
NCCLCHECK(selectTransport<1>(comm->topo, graph, comm->peerInfo+comm->rank, comm->peerInfo+peer, &connect, conn, channel->id));
|
|
NCCLCHECK(bootstrapSend(comm->bootstrap, peer, &connect, sizeof(struct ncclConnect)));
|
|
}
|
|
for (int i=0; i<nsend; i++) {
|
|
int peer = peerSend[i];
|
|
if (peer == -1 || peer >= comm->nRanks) continue;
|
|
conn = &channel->peers[peer].send;
|
|
if (conn->connected) {++nSkippedSend; continue; }
|
|
memset(&connect, 0, sizeof(connect));
|
|
NCCLCHECK(bootstrapRecv(comm->bootstrap, peer, &connect, sizeof(struct ncclConnect)));
|
|
NCCLCHECK(conn->transportComm->connect(&connect, 1, comm->rank, conn));
|
|
conn->connected = 1;
|
|
CUDACHECK(hipMemcpy(&channel->devPeers[peer].send, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice));
|
|
}
|
|
for (int i=0; i<nrecv; i++) {
|
|
int peer = peerRecv[i];
|
|
if (peer == -1 || peer >= comm->nRanks) continue;
|
|
conn = &channel->peers[peer].recv;
|
|
if (conn->connected) {++nSkippedRecv; continue; }
|
|
memset(&connect, 0, sizeof(connect));
|
|
NCCLCHECK(bootstrapRecv(comm->bootstrap, peer, &connect, sizeof(struct ncclConnect)));
|
|
NCCLCHECK(conn->transportComm->connect(&connect, 1, comm->rank, conn));
|
|
conn->connected = 1;
|
|
CUDACHECK(hipMemcpy(&channel->devPeers[peer].recv, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice));
|
|
}
|
|
TRACE(NCCL_INIT, "nsend %d nrecv %d nSkippedSend %u nSkippedRecv %u - DONE", nsend, nrecv, nSkippedSend, nSkippedRecv);
|
|
return ncclSuccess;
|
|
}
|