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

Add gather, scatter and alltoall collectives

Browse Source 
Introducing 3 new APIs:
ncclResult_t  ncclGather(const void* sendbuff, void* recvbuff, size_t sendcount,
    ncclDataType_t datatype, int root, ncclComm_t comm, hipStream_t stream);
ncclResult_t  ncclScatter(const void* sendbuff, void* recvbuff,
    size_t recvcount, ncclDataType_t datatype, int root, ncclComm_t comm,
    hipStream_t stream);
ncclResult_t  ncclAllToAll(const void* sendbuff, void* recvbuff, size_t count,
    ncclDataType_t datatype, ncclComm_t comm, hipStream_t stream);

Only out of place operation is supported.
Preprocessor symbol RCCL_GATHER_SCATTER=1 indicates API availibility.
By default the APIs launche RCCL kernel implementation, which can be disabled by
RCCL_ALLTOALL_KERNEL_DISABLE=1. Then the APIs use wrapper around ncclSend and ncclRecv.
This commit is contained in:
Wenkai Du committed 2020-06-09 17:44:08 -07:00
1 parent 71ec3e09df
commit e80e29573c
24 files changed
+617 -15
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/misc/argcheck.cc
+5 -2
View File
@@ -46,11 +46,14 @@ ncclResult_t ArgsCheck(struct ncclInfo* info) {
}
// Type is OK, compute nbytes. Convert Allgather/Broadcast/P2P calls to chars.
info->nBytes = info->count * ncclTypeSize(info->datatype);
if (info->coll == ncclCollAllGather || info->coll == ncclCollBroadcast) {
if (info->coll == ncclCollAllGather || info->coll == ncclCollBroadcast
|| info->coll == ncclCollGather || info->coll == ncclCollScatter || info->coll == ncclCollAllToAll) {
info->count = info->nBytes;
info->datatype = ncclInt8;
}
if (info->coll == ncclCollAllGather || info->coll == ncclCollReduceScatter) info->nBytes *= info->comm->nRanks; // count is per rank
if (info->coll == ncclCollAllGather || info->coll == ncclCollReduceScatter
|| info->coll == ncclCollGather || info->coll == ncclCollScatter || info->coll == ncclCollAllToAll)
info->nBytes *= info->comm->nRanks; // count is per rank
if (info->op < 0 || info->op >= ncclNumOps) {
WARN("%s : invalid reduction operation %d", info->opName, info->op);