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

Skip unused peer connection in scatter and gather (#301)

Browse Source 
[ROCm/rccl commit: 4e68229c8b]
This commit is contained in:
Wenkai Du
2020-11-12 15:47:34 -08:00
committed by GitHub
parent a7ef699687
commit 62d21047b8
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/collectives/device/gather.h
+3 -2
View File
@@ -52,9 +52,10 @@ __device__ void ncclGatherKernel(struct CollectiveArgs* args) {
else {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
int peerNone = -1;
if (rank == root) {
ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
prims(tid, nthreads, &peerRecv, &peerNone, NULL, stepSize, channel, comm);
ssize_t recv_offset = chunkOffset + peerRecv*size;
prims.recv(thisOutput+recv_offset, nelem);
@@ -62,7 +63,7 @@ __device__ void ncclGatherKernel(struct CollectiveArgs* args) {
else {
if (peerSend == root) {
ncclPrimitives<UNROLL, GATHER_CHUNKSTEPS/GATHER_SLICESTEPS, GATHER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
prims(tid, nthreads, &peerNone, &peerSend, NULL, stepSize, channel, comm);
ssize_t send_offset = chunkOffset;
prims.send(thisInput+send_offset, nelem);
projects/rccl/src/collectives/device/scatter.h
+3 -2
View File
@@ -52,9 +52,10 @@ __device__ void ncclScatterKernel(struct CollectiveArgs* args) {
else {
int peerSend = (rank+(blockIdx.x*peersPerChan)+i)%nranks;
int peerRecv = (2*nranks+rank-((blockIdx.x*peersPerChan)%nranks)-(i%nranks))%nranks;
int peerNone = -1;
if (rank == root) {
ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
prims(tid, nthreads, &peerNone, &peerSend, NULL, stepSize, channel, comm);
ssize_t send_offset = chunkOffset + peerSend*size;
prims.send(thisInput+send_offset, nelem);
@@ -62,7 +63,7 @@ __device__ void ncclScatterKernel(struct CollectiveArgs* args) {
else {
if (peerRecv == root) {
ncclPrimitives<UNROLL, SCATTER_CHUNKSTEPS/SCATTER_SLICESTEPS, SCATTER_SLICESTEPS, T, 1, 1, 0, FUNC>
prims(tid, nthreads, &peerRecv, &peerSend, NULL, stepSize, channel, comm);
prims(tid, nthreads, &peerRecv, &peerNone, NULL, stepSize, channel, comm);
ssize_t recv_offset = chunkOffset;
prims.recv(thisOutput+recv_offset, nelem);