AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

2.15.1-1

Browse Source 
Add support for H100 (sm90).
Make sure NCCL kernel honor user stream priorities.
This commit is contained in:
Sylvain Jeaugey
2022-09-27 02:31:13 -07:00
parent 78313a6d21
commit da8152e57a
22 changed files with 303 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transport.cc
+8 -8
View File
@@ -42,7 +42,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;
@@ -77,15 +77,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];
uint32_t sendMask = comm->connectSend[sendPeer];
uint64_t recvMask = comm->connectRecv[recvPeer];
uint64_t sendMask = comm->connectSend[sendPeer];
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;
}
@@ -94,7 +94,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;
}
@@ -119,7 +119,7 @@ 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;
@@ -129,7 +129,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
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;
@@ -137,7 +137,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
}
}
TIME_STOP(4);
comm->connectRecv[recvPeer] = comm->connectSend[sendPeer] = 0;
comm->connectRecv[recvPeer] = comm->connectSend[sendPeer] = 0UL;
}
CUDACHECK(cudaStreamSynchronize(transportSetupStream));
CUDACHECK(cudaStreamDestroy(transportSetupStream));