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

Merge remote-tracking branch 'nccl/master' into develop

Browse Source
This commit is contained in:
Wenkai Du
2022-04-18 11:14:51 -07:00
parent fd2f1b3b88 9bfc1c6e35
commit d28e1cb44f
81 changed files with 7358 additions and 4468 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/collectives/device/all_gather.h
+9 -9
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -13,9 +13,9 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclRing *ring = &ncclShmem->channel.ring;
const int *ringRanks = ring->devUserRanks;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLGATHER_CHUNKSTEPS : 1));
@@ -23,12 +23,12 @@ namespace {
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2);
const int nranks = ncclShmem->comm.nRanks;
const ssize_t loopSize = nChannels*int(chunkSize);
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
T *inputBuf = (T*)args->sendbuff;
T *outputBuf = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->coll.redOpArg, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->redOpArg, args->connIndex << 16);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
@@ -37,7 +37,7 @@ namespace {
realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
realChunkSize = size-gridOffset < loopSize ? args->lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
realChunkSize = int(realChunkSize);
src/collectives/device/all_reduce.cu
+2 -5
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2021 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -8,7 +8,4 @@
#include "common.h"
#include "collectives.h"
// [RCCL]
// IMPL_COLL_R(AllReduce);
IMPL_COLL_CLIQUE(AllReduce);
// [/RCCL]
IMPL_COLL_R(AllReduce);
src/collectives/device/all_reduce.h
+54 -49
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -8,26 +8,26 @@
#include "devcomm.h"
#include "collectives.h"
#include "primitives.h"
#include "clique/AllReduceCliqueKernel.h" // [RCCL] AllReduce Clique-based kernel support
//#include "clique/AllReduceCliqueKernel.h" // [RCCL] AllReduce Clique-based kernel support
namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclRing *ring = &ncclShmem->channel.ring;
int ringIx = ring->index;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLREDUCE_CHUNKSTEPS : 1));
const int nranks = ncclShmem->comm.nRanks;
const ssize_t loopSize = nChannels*nranks*chunkSize;
const ssize_t size = args->coll.count;
#ifdef ENABLE_PROFILING
auto devProf = ncclShmem->comm.devProf;
uint64_t clk, t0 = 0ULL, ws;
if (tid == 0) clk = __builtin_amdgcn_s_memrealtime();
#endif
const ssize_t size = args->count;
int minChunkSize;
if (Proto::Id == NCCL_PROTO_LL)
@@ -37,8 +37,8 @@ namespace {
minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2;
}
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto> prims
(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->coll.redOpArg, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->redOpArg, args->connIndex << 16);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
@@ -110,32 +110,35 @@ namespace {
ACCUMULATE_COUNTER(directRecv);
}
#ifdef ENABLE_PROFILING
if (tid == 0 && args->coll.opCount) devProf->elems[blockIdx.x].total_cycle += (__builtin_amdgcn_s_memrealtime() - clk);
if (tid == 0) {
struct ncclProfElem *elem = devProf.elems+args->opCount;
elem->elem[blockIdx.x].total_cycle += (__builtin_amdgcn_s_memrealtime() - clk);
}
#endif
}
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runTreeUpDown(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclTree *tree = &ncclShmem->channel.tree;
ssize_t chunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? args->coll.lastChunkSize
Proto::Id == NCCL_PROTO_SIMPLE ? args->lastChunkSize
/* LL & LL128 */ : Proto::calcBytePerStep()/sizeof(T));
const ssize_t minChunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? nthreads*8*(sizeof(uint64_t)/sizeof(T))
/* LL & LL128 */ : nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const ssize_t loopSize = int(nChannels*chunkSize);
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
if (loopSize > size)
chunkSize = divUp((int)size, int(nChannels*minChunkSize))*int(minChunkSize);
{ // Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto> prims
(tid, nthreads, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->coll.redOpArg);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto, 0> prims
(tid, nthreads, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->redOpArg);
if (tree->up == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
@@ -160,8 +163,8 @@ namespace {
}
{ // Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto> prims
(tid, nthreads, &tree->up, tree->down, args->sendbuff, args->recvbuff, args->coll.redOpArg);
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0> prims
(tid, nthreads, &tree->up, tree->down, args->sendbuff, args->recvbuff, args->redOpArg);
if (tree->up == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
@@ -189,19 +192,19 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runTreeSplit(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclTree *tree = &ncclShmem->channel.tree;
ssize_t chunkSize = int(
Proto::Id != NCCL_PROTO_LL ? args->coll.lastChunkSize
Proto::Id != NCCL_PROTO_LL ? args->lastChunkSize
: Proto::calcBytePerStep()/sizeof(T));
const ssize_t minChunkSize = int(
Proto::Id == NCCL_PROTO_SIMPLE ? nthreads*8*(sizeof(uint64_t)/sizeof(T)) :
Proto::Id == NCCL_PROTO_LL ? nthreads*(Proto::calcBytePerGrain()/sizeof(T))
/* LL128 */ : nthreads*(Proto::calcBytePerGrain()/sizeof(T))/8);
const ssize_t loopSize = int(nChannels*chunkSize);
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
int nthreadsSplit;
if (Proto::Id == NCCL_PROTO_SIMPLE) {
@@ -218,8 +221,8 @@ namespace {
if (tree->up == -1) {
// Reduce and broadcast. Max number of recv is 3, max number of send is 3
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto>
prims(tid, nthreads, tree->down, tree->down, args->sendbuff, args->recvbuff, args->coll.redOpArg);
Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid, nthreads, tree->down, tree->down, args->sendbuff, args->recvbuff, args->redOpArg);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
int nelem = min(chunkSize, size-offset);
@@ -235,8 +238,8 @@ namespace {
* into DirectRecv and DirectSend capabilities, this ctor would have both=0,
* but the ctor above for tree roots would be DirectRecv=0 DirectSend=1.
*/
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto>
prims(tid, nthreadsSplit, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->coll.redOpArg, 0*Proto::MaxGroupWidth);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto, 0>
prims(tid, nthreadsSplit, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->redOpArg, 0*Proto::MaxGroupWidth);
if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
@@ -254,8 +257,8 @@ namespace {
}
else {
// Broadcast down. Max number of recv is 1, max number of send is 3 (binary tree + local)
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &tree->up, tree->down, args->sendbuff, args->recvbuff, args->coll.redOpArg, 1*Proto::MaxGroupWidth);
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &tree->up, tree->down, args->sendbuff, args->recvbuff, args->redOpArg, 1*Proto::MaxGroupWidth);
if (tree->down[0] == -1) {
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*int(chunkSize);
@@ -294,11 +297,11 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
static constexpr int COLLNET_COPY_THREADS = 64;
const int tid = threadIdx.x;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int bid = args->bid;
const int nChannels = args->nChannels;
struct ncclDirect* tree = &ncclShmem->channel.collTree;
const ssize_t chunkSize = int(args->coll.lastChunkSize);
const ssize_t size = args->coll.count;
const ssize_t chunkSize = int(args->lastChunkSize);
const ssize_t size = args->count;
const ssize_t loopSize = nChannels*tree->nHeads*chunkSize;
const int hasUp = (tree->up[0] >= 0) ? 1 : 0;
@@ -306,7 +309,7 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
const int nThreadsScatter = ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 2*COLLNET_COPY_THREADS : 0);
const int nThreadsGather = ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 1*COLLNET_COPY_THREADS : 0);
const int nThreadsBcast = ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 0 : 1*COLLNET_COPY_THREADS);
const int nThreadsReduce = args->nThreads - nThreadsScatter - nThreadsGather - nThreadsBcast;
const int nThreadsReduce = args->header.nWarps*WARP_SIZE - nThreadsScatter - nThreadsGather - nThreadsBcast;
const int tidStartBcast = nThreadsGather;
const int tidStartScatter = tidStartBcast + nThreadsBcast;
const int tidStartReduce = tidStartScatter + nThreadsScatter;
@@ -316,8 +319,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
if (tid >= tidStartScatter && tid < tidStartReduce && hasUp) {
// Scatter
int group = (2*Proto::MaxGroupWidth) | (1<<16);
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto>
prims(tid-tidStartScatter, nThreadsScatter, NULL, tree->up, args->sendbuff, args->recvbuff, args->coll.redOpArg, group, args);
Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid-tidStartScatter, nThreadsScatter, NULL, tree->up, args->sendbuff, args->recvbuff, args->redOpArg, group, args);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*tree->nHeads*chunkSize;
int nelem = min(tree->nHeads*chunkSize, size-offset);
@@ -331,8 +334,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
int group = (3*Proto::MaxGroupWidth) | (1<<16);
if (hasDn) {
// Reduce, send to network
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto>
prims(tid-tidStartReduce, nThreadsReduce, tree->down, &tree->out, args->sendbuff, args->recvbuff, args->coll.redOpArg, group, args);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>
prims(tid-tidStartReduce, nThreadsReduce, tree->down, &tree->out, args->sendbuff, args->recvbuff, args->redOpArg, group, args);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
@@ -344,8 +347,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
}
} else {
// Directly send to network
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto>
prims(tid-tidStartReduce, nThreadsReduce, nullptr, &tree->out, args->sendbuff, args->recvbuff, args->coll.redOpArg, group);
Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
prims(tid-tidStartReduce, nThreadsReduce, nullptr, &tree->out, args->sendbuff, args->recvbuff, args->redOpArg, group);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
@@ -355,8 +358,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
} else if (tid < tidStartBcast && hasUp) {
// Gather
int group = (0*Proto::MaxGroupWidth) | (0<<16);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/0, Proto>
prims(tid, nThreadsGather, tree->up, NULL, args->sendbuff, args->recvbuff, args->coll.redOpArg, group, args);
Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/0, Proto, 0>
prims(tid, nThreadsGather, tree->up, NULL, args->sendbuff, args->recvbuff, args->redOpArg, group, args);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + bid*tree->nHeads*chunkSize;
int nelem = min(tree->nHeads*chunkSize, size-offset);
@@ -366,8 +369,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
int group = (1*Proto::MaxGroupWidth) | (0<<16);
if (hasDn) {
// Recv from network, broadcast
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, tree->down, args->sendbuff, args->recvbuff, args->coll.redOpArg, group, args);
Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, tree->down, args->sendbuff, args->recvbuff, args->redOpArg, group, args);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
@@ -375,8 +378,8 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET, NCCL_PROTO
}
} else {
// Recv from network (no post thread needed)
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, nullptr, args->sendbuff, args->recvbuff, args->coll.redOpArg, group);
Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
prims(tid-tidStartBcast, nThreadsBcast, &tree->out, nullptr, args->sendbuff, args->recvbuff, args->redOpArg, group);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t offset = gridOffset + (bid*tree->nHeads+tree->headRank)*chunkSize;
int nelem = min(chunkSize, size-offset);
@@ -404,13 +407,15 @@ struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
runRing<T, RedOp, ProtoLL128>(args);
//LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
}
};
template<typename T, typename RedOp>
struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL128> {
__device__ __attribute__((noinline)) void run(ncclWorkElem *args) {
LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
runTreeSplit<T, RedOp, ProtoLL128>(args);
//LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
}
};
src/collectives/device/alltoall_pivot.h
+8 -8
View File
@@ -12,24 +12,24 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nranks = ncclShmem->comm.nRanks;
const ncclRing *ring = &ncclShmem->channel.ring;
const int num_bi_rings = args->coll.pivotA2ANumBiRings;
const int num_bi_rings = args->pivotA2ANumBiRings;
const int num_uni_rings = num_bi_rings * 2;
const int num_chunks = args->coll.nChannels / 2;
const int num_chunks = args->nChannels / 2;
const int chunk_id = (bid % num_bi_rings) + (bid / num_uni_rings * num_bi_rings);
const int elem_size = args->coll.count % 256 ? 1 : 256;
const ssize_t num_elems = args->coll.count / elem_size;
const int elem_size = args->count % 256 ? 1 : 256;
const ssize_t num_elems = args->count / elem_size;
const int num_padding_chunks = num_elems % num_chunks;
const ssize_t chunk_offset = elem_size * (num_elems / num_chunks * chunk_id + (chunk_id < num_padding_chunks ? chunk_id : num_padding_chunks));
const ssize_t chunk_size = elem_size * (num_elems / num_chunks + (chunk_id < num_padding_chunks ? 1 : 0));
const int pivot_direction = (bid % num_uni_rings) / num_bi_rings;
const ssize_t prims_size = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLTOALL_PIVOT_CHUNKSTEPS : 1));
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto> prims
(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, /*redOpArg(ignored)=*/0, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, /*redOpArg(ignored)=*/0, args->connIndex << 16);
for (int num_hops = 0; num_hops <= nranks / 2; num_hops++) {
const int src_rank = ring->devUserRanks[(nranks - num_hops) % nranks];
src/collectives/device/broadcast.h
+13 -10
View File
@@ -1,5 +1,5 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -12,27 +12,27 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclRing *ring = &ncclShmem->channel.ring;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? BROADCAST_CHUNKSTEPS : 1));
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
const int rank = ring->devUserRanks[0];
const int nextRank = ring->devUserRanks[1];
const int root = args->coll.root;
#ifdef ENABLE_PROFILING
auto devProf = ncclShmem->comm.devProf;
uint64_t clk, t0 = 0ULL, ws;
if (tid == 0) clk = __builtin_amdgcn_s_memrealtime();
#endif
const int root = args->root;
T *inputBuf = (T*)args->sendbuff;
T *outputBuf = (T*)args->recvbuff;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->coll.redOpArg, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->redOpArg, args->connIndex << 16);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
@@ -41,7 +41,7 @@ namespace {
realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
realChunkSize = size-gridOffset < loopSize ? args->lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
realChunkSize = int(realChunkSize);
@@ -70,7 +70,10 @@ namespace {
}
}
#ifdef ENABLE_PROFILING
if (tid == 0 && args->coll.opCount) devProf->elems[blockIdx.x].total_cycle += (__builtin_amdgcn_s_memrealtime() - clk);
if (tid == 0) {
struct ncclProfElem *elem = devProf.elems+args->opCount;
elem->elem[blockIdx.x].total_cycle += (__builtin_amdgcn_s_memrealtime() - clk);
}
#endif
}
}
src/collectives/device/common.h
+201 -270
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -10,6 +10,7 @@
#include "collectives.h"
#include "devcomm.h"
#include "op128.h"
#define COLL_UNROLL 2
#define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1) // Using balanced tree instead of split tree
@@ -67,40 +68,8 @@
NCCL_FUNCS3B(func, Sum), \
NCCL_FUNCS3B(func, Sum)
// [RCCL] Adding clique-based kernels for AllReduce, in-place of unused RingLL28 kernels
#define NCCL_FUNC5B(func, algo, devredop, type, nullify) \
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL, devredop, type)), \
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, LL128, devredop, type)), \
MACRO_IF(nullify, nullptr, NCCL_FUNC_NAME(func, algo, SIMPLE, devredop, type))
#define NCCL_FUNC4B(func, devredop, type, nullify) \
NCCL_FUNC5B(func, TREE, devredop, type, nullify), \
NCCL_FUNC5B(func, RING, devredop, type, nullify), \
NCCL_FUNC5B(func, COLLNET, devredop, type, nullify)
#define NCCL_FUNCS3C(func, devredop, nullForFloat) \
NCCL_FUNC4B(func, devredop, int8_t, 0), \
NCCL_FUNC4B(func, devredop, uint8_t, 0), \
NCCL_FUNC4B(func, devredop, int32_t, 0), \
NCCL_FUNC4B(func, devredop, uint32_t, 0), \
NCCL_FUNC4B(func, devredop, int64_t, 0), \
NCCL_FUNC4B(func, devredop, uint64_t, 0), \
NCCL_FUNC4B(func, devredop, half, nullForFloat), \
NCCL_FUNC4B(func, devredop, float, nullForFloat), \
NCCL_FUNC4B(func, devredop, double, nullForFloat), \
NCCL_FUNC4B(func, devredop, rccl_bfloat16, nullForFloat)
#define NCCL_FUNCS2C(func) \
NCCL_FUNCS3C(func, Sum, /*nullForFloat=*/0), \
NCCL_FUNCS3C(func, Prod, /*nullForFloat=*/0), \
NCCL_FUNCS3C(func, Max, /*nullForFloat=*/0), \
NCCL_FUNCS3C(func, Min, /*nullForFloat=*/0), \
NCCL_FUNCS3C(func, PreMulSum, /*nullForFloat=*/0), \
NCCL_FUNCS3C(func, SumPostDiv, /*nullForFloat=*/1)
// Must be consistent with the ncclFuncSet enum
using ncclKernelFunc_t = void (*)(struct ncclWorkElem* args);
using ncclKernelFunc_t = void (*)();
static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
// Don't try to initialize the host shadow copy of this device-side global
@@ -108,13 +77,13 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
// confuses clang. This will be fixed in the next clang release.
#if defined(__HIP_DEVICE_COMPILE__)
#if defined(BUILD_ALLREDUCE_ONLY)
NCCL_FUNC4B(AllReduce, Sum, float, 0),
NCCL_FUNC4(AllReduce, Sum, float, 0),
#else
NCCL_FUNCS2B(Broadcast),
NCCL_FUNCS2A(Reduce),
NCCL_FUNCS2B(AllGather),
NCCL_FUNCS2A(ReduceScatter),
NCCL_FUNCS2C(AllReduce),
NCCL_FUNCS2A(AllReduce),
NCCL_ONERANK_REDUCE_NAME(PreMulSum, int8_t),
NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint8_t),
NCCL_ONERANK_REDUCE_NAME(PreMulSum, int32_t),
@@ -136,18 +105,18 @@ static const __device__ constexpr ncclKernelFunc_t ncclFuncs[]{
template<unsigned short f, unsigned short l>
struct Caller {
static __device__ __host__
void call(struct ncclWorkElem* const c) noexcept
void call(unsigned short funcIndex) noexcept
{
constexpr unsigned short m = f + (l - f) / 2;
return (c->funcIndex < m) ? Caller<f, m>::call(c) : Caller<m, l>::call(c);
return (funcIndex < m) ? Caller<f, m>::call(funcIndex) : Caller<m, l>::call(funcIndex);
}
};
template<unsigned short f>
struct Caller<f, f + 1>{
static __device__ __host__
void call(struct ncclWorkElem* const c) noexcept { ncclFuncs[f](c); }
void call(unsigned short funcIndex) noexcept { ncclFuncs[f](); }
};
static_assert(FUNC_INDEX_P2P == 2710, "Wrong P2P function index");
@@ -155,86 +124,86 @@ static_assert(FUNC_INDEX_ALLTOALL_PIVOT == 2711, "Wrong AllToAllPivot function i
inline
__device__
void NCCL_CALL_FUNCTIONS(struct ncclWorkElem* const c) noexcept {
void NCCL_CALL_FUNCTIONS(unsigned short funcIndex) noexcept {
#if defined(BUILD_ALLREDUCE_ONLY)
if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_RING_SIMPLE_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL))
ncclFunction_AllReduce_RING_LL_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL128))
ncclFunction_AllReduce_RING_LL128_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_TREE_SIMPLE_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_LL))
ncclFunction_AllReduce_TREE_LL_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_COLLNET_SIMPLE_Sum_float(c);
else if (c->funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_LL))
ncclFunction_AllReduce_COLLNET_LL_Sum_float(c);
if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_RING_SIMPLE_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL))
ncclFunction_AllReduce_RING_LL_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_RING, NCCL_PROTO_LL128))
ncclFunction_AllReduce_RING_LL_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_TREE_SIMPLE_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_TREE, NCCL_PROTO_LL))
ncclFunction_AllReduce_TREE_LL_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE))
ncclFunction_AllReduce_COLLNET_SIMPLE_Sum_float();
else if (funcIndex == FUNC_INDEX(ncclFuncAllReduce, ncclSum, ncclFloat32, NCCL_ALGO_COLLNET, NCCL_PROTO_LL))
ncclFunction_AllReduce_COLLNET_LL_Sum_float();
else
assert("Unsupported function index");
#else
if (c->funcIndex < 540) {
if (c->funcIndex % 9 == 0) ncclFunction_Broadcast_TREE_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 2) ncclFunction_Broadcast_TREE_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex % 9 == 3) ncclFunction_Broadcast_RING_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 4) ncclFunction_Broadcast_RING_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 5) ncclFunction_Broadcast_RING_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex % 9 == 6) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 7) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t(c);
else ncclFunction_Broadcast_COLLNET_SIMPLE_Sum_int8_t(c);
if (funcIndex < 540) {
if (funcIndex % 9 == 0) ncclFunction_Broadcast_TREE_LL_Sum_int8_t();
else if (funcIndex % 9 == 1) ncclFunction_Broadcast_TREE_LL_Sum_int8_t();
else if (funcIndex % 9 == 2) ncclFunction_Broadcast_TREE_SIMPLE_Sum_int8_t();
else if (funcIndex % 9 == 3) ncclFunction_Broadcast_RING_LL_Sum_int8_t();
else if (funcIndex % 9 == 4) ncclFunction_Broadcast_RING_LL_Sum_int8_t();
else if (funcIndex % 9 == 5) ncclFunction_Broadcast_RING_SIMPLE_Sum_int8_t();
else if (funcIndex % 9 == 6) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t();
else if (funcIndex % 9 == 7) ncclFunction_Broadcast_COLLNET_LL_Sum_int8_t();
else ncclFunction_Broadcast_COLLNET_SIMPLE_Sum_int8_t();
}
else if (c->funcIndex < 1080) Caller<540, 1080>::call(c);
else if (c->funcIndex < 1620) {
if (c->funcIndex % 9 == 0) ncclFunction_AllGather_TREE_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 1) ncclFunction_AllGather_TREE_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 2) ncclFunction_AllGather_TREE_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex % 9 == 3) ncclFunction_AllGather_RING_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 4) ncclFunction_AllGather_RING_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 5) ncclFunction_AllGather_RING_SIMPLE_Sum_int8_t(c);
else if (c->funcIndex % 9 == 6) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t(c);
else if (c->funcIndex % 9 == 7) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t(c);
else ncclFunction_AllGather_COLLNET_SIMPLE_Sum_int8_t(c);
else if (funcIndex < 1080) Caller<540, 1080>::call(funcIndex);
else if (funcIndex < 1620) {
if (funcIndex % 9 == 0) ncclFunction_AllGather_TREE_LL_Sum_int8_t();
else if (funcIndex % 9 == 1) ncclFunction_AllGather_TREE_LL_Sum_int8_t();
else if (funcIndex % 9 == 2) ncclFunction_AllGather_TREE_SIMPLE_Sum_int8_t();
else if (funcIndex % 9 == 3) ncclFunction_AllGather_RING_LL_Sum_int8_t();
else if (funcIndex % 9 == 4) ncclFunction_AllGather_RING_LL_Sum_int8_t();
else if (funcIndex % 9 == 5) ncclFunction_AllGather_RING_SIMPLE_Sum_int8_t();
else if (funcIndex % 9 == 6) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t();
else if (funcIndex % 9 == 7) ncclFunction_AllGather_COLLNET_LL_Sum_int8_t();
else ncclFunction_AllGather_COLLNET_SIMPLE_Sum_int8_t();
}
else if (c->funcIndex < 2700) Caller<1620, 2700>::call(c);
else if (funcIndex < 2700) Caller<1620, 2700>::call(funcIndex);
else {
switch (c->funcIndex - 2700) {
switch (funcIndex - 2700) {
case 0:
ncclFunction_OneRankReduce_PreMulSum_int8_t(c);
ncclFunction_OneRankReduce_PreMulSum_int8_t();
break;
case 1:
ncclFunction_OneRankReduce_PreMulSum_uint8_t(c);
ncclFunction_OneRankReduce_PreMulSum_uint8_t();
break;
case 2:
ncclFunction_OneRankReduce_PreMulSum_int32_t(c);
ncclFunction_OneRankReduce_PreMulSum_int32_t();
break;
case 3:
ncclFunction_OneRankReduce_PreMulSum_uint32_t(c);
ncclFunction_OneRankReduce_PreMulSum_uint32_t();
break;
case 4:
ncclFunction_OneRankReduce_PreMulSum_int64_t(c);
ncclFunction_OneRankReduce_PreMulSum_int64_t();
break;
case 5:
ncclFunction_OneRankReduce_PreMulSum_uint64_t(c);
ncclFunction_OneRankReduce_PreMulSum_uint64_t();
break;
case 6:
ncclFunction_OneRankReduce_PreMulSum_half(c);
ncclFunction_OneRankReduce_PreMulSum_half();
break;
case 7:
ncclFunction_OneRankReduce_PreMulSum_float(c);
ncclFunction_OneRankReduce_PreMulSum_float();
break;
case 8:
ncclFunction_OneRankReduce_PreMulSum_double(c);
ncclFunction_OneRankReduce_PreMulSum_double();
break;
case 9:
ncclFunction_OneRankReduce_PreMulSum_rccl_bfloat16(c);
ncclFunction_OneRankReduce_PreMulSum_rccl_bfloat16();
break;
case 10:
ncclFunction_SendRecv_RING_SIMPLE_Sum_int8_t(c);
ncclFunction_SendRecv_RING_SIMPLE_Sum_int8_t();
break;
case 11:
ncclFunction_AllToAllPivot_RING_SIMPLE_Sum_int8_t(c);
ncclFunction_AllToAllPivot_RING_SIMPLE_Sum_int8_t();
default:
break;
}
@@ -249,45 +218,49 @@ class ncclFunction {
};
#ifdef ENABLE_COLLTRACE
#define traceColl(fIdx) \
uint32_t pos = __atomic_fetch_add(shmem.comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
shmem.comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
shmem.comm.collTrace[pos].bid = bid; \
shmem.comm.collTrace[pos].funcIndex = fIdx; \
if (fIdx == FUNC_INDEX_P2P) { \
shmem.comm.collTrace[pos].opCount = elems[0].p2p.opCount; \
shmem.comm.collTrace[pos].p2p.nThreads = elems[0].p2p.nThreads; \
shmem.comm.collTrace[pos].p2p.delta = (uint16_t)(elems[0].p2p.delta); \
#define traceColl(elem,launch_type) \
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
ncclShmem->comm.collTrace[pos].bid = blockIdx.x; \
ncclShmem->comm.collTrace[pos].funcIndex = ncclShmem->work.header.funcIndex; \
asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (ncclShmem->comm.collTrace[pos].data_0)); \
if (elem.header.type == ncclWorkTypeP2p) { \
struct ncclWorkElemP2p *p2pElems = (struct ncclWorkElemP2p *)&elem; \
ncclShmem->comm.collTrace[pos].p2p[0].connIndex = p2pElems[0].connIndex; \
ncclShmem->comm.collTrace[pos].p2pOpCount[0] = p2pElems[0].opCount; \
ncclShmem->comm.collTrace[pos].p2p[0].ngroups = p2pElems[0].ngroups; \
ncclShmem->comm.collTrace[pos].p2p[0].nWarps = p2pElems[0].nWarps; \
ncclShmem->comm.collTrace[pos].p2p[0].warpStart = p2pElems[0].warpStart; \
ncclShmem->comm.collTrace[pos].p2p[0].peer = (uint16_t)(p2pElems[0].peer); \
ncclShmem->comm.collTrace[pos].p2p[1].connIndex = p2pElems[1].connIndex; \
ncclShmem->comm.collTrace[pos].p2pOpCount[1] = p2pElems[1].opCount; \
ncclShmem->comm.collTrace[pos].p2p[1].ngroups = p2pElems[1].ngroups; \
ncclShmem->comm.collTrace[pos].p2p[1].nWarps = p2pElems[1].nWarps; \
ncclShmem->comm.collTrace[pos].p2p[1].warpStart = p2pElems[1].warpStart; \
ncclShmem->comm.collTrace[pos].p2p[1].peer = (uint16_t)(p2pElems[1].peer); \
ncclShmem->comm.collTrace[pos].type = (ncclCollTraceP2pElemType|launch_type); \
} else { \
shmem.comm.collTrace[pos].opCount = elems[0].coll.opCount; \
shmem.comm.collTrace[pos].coll.nThreads = elems[0].nThreads; \
shmem.comm.collTrace[pos].coll.bid = elems[0].coll.bid; \
shmem.comm.collTrace[pos].coll.nChannels = elems[0].coll.nChannels; \
ncclShmem->comm.collTrace[pos].opCount = elem.opCount; \
ncclShmem->comm.collTrace[pos].coll.nWarps = elem.header.nWarps; \
ncclShmem->comm.collTrace[pos].coll.bid = elem.bid; \
ncclShmem->comm.collTrace[pos].coll.nChannels = elem.nChannels; \
ncclShmem->comm.collTrace[pos].type = (ncclCollTraceCollElemType|launch_type); \
}
#define traceKernelLaunch(fIdx) { \
if (!(fIdx == FUNC_INDEX_P2P && elems[0].p2p.nThreads == 0)) { \
traceColl(fIdx); \
asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (shmem.comm.collTrace[pos].data_0)); \
shmem.comm.collTrace[pos].type = ncclCollTraceKernelLaunchType; \
} \
#define traceKernelLaunch(elem,firstLaunch) { \
traceColl(elem,(firstLaunch?ncclCollTraceKernelLaunchType:ncclCollTraceCollLaunchType)); \
}
#define traceCollEnd(fIdx) { \
if (!(fIdx == FUNC_INDEX_P2P && elems[0].p2p.nThreads == 0)) { \
traceColl(fIdx); \
shmem.comm.collTrace[pos].type = ncclCollTraceCollEndType; \
} \
#define traceKernelEnd() { \
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
ncclShmem->comm.collTrace[pos].bid = bid; \
ncclShmem->comm.collTrace[pos].type = ncclCollTraceKernelEndType; \
}
#define traceKernelEnd(fIdx) { \
if (!(fIdx == FUNC_INDEX_P2P && elems[0].p2p.nThreads == 0)) { \
traceColl(fIdx); \
shmem.comm.collTrace[pos].type = ncclCollTraceKernelEndType; \
} \
}
#define traceAbort(fIdx) { \
if (!(fIdx == FUNC_INDEX_P2P && elems[0].p2p.nThreads == 0)) { \
traceColl(fIdx); \
shmem.comm.collTrace[pos].type = ncclCollTraceAbortType; \
} \
#define traceAbort() { \
uint32_t pos = __atomic_fetch_add(ncclShmem->comm.collTraceTail, 1, __ATOMIC_SEQ_CST)%COLLTRACE_NUM_ITEMS; \
ncclShmem->comm.collTrace[pos].timeStamp = __builtin_amdgcn_s_memrealtime(); \
ncclShmem->comm.collTrace[pos].bid = bid; \
ncclShmem->comm.collTrace[pos].type = ncclCollTraceAbortType; \
}
// traceData(int16_t data2, uint32_t data4, uint64_t data8_0, uint64_t data8_1)
#define traceData(data2, data4, data8_0, data8_1) { \
@@ -301,9 +274,8 @@ class ncclFunction {
ncclShmem->comm.collTrace[pos].type = ncclCollTraceDataType; \
}
#else
#define traceKernelLaunch(fIdx)
#define traceCollEnd(fIdx)
#define traceAbort(fIdx)
#define traceKernelLaunch()
#define traceAbort()
#define traceData(data2, data4, data8_0, data8_1)
#endif
@@ -313,6 +285,28 @@ __device__ inline bool barrierReduceAny(int bit, uint32_t* abortCount) {
return atomicAdd(abortCount, 0) != 0;
}
// Copy src to dst and fill extra size with zeroes
template<typename Tdst, typename Tsrc>
__device__ void copyToShmem(Tdst *dst, Tsrc const *src, int tid, int nthreads) {
static_assert(sizeof(Tdst)%(2*sizeof(uint64_t)) == 0 && sizeof(Tsrc)%(2*sizeof(uint64_t)) == 0,
"copyToShmem needs sizes which are multiple of 16B");
static_assert(sizeof(Tdst) >= sizeof(Tsrc), "Tdst size is too small");
static_assert(sizeof(Tdst) <= WARP_SIZE*2*sizeof(uint64_t), "copyToShmem limited to 512B to make sure it can always be done in one cycle");
uint64_t *d = reinterpret_cast<uint64_t*>(dst);
uint64_t const *s = reinterpret_cast<uint64_t const*>(src);
uint64_t *shmemPtr = d;
int offset = 2*tid;
uint64_t v0, v1;
if (offset >= sizeof(Tsrc)/sizeof(uint64_t)) {
v0 = v1 = 0ULL;
} else {
v0 = s[offset] ; v1 = s[offset+1];
}
if (offset < sizeof(Tdst)/sizeof(uint64_t)) {
shmemPtr[offset] = v0; shmemPtr[offset+1] = v1;
}
}
template<typename T>
__device__ int copyToShmem(T *dst, T const *src, int turn=0) {
static_assert(sizeof(uint64_t) <= alignof(T), "Uhoh");
@@ -352,47 +346,21 @@ struct RunWorkElement {
}
};
#if CUDART_VERSION >= 11030
__device__ constexpr int ncclWorkElemFactors[NCCL_NUM_ALGORITHMS] =
#else
static __device__ __constant__ int ncclWorkElemFactors[NCCL_NUM_ALGORITHMS] =
#endif
{/*Tree*/1, /*Ring and P2P*/1, /*CollNet*/NCCL_REG_ELEM_FACTOR};
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
struct RunWork {
// This __forceinline__ is necessary. The compiler was inserting a function call
// here from the LL ncclKernel.
__device__ __forceinline__ void run(ncclWork *w) {
int tid = threadIdx.x;
/* Some invariants that must hold:
* 1. All elems[] have same funcIndex.
* 2. All elems[] have same nThreads.
* 3. The thread-to-group relation (as in prims group numbers) is the same
* for all elems[].
*
* If (1) isn't true then we might be in the wrong function since dispatch
* on ncclFuncs[w->funcIndex] is how we got here.
*
* If (2) or (3) aren't true, then threads from different work elements
* could race for barrier resources (barrier numbers 0...15) which is fatal.
*
* IMPORTANT!!! To ensure (3), implementations of
* `RunWorkElement<Fn,T,RedOp,Algo,Proto>::run()` may only use the following
* when deciding how to map threads to groups:
* Fn, T, RedOp, Algo, Proto, nThreads
*
* This last one is difficult to enforce so I hope everyone reads this.
*/
if (tid < w->elems[0].nThreads) {
#pragma unroll 1
for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].active != 0; e+=ncclWorkElemFactors[Algo])
int wid = threadIdx.x / WARP_SIZE;
int inc = w->header.type == ncclWorkTypeRegColl ? sizeof(ncclWorkElemReg) / sizeof(ncclWorkElem) : 1;
#pragma unroll 1
for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].header.type != ncclWorkTypeUnused; e += inc) {
if (wid < w->header.nWarps)
RunWorkElement<Fn, T, RedOp, Algo, Proto>().run(&w->elems[e]);
}
}
};
#define MAXWARPS (NCCL_MAX_NTHREADS/WARP_SIZE)
struct ncclShmemGroup {
ncclConnInfo *recvConns[NCCL_MAX_DIRECT_ARITY];
ncclConnInfo *sendConns[NCCL_MAX_DIRECT_ARITY];
@@ -400,7 +368,7 @@ struct ncclShmemGroup {
void* dsts[NCCL_MAX_DIRECT_ARITY+1];
int totalSendSize[NCCL_MAX_SLICE_PER_CHUNK];
uint64_t barrier;
uint64_t barrier_next[MAXWARPS];
uint64_t barrier_next[NCCL_MAX_GROUPS];
};
struct ncclShmemData {
@@ -408,18 +376,41 @@ struct ncclShmemData {
uint64_t ll128warp[NCCL_MAX_GROUPS][NCCL_MAX_GROUPS];
struct ncclShmemGroup groups[NCCL_MAX_GROUPS];
};
uint32_t sync[MAXWARPS];
uint32_t sync[NCCL_MAX_GROUPS];
uint64_t redOpArgs[NCCL_MAX_DIRECT_ARITY+1];
ncclDevComm comm;
ncclChannel channel;
ncclWork work;
struct ncclDevComm comm;
struct ncclChannel channel;
uint64_t pad;
struct ncclWork work;
};
static __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we) {
if (we->header.type != ncclWorkTypeUnused && we->redOpArgIsPtr) {
/* redOpArg is a pointer to the scalar value, so we'll dereference it
* here so that redOpArg holds the bits of the scalar going forward.
* The tricky thing is we don't know its type T since that's encoded in
* the funcIndex. Because it would be difficult to get sizeof(T) from
* funcIndex, we'll cheat and just dereference the largest possible size
* given the alignment of the pointer. We might be reading in more bytes
* than we need but that's harmless.
*/
if (we->redOpArg%2 != 0)
we->redOpArg = *reinterpret_cast<uint8_t*>(we->redOpArg);
else if (we->redOpArg%4 != 0)
we->redOpArg = *reinterpret_cast<uint16_t*>(we->redOpArg);
else if (we->redOpArg%8 != 0)
we->redOpArg = *reinterpret_cast<uint32_t*>(we->redOpArg);
else
we->redOpArg = *reinterpret_cast<uint64_t*>(we->redOpArg);
}
}
extern __device__ struct ncclShmemData *ncclShmem;
template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int FnIndex, bool COLLTRACE>
__device__ void ncclKernel(ncclWorkElem first) {
__device__ void ncclKernel(struct ncclDevComm* comm, ncclWorkElem first) {
int tid = threadIdx.x;
int nthreads = blockDim.x;
int bid = blockIdx.x;
__shared__ struct ncclShmemData shmem;
ncclShmem = &shmem;
@@ -428,110 +419,90 @@ __device__ void ncclKernel(ncclWorkElem first) {
abortCount = 0;
for (auto i = 0; i < NCCL_MAX_GROUPS; i++) {
shmem.groups[i].barrier = 0;
for (auto j = 0; j < MAXWARPS; j++) shmem.groups[i].barrier_next[j] = 0;
for (auto j = 0; j < NCCL_MAX_GROUPS; j++) shmem.groups[i].barrier_next[j] = 0;
}
}
__syncthreads();
int turn = copyToShmem(&shmem.comm, first.comm);
int turn = copyToShmem(&ncclShmem->comm, comm);
// get address of channel without incurring indirect load from ncclDevCom::channels
ncclChannel *channel = &((ncclDevCommAndChannels*)first.comm)->channels[bid];
turn = copyToShmem(&shmem.channel, channel, turn);
ncclChannel *channel = &((ncclDevCommAndChannels*)comm)->channels[bid];
turn = copyToShmem(&ncclShmem->channel, channel, turn);
// To optimize for latency, (only) the first operation is passed as argument.
if (bid == 0 && first.active != 0) {
turn = copyToShmem(&shmem.work.elems[0], &first, turn);
if (1 <= tid && tid < NCCL_MAX_WORK_ELEMENTS && tid % ncclWorkElemFactors[Algo] == 0) {
shmem.work.elems[tid].active = 0;
shmem.work.elems[tid].redOpArgIsPtr = 0;
}
if (bid == 0 && first.header.type != ncclWorkTypeUnused) {
// Copy first elem to work and zero out the rest
copyToShmem(&ncclShmem->work, &first, tid, nthreads);
}
struct ncclWorkElem* elems = shmem.work.elems;
__syncthreads(); // publish shmem
__syncthreads(); // publish ncclShmem
ncclWork *workFifoHost = shmem.channel.workFifo;
ncclWork *workFifoDev = shmem.channel.workFifoDev;
int workFifoIx = shmem.channel.index;
ncclWork *workFifoHost = ncclShmem->channel.workFifo;
ncclWork *workFifoDev = ncclShmem->channel.workFifoDev;
int workFifoIx = ncclShmem->channel.index;
bool skipLoadWork = false, firstLaunch = true;
if (bid == 0 && first.active != 0)
if (bid == 0 && first.header.type != ncclWorkTypeUnused)
skipLoadWork = true;
while (true) {
if (!skipLoadWork) {
copyToShmem(&shmem.work, &workFifoDev[workFifoIx]); // turn no longer helps
// Check whether the last operation was aborted and make sure all threads exit
int aborted = tid == 0 ? *shmem.comm.abortFlag : 0;
if (barrierReduceAny(aborted, &abortCount)) { // publish shmem.work
if (COLLTRACE && tid == 0) traceAbort(elems->funcIndex);
break;
copyToShmem(&ncclShmem->work, &workFifoDev[workFifoIx], tid, nthreads);
{ // Check whether the last operation was aborted and make sure all threads exit
int aborted = tid == 0 ? *comm->abortFlag : 0;
if (barrierReduceAny(aborted, &abortCount)) { // publish ncclShmem->work
if (COLLTRACE && tid == 0) traceAbort();
break;
}
if (tid == 0)
workFifoHost[workFifoIx].header.type = ncclWorkTypeUnused;
}
if (tid == 0)
workFifoHost[workFifoIx].elems[0].active = 0;
if (COLLTRACE && tid == 0) {
if (firstLaunch) traceKernelLaunch(elems->funcIndex);
if (!firstLaunch) traceCollEnd(elems->funcIndex);
firstLaunch = false;
}
} else if (COLLTRACE && tid == 0) {
traceKernelLaunch(elems->funcIndex);
firstLaunch = false;
}
workFifoIx = (workFifoIx + 1)%NCCL_MAX_OPS;
if (tid == 0)
channel->index = workFifoIx; // write back to real channel, not shmem shadow
if (tid < NCCL_MAX_WORK_ELEMENTS && tid % ncclWorkElemFactors[Algo] == 0) {
ncclWorkElem *we = &shmem.work.elems[tid];
if (we->redOpArgIsPtr && we->active != 0) {
/* redOpArg is a pointer to the scalar value, so we'll dereference it
* here so that redOpArg holds the bits of the scalar going forward.
* The tricky thing is we don't know its type T since that's encoded in
* the funcIndex. Because it would be difficult to get sizeof(T) from
* funcIndex, we'll cheat and just dereference the largest possible size
* given the alignment of the pointer. We might be reading in more bytes
* than we need but that's harmless.
*/
if (we->coll.redOpArg%2 != 0)
we->coll.redOpArg = *reinterpret_cast<uint8_t*>(we->coll.redOpArg);
else if (we->coll.redOpArg%4 != 0)
we->coll.redOpArg = *reinterpret_cast<uint16_t*>(we->coll.redOpArg);
else if (we->coll.redOpArg%8 != 0)
we->coll.redOpArg = *reinterpret_cast<uint32_t*>(we->coll.redOpArg);
else
we->coll.redOpArg = *reinterpret_cast<uint64_t*>(we->coll.redOpArg);
}
__syncwarp();
if (ncclShmem->work.header.type == ncclWorkTypeColl) {
if (tid < NCCL_MAX_WORK_ELEMENTS) ncclRedopPtrDeref(&ncclShmem->work.elems[tid]);
} else if (ncclShmem->work.header.type == ncclWorkTypeRegColl) {
if (tid < NCCL_MAX_WORK_ELEMENTS_REG) ncclRedopPtrDeref(&ncclShmem->work.regElems[tid].elem);
}
__syncthreads();
if (shmem.work.elems[0].funcIndex == FnIndex)
RunWork<Fn, T, RedOp, Algo, Proto>().run(&shmem.work);
else
NCCL_CALL_FUNCTIONS(&elems[0]);
if (shmem.work.elems[0].active == 2) {
if (COLLTRACE && tid == 0) traceKernelEnd(elems->funcIndex)
break;
if (COLLTRACE && tid == 0) {
traceKernelLaunch(ncclShmem->work.elems[0],firstLaunch);
firstLaunch = false;
#pragma unroll 1
for(int e=1; e < NCCL_MAX_WORK_ELEMENTS && ncclShmem->work.elems[e].header.type != ncclWorkTypeUnused; e ++) {
traceColl(ncclShmem->work.elems[e], 0);
}
}
if (ncclShmem->work.header.funcIndex == FnIndex)
RunWork<Fn, T, RedOp, Algo, Proto>().run(&ncclShmem->work);
else
NCCL_CALL_FUNCTIONS(ncclShmem->work.header.funcIndex);
if (ncclShmem->work.header.isLast) break;
__syncthreads();
skipLoadWork = false;
}
if (COLLTRACE && tid == 0) traceKernelEnd()
}
#define IMPL_COLL_KERN(func, algo, proto, devredop, type, fIndex) \
__launch_bounds__(NCCL_MAX_NTHREADS, 1) \
__global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(ncclWorkElem first) { \
if (first.comm->collTraceThread) \
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, true>(first); \
__global__ void NCCL_KERN_NAME(func, algo, proto, devredop, type)(struct ncclDevComm* comm, ncclWorkElem first) { \
if (comm->collTraceThread) \
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, true>(comm, first); \
else \
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, false>(first); \
ncclKernel<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto, fIndex, false>(comm, first); \
}
// Examples : AllReduce, RING, LL, Sum, uint8
/* Functions for aggregation case */
#define IMPL_COLL_FUNC(func, algo, proto, devredop, type) \
__device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, devredop, type)(struct ncclWorkElem* args) { \
__device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, devredop, type)() { \
RunWork<ncclFunc##func, type, Func##devredop<type>, NCCL_ALGO_##algo, NCCL_PROTO_##proto>().run(&ncclShmem->work); \
}
@@ -574,46 +545,6 @@ __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, dev
IMPL_COLL2(func, PreMulSum) \
IMPL_COLL2A(func, SumPostDiv)
// [RCCL] Define clique-based implementations (repurposed LL128)
#define IMPL_COLL4_CLIQUE(func, algo, devredop, type, ncclType) \
IMPL_COLL_FUNC(func, algo, LL, devredop, type) \
IMPL_COLL_FUNC(func, algo, LL128, devredop, type) \
IMPL_COLL_FUNC(func, algo, SIMPLE, devredop, type) \
#define IMPL_COLL3_CLIQUE(func, devredop, type, ncclType) \
IMPL_COLL4_CLIQUE(func, TREE, devredop, type, ncclType) \
IMPL_COLL4_CLIQUE(func, RING, devredop, type, ncclType) \
IMPL_COLL4_CLIQUE(func, COLLNET, devredop, type, ncclType)
#define IMPL_COLL2_CLIQUE(func, devredop) \
IMPL_COLL3_CLIQUE(func, devredop, int8_t, ncclInt8) \
IMPL_COLL3_CLIQUE(func, devredop, uint8_t, ncclUint8) \
IMPL_COLL3_CLIQUE(func, devredop, int32_t, ncclInt32) \
IMPL_COLL3_CLIQUE(func, devredop, uint32_t, ncclUint32) \
IMPL_COLL3_CLIQUE(func, devredop, int64_t, ncclInt64) \
IMPL_COLL3_CLIQUE(func, devredop, uint64_t, ncclUint64) \
IMPL_COLL3_CLIQUE(func, devredop, half, ncclFloat16) \
IMPL_COLL3_CLIQUE(func, devredop, float, ncclFloat32) \
IMPL_COLL3_CLIQUE(func, devredop, double, ncclFloat64) \
IMPL_COLL3_CLIQUE(func, devredop, rccl_bfloat16, ncclBfloat16)
#define IMPL_COLL2A_CLIQUE(func, devredop) \
IMPL_COLL3_CLIQUE(func, devredop, int8_t, ncclInt8) \
IMPL_COLL3_CLIQUE(func, devredop, uint8_t, ncclUint8) \
IMPL_COLL3_CLIQUE(func, devredop, int32_t, ncclInt32) \
IMPL_COLL3_CLIQUE(func, devredop, uint32_t, ncclUint32) \
IMPL_COLL3_CLIQUE(func, devredop, int64_t, ncclInt64) \
IMPL_COLL3_CLIQUE(func, devredop, uint64_t, ncclUint64)
#define IMPL_COLL_CLIQUE(func) \
IMPL_COLL2_CLIQUE(func, Sum) \
IMPL_COLL2_CLIQUE(func, Prod) \
IMPL_COLL2_CLIQUE(func, Min) \
IMPL_COLL2_CLIQUE(func, Max) \
IMPL_COLL2_CLIQUE(func, PreMulSum) \
IMPL_COLL2A_CLIQUE(func, SumPostDiv)
// [/RCCL]
// Copy primitives only define one function for copy
#define IMPL_COLL_C(func) IMPL_COLL3(func, Sum, int8_t, ncclInt8);
src/collectives/device/common_kernel.h
+14 -8
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -17,9 +17,15 @@
// Define min for ssize_t
static __device__ int min(int a, ssize_t b) { return (a < b) ? a : b; }
template <typename T>
inline __device__ void loadPtr(void** ptr, T* &v) {
inline __device__ int loadInt(int* ptr) {
int v;
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
v = LOAD(ptr);
#else
asm volatile("ld.volatile.global.u32 %0, [%1];"
: "=r"(v) : "l"(ptr));
#endif
return v;
}
typedef uint64_t PackType;
@@ -485,16 +491,16 @@ struct MULTI128 {
inline __device__ void Fetch128(Pack128& v, const Pack128* p) {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
v.x = p->x;
v.y = p->y;
v.x = __builtin_nontemporal_load(&p->x);
v.y = __builtin_nontemporal_load(&p->y);
#else
asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];" : "=l"(v.x), "=l"(v.y) : "l"(p) : "memory");
#endif
}
inline __device__ void Store128(Pack128* p, Pack128& v) {
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
p->x = v.x;
p->y = v.y;
__builtin_nontemporal_store(v.x, &p->x);
__builtin_nontemporal_store(v.y, &p->y);
#else
asm volatile("st.volatile.global.v2.u64 [%0], {%1,%2};" :: "l"(p), "l"(v.x), "l"(v.y) : "memory");
#endif
src/collectives/device/onerank_reduce.cu
+8 -8
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -17,11 +17,11 @@ namespace {
int tid = threadIdx.x;
int tn = blockDim.x;
#pragma unroll 1
for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].active != 0; e++) {
for(int e=0; e < NCCL_MAX_WORK_ELEMENTS && w->elems[e].header.type != ncclWorkTypeUnused; e++) {
ncclWorkElem *we = &w->elems[e];
intptr_t eltN = we->coll.count;
int bid = we->coll.bid;
int bn = we->coll.nChannels;
intptr_t eltN = we->count;
int bid = we->bid;
int bn = we->nChannels;
T const *src = (T const*)we->sendbuff;
T *dst = (T*)we->recvbuff;
@@ -37,13 +37,13 @@ namespace {
src += i0;
dst += i0;
ReduceOrCopyMulti<COLL_UNROLL, RedOp, T, 1, 1, 1, 1, 1>
(tid, tn, &(we->coll.redOpArg), true, 1, &src, 1, &dst, i1-i0);
(tid, tn, &(we->redOpArg), true, 1, &src, 1, &dst, i1-i0);
}
}
}
#define INSTANTIATE(devredop, type) \
__device__ void NCCL_ONERANK_REDUCE_NAME(devredop, type)(struct ncclWorkElem* args) { \
__device__ void NCCL_ONERANK_REDUCE_NAME(devredop, type)() { \
oneRankReduce<type, Func##devredop<type>>(); \
}
src/collectives/device/primitives.h
+7 -15
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -125,7 +125,7 @@ struct FanSymmetric {
};
// The primitives class. Specialized per protocol in the other headers.
template<typename T, typename RedOp, typename Fan, int Direct, typename Proto>
template<typename T, typename RedOp, typename Fan, int Direct, typename Proto, int P2p>
class Primitives;
// Used by LL & LL128 to implement direct members in the naive way.
@@ -157,20 +157,12 @@ struct PrimitivesWithoutDirect {
#include "prims_ll128.h"
#ifdef ENABLE_PROFILING
#ifdef ENABLE_TIMING_PROFILE
#define INIT_COUNTER \
if (tid == 0) { t0 = __builtin_amdgcn_s_memrealtime(); }
if (tid == 0) { struct ncclProfElem *elem = devProf.elems+args->opCount%PROFILE_NUM_ITEMS; t0 = __builtin_amdgcn_s_memrealtime(); ws = elem->elem[blockIdx.x].wait_cycle; }
#define ACCUMULATE_COUNTER(prim) \
if (tid == 0 && args->coll.opCount) { devProf->elems[blockIdx.x].prim##_cycle += (__builtin_amdgcn_s_memrealtime() - t0); \
devProf->elems[blockIdx.x].prim##_byte += nelem * sizeof(T); }
#else
#define INIT_COUNTER \
if (tid == 0) { t0 = __builtin_amdgcn_s_memrealtime(); ws = devProf->elems[blockIdx.x].wait_cycle; }
#define ACCUMULATE_COUNTER(prim) \
if (tid == 0 && args->coll.opCount) { devProf->elems[blockIdx.x].prim##_cycle += (__builtin_amdgcn_s_memrealtime() - t0 \
+ ws - devProf->elems[blockIdx.x].wait_cycle); \
devProf->elems[blockIdx.x].prim##_byte += nelem * sizeof(T); }
#endif
if (tid == 0) { struct ncclProfElem *elem = devProf.elems+args->opCount%PROFILE_NUM_ITEMS; elem->elem[blockIdx.x].prim##_cycle += (__builtin_amdgcn_s_memrealtime() - t0 \
+ ws - elem->elem[blockIdx.x].wait_cycle); \
elem->elem[blockIdx.x].prim##_byte += nelem * sizeof(T); elem->elem[blockIdx.x].opCount = args->opCount;}
#else
#define INIT_COUNTER
#define ACCUMULATE_COUNTER(prim)
src/collectives/device/prims_ll.h
+8 -15
View File
@@ -1,13 +1,13 @@
/*************************************************************************
* Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
template<typename T, typename RedOp, typename Fan, int Direct>
class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL>> {
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
@@ -45,7 +45,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(1+group));
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15-group));
#endif
}
@@ -123,7 +123,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
template<int BeginIx>
__device__ void readLLBeginAll(int offset, ncclLLFifoLine(&line)[MaxRecv]) {
#pragma unroll 1
#pragma unroll
for (int i=BeginIx; i < MaxRecv; i++) {
if (i < fan.nrecv()) {
union ncclLLFifoLine* src = recvPtr(i) + offset;
@@ -290,14 +290,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
// Always waitSend in case of cleanup
nelem = nelem < 0 ? 0 : nelem;
#ifdef ENABLE_PROFILING
uint64_t t0;
if (tid == 0) t0 = __builtin_amdgcn_s_memrealtime();
#endif
if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine));
#ifdef ENABLE_PROFILING
if (SEND && tid == 0) ncclShmem->comm.devProf->elems[blockIdx.x].wait_cycle = (__builtin_amdgcn_s_memrealtime() - t0);
#endif
nelem -= tid*EltPerLine;
srcElts += tid*EltPerLine;
@@ -324,7 +317,7 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL>:
}
if (RECV) {
data = !SRC ? peerData : MULTI<RedOp,T>()(redOp, peerData, data);
#pragma unroll 1
#pragma unroll MaxRecv
for (int i=1; i < MaxRecv && i < fan.nrecv(); i++) {
peerData = readLLFinish(offset, line, i);
data = MULTI<RedOp,T>()(redOp, peerData, data);
src/collectives/device/prims_ll128.h
+10 -6
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -11,9 +11,9 @@
#define __any_sync(WARP_MASK, needReload) (true)
template<typename T, typename RedOp, typename Fan, int Direct>
class Primitives<T, RedOp, Fan, Direct, ProtoLL128>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128>> {
template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
@@ -52,7 +52,11 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128>:
inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; }
inline __device__ void barrier() {
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(1+group));
#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
__syncthreads();
#else
asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15-group));
#endif
}
uint32_t abort = 0;
src/collectives/device/prims_simple.h
+35 -40
View File
@@ -1,14 +1,14 @@
/*************************************************************************
* Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
template<typename T, typename RedOp, typename Fan, int Direct,
int SlicePerChunk, int StepPerSlice, int Unroll>
int SlicePerChunk, int StepPerSlice, int Unroll, int P2p>
class Primitives<
T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll>
T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll>, P2p
> {
static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
static constexpr int Input=0, Output=1;
@@ -19,7 +19,7 @@ class Primitives<
RolePostSend = 0x10,
RolePostRecv = 0x20,
Aborted = 0x40,
PtrsFifoEnabled = 0x80,
OffsFifoEnabled = 0x80,
SizesFifoEnabled = 0x100,
DirectWrite = 0x200,
DirectRead = 0x400,
@@ -33,10 +33,10 @@ class Primitives<
int flags;
int group;
uint64_t step;
int *connOffsFifoPtr; // (flags & OffsFifoEnabled)
union {
void **connPtrsFifoPtr; // (flags & PtrsFifoEnabled)
T *userBuff; // (flags & (RoleInput|RoleOutput))
T *connEltsFifo; // !(flags & (PtrsFifoEnabled|RoleInput|RoleOutput))
T *connEltsFifo; // !(flags & (RoleInput|RoleOutput))
};
union {
int volatile *connSizesFifoPtr; // (flags & SizesFifoEnabled)
@@ -59,7 +59,7 @@ class Primitives<
if (nthreads == WARP_SIZE)
__syncwarp();
else
asm volatile("bar.sync %0, %1;" :: "r"(group+1), "r"(nthreads));
asm volatile("bar.sync %0, %1;" :: "r"(15-group), "r"(nthreads));
#endif
flags |= ThreadsSynced;
}
@@ -70,7 +70,7 @@ class Primitives<
if (nworkers == nthreads)
barrier();
else
asm volatile("bar.sync %0, %1;" :: "r"(group+2), "r"(nworkers));
asm volatile("bar.sync %0, %1;" :: "r"(8-group), "r"(nworkers));
#endif
}
@@ -88,9 +88,6 @@ class Primitives<
const bool isSendNotRecv = (Send && Recv) ? (flags & RoleWaitSend) : Send;
const bool noRecvWait = DirectRecv && Src && (flags & DirectRead); // no wait when directly reading from remote input
const bool noSendWait = DirectSend && (flags & (DirectRead|DirectWrite)); // no wait in empty send (e.g. directScatter) or direct remote write
#if defined(ENABLE_PROFILING) && !defined(ENABLE_TIMING_PROFILE)
uint64_t t0 = __builtin_amdgcn_s_memrealtime();
#endif
if (((flags & (Recv*RoleWaitRecv)) && !noRecvWait) ||
((flags & (Send*RoleWaitSend)) && !noSendWait)) {
int spins = 0;
@@ -109,8 +106,8 @@ class Primitives<
void **ptrs = isSendNotRecv ? (ncclShmem->groups[group].dsts + Dst)
: (ncclShmem->groups[group].srcs + Src);
if (flags & PtrsFifoEnabled)
loadPtr(connPtrsFifoPtr + step%NCCL_STEPS, ptrs[index]);
if (flags & OffsFifoEnabled)
ptrs[index] = connEltsFifo + loadInt(connOffsFifoPtr + (step%NCCL_STEPS))/sizeof(T);
else if (isSendNotRecv && DirectSend) {
if (flags & DirectWrite) {
ptrs[index] = directBuff + remoteIx + offset;
@@ -132,14 +129,6 @@ class Primitives<
ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
}
step += StepPerSlice;
#if defined(ENABLE_PROFILING) && !defined(ENABLE_TIMING_PROFILE)
if (opCount) {
if (isSendNotRecv)
ncclShmem->comm.devProf->elems[blockIdx.x].wait_send_cycle += (__builtin_amdgcn_s_memrealtime() - t0);
else
ncclShmem->comm.devProf->elems[blockIdx.x].wait_recv_cycle += (__builtin_amdgcn_s_memrealtime() - t0);
}
#endif
}
}
@@ -204,7 +193,10 @@ class Primitives<
waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(dstIx, remoteIx, offset, sliceSize);
subBarrier();
#ifdef ENABLE_PROFILING
if (tid == 0 && opCount) ncclShmem->comm.devProf->elems[blockIdx.x].wait_cycle += (__builtin_amdgcn_s_memrealtime() - t0);
if (tid == 0) {
struct ncclProfElem *elem = ncclShmem->comm.devProf.elems+opCount%PROFILE_NUM_ITEMS;
elem->elem[blockIdx.x].wait_cycle += (__builtin_amdgcn_s_memrealtime() - t0);
}
#endif
if (DirectRecv && ncclShmem->groups[group].srcs[0] == ncclShmem->groups[group].dsts[0]) {
// We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
@@ -262,6 +254,8 @@ class Primitives<
}
// Scatter/Gather generic op
// skip: my own rank order in the buffer chunks
// shift: peer offset to avoid all ranks sending to or receiving from same peer
template <int DirectRecv1, int DirectSend1, int Recv, int Send>
__device__ __forceinline__ void
ScatterGatherOp(intptr_t inpIx, intptr_t outIx, int totalElem, int peerElem, int skip, int shift, bool postOp) {
@@ -287,11 +281,13 @@ class Primitives<
for (int j=0; j<fan.nsend(); j++) {
int i = (j+shift)%fan.nsend();
int peerOffset = i*peerElem;
// Skip the data I am responsible of reducing myself
if (skip >= 0 && i >= skip) peerOffset += peerElem;
const T* src0 = (T*)ncclShmem->groups[group].srcs[0] + peerOffset;
int realPeerSize = min(realSize, totalElem-peerOffset);
if (realPeerSize > 0 && ncclShmem->groups[group].dsts[i] != nullptr) {
ReduceOrCopyMulti<Unroll, RedOp, T, 1, 1, 1, 1, PreOpN>(tid, nworkers, ncclShmem->redOpArgs, false, 1, &src0, 1, (T**)ncclShmem->groups[group].dsts+i, realPeerSize);
// Mark for threadfence at the end
if (tid == 0) ncclShmem->groups[group].totalSendSize[slice] += realPeerSize;
}
}
@@ -319,6 +315,7 @@ class Primitives<
}
}
barrier();
// If we indeed send something, threadfence
if (Send && (flags & RolePostSend) && ncclShmem->groups[group].totalSendSize[slice] > 0 && index == 0)
__threadfence_system();
__syncwarp();
@@ -340,18 +337,18 @@ class Primitives<
ncclShmem->groups[group].recvConns[index] = conn; // WaitRecv role saves since that's who needs it in setDataPtrs()
connStepPtr = conn->tail;
connStepCache = LOAD(connStepPtr);
flags |= (conn->ptrsFifo != nullptr) ? PtrsFifoEnabled : 0;
flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0;
if (Direct) {
// User buffers have been registered
if ((conn->direct & (NCCL_IPC_READ|NCCL_IPC_WRITE)) && e != nullptr && e->regUsed) {
if (connIndex == 1) {
if (connIndex == 1 && P2p == 0) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
flags |= (e->direct & NCCL_DIRECT_WRITE) ? DirectWrite :
(e->direct & NCCL_DIRECT_READ) ? DirectRead : 0;
}
} else if (conn->direct & (NCCL_DIRECT_WRITE|NCCL_DIRECT_READ)) {
if (connIndex == 1) {
if (connIndex == 1 && P2p == 0) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
// direct read not allowed in non-register case
@@ -360,10 +357,9 @@ class Primitives<
}
}
}
if (flags & PtrsFifoEnabled)
connPtrsFifoPtr = conn->ptrsFifo;
else
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (flags & OffsFifoEnabled)
connOffsFifoPtr = conn->offsFifo;
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
}
}
}
@@ -380,11 +376,10 @@ class Primitives<
ncclShmem->groups[group].sendConns[index] = conn; // WaitSend role saves since that's who needs it in setDataPtrs()
connStepPtr = conn->head;
connStepCache = LOAD(connStepPtr);
flags |= (conn->ptrsFifo != nullptr) ? PtrsFifoEnabled : 0;
if (flags & PtrsFifoEnabled)
connPtrsFifoPtr = conn->ptrsFifo;
else
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0;
if (flags & OffsFifoEnabled)
connOffsFifoPtr = conn->offsFifo;
connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
if (conn->sizesFifo != nullptr) {
flags |= SizesFifoEnabled;
@@ -392,14 +387,14 @@ class Primitives<
} else if (Direct) {
// User buffers have been registered
if ((conn->direct & (NCCL_IPC_READ|NCCL_IPC_WRITE)) && e != nullptr && e->regUsed) {
if (connIndex == 1) {
if (connIndex == 1 && P2p == 0) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
flags |= (e->direct & NCCL_DIRECT_WRITE) ? DirectWrite :
(e->direct & NCCL_DIRECT_READ) ? DirectRead : 0;
}
} else if (conn->direct & (NCCL_DIRECT_WRITE|NCCL_DIRECT_READ)) {
if (connIndex == 1) {
if (connIndex == 1 && P2p == 0) {
flags |= DirectRead; // scatter-reduce use direct pull
} else {
// direct read not allowed in non-register case
@@ -419,7 +414,7 @@ class Primitives<
):
tid(tid),
stepSize(ncclShmem->comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T)),
opCount(ncclShmem->work.elems[0].coll.opCount) {
opCount(ncclShmem->work.elems[0].opCount) {
// For send operations, we need an extra warp to overlap the threadfence and the copy
this->nthreads = nthreads;
@@ -460,7 +455,7 @@ class Primitives<
loadRecvConn(&ncclShmem->channel.devPeers[peer], connIndex, e);
loadSendConn(&ncclShmem->channel.devPeers[peer], connIndex, e);
setDataPtrs(inputBuf, outputBuf, redOpArg, (struct ncclWorkRegElem*)e);
setDataPtrs(inputBuf, outputBuf, redOpArg, (struct ncclWorkElemReg*)e);
}
__device__ ~Primitives() {
@@ -477,7 +472,7 @@ class Primitives<
barrier();
}
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf, uint64_t redOpArg, struct ncclWorkRegElem* e) {
__device__ void setDataPtrs(void const *inputBuf, void *outputBuf, uint64_t redOpArg, struct ncclWorkElemReg* e) {
if (flags & RoleInput) {
userBuff = (T*)inputBuf;
ncclShmem->redOpArgs[0] = redOpArg; // scaler for local input
src/collectives/device/reduce.h
+10 -10
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -13,21 +13,21 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclRing *ring = &ncclShmem->channel.ring;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCE_CHUNKSTEPS : 1));
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
const int nranks = ncclShmem->comm.nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
const int rank = ncclShmem->comm.rank;
const int prevRank = ring->devUserRanks[nranks-1];
const int root = args->coll.root;
const int root = args->root;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->coll.redOpArg, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->redOpArg, args->connIndex << 16);
auto calcChunkSize = [&]__device__(ssize_t gridOffset)->int {
int realChunkSize;
@@ -36,7 +36,7 @@ namespace {
realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
realChunkSize = size-gridOffset < loopSize ? args->lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128, chunkSize);
return realChunkSize;
src/collectives/device/reduce_scatter.h
+9 -9
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -13,9 +13,9 @@ namespace {
template<typename T, typename RedOp, typename Proto>
__device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
const int tid = threadIdx.x;
const int nthreads = args->nThreads;
const int bid = args->coll.bid;
const int nChannels = args->coll.nChannels;
const int nthreads = args->header.nWarps*WARP_SIZE;
const int bid = args->bid;
const int nChannels = args->nChannels;
ncclRing *ring = &ncclShmem->channel.ring;
int const *ringRanks = ring->devUserRanks;
const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCESCATTER_CHUNKSTEPS : 1));
@@ -23,10 +23,10 @@ namespace {
const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2);
const int nranks = ncclShmem->comm.nRanks;
const ssize_t loopSize = nChannels*chunkSize;
const ssize_t size = args->coll.count;
const ssize_t size = args->count;
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->coll.redOpArg, args->coll.connIndex << 16);
Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->redOpArg, args->connIndex << 16);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
ssize_t realChunkSize;
@@ -35,7 +35,7 @@ namespace {
realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
}
else if (Proto::Id == NCCL_PROTO_LL)
realChunkSize = size-gridOffset < loopSize ? args->coll.lastChunkSize : chunkSize;
realChunkSize = size-gridOffset < loopSize ? args->lastChunkSize : chunkSize;
else if (Proto::Id == NCCL_PROTO_LL128)
realChunkSize = min(divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128, chunkSize);
realChunkSize = int(realChunkSize);
src/collectives/device/sendrecv.h
+62 -68
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -11,73 +11,67 @@
template<typename T, typename RedOp>
struct RunWork<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
__device__ __attribute__((noinline)) void run(ncclWork *work) {
int tid = threadIdx.x;
int group = 0;
const int rank = ncclShmem->comm.rank;
const int nRanks = ncclShmem->comm.nRanks;
using Proto = ProtoSimple<1, 1>;
for (int s=0; s<NCCL_MAX_WORK_ELEMENTS; s++) {
ncclWorkElem *args = &work->elems[s];
int nThreadsSegment = args->p2p.nThreads;
if (args->active == 0 || nThreadsSegment == 0) break;
int nThreadsSplit = nThreadsSegment/2;
int groupRecv = group;
group += Proto::calcGroupWidth(/*send=*/false, nThreadsSplit);
int groupSend = group;
group += Proto::calcGroupWidth(/*send=*/true, nThreadsSegment - nThreadsSplit);
if (tid < nThreadsSegment) {
// Compute pointers
T const* sendbuff = (const T*)args->sendbuff;
T* recvbuff = (T*)args->recvbuff;
ssize_t const sendCount = args->p2p.sendCount;
ssize_t const recvCount = args->p2p.recvCount;
int const delta = args->p2p.delta;
if (delta == 0) {
if (sendbuff != recvbuff) {
ReduceOrCopyMulti<COLL_UNROLL, RedOp, T, 1, 1, 1, 1, 0>(tid, nThreadsSegment, nullptr, false, 1, &sendbuff, 1, &recvbuff, sendCount);
}
}
else {
if ((tid < nThreadsSplit) && recvCount >= 0) {
int const peer = (rank - delta + nRanks)%nRanks;
int const t0 = 0;
int const nt = nThreadsSplit;
int const chunkSize = args->p2p.recvChunkSize/sizeof(T);
Primitives<T, RedOp, FanAsymmetric<1, 0>, 0, Proto> prims
(tid-t0, nt, &peer, nullptr, nullptr, recvbuff, /*redOpArg(ignored)=*/0, groupRecv | (args->p2p.recvIdx << 16));
ssize_t offset = 0;
do {
int nelem = roundUp(chunkSize, nt*(sizeof(uint64_t)/sizeof(T)));
nelem = min(chunkSize, recvCount-offset);
prims.directRecv(offset, nelem);
offset += nelem;
} while(offset < recvCount);
}
if ((tid >= nThreadsSplit) && sendCount >= 0) {
int const peer = (rank + delta)%nRanks;
int const t0 = nThreadsSplit;
int const nt = nThreadsSegment - nThreadsSplit;
int const chunkSize = args->p2p.sendChunkSize/sizeof(T);
Primitives<T, RedOp, FanAsymmetric<0, 1>, 0, Proto> prims
(tid-t0, nt, nullptr, &peer, sendbuff, nullptr, /*redOpArg(ignored)=*/0, groupSend | (args->p2p.sendIdx << 16));
ssize_t offset = 0;
do {
int nelem = roundUp(chunkSize, nt*(sizeof(uint64_t)/sizeof(T)));
nelem = min(chunkSize, sendCount-offset);
prims.directSend(offset, offset, nelem);
offset += nelem;
} while(offset < sendCount);
}
}
break;
__device__ __forceinline__ void runSend(const int tid, const int nthreads, const int group, struct ncclWorkElemP2p* args) {
if (args->peer == ncclShmem->comm.rank) {
struct ncclWorkElemP2p* recvArgs = args-1;
if (args->buff != recvArgs->buff) {
ReduceOrCopyMulti<COLL_UNROLL, RedOp, T, 1, 1, 1, 1, 0>(tid, nthreads, nullptr, false, 1, (const T**)&args->buff, 1, (T**)&recvArgs->buff, args->count);
}
tid -= nThreadsSegment;
} else {
using Proto = ProtoSimple<1, 1>;
ssize_t const count = args->count;
int const chunkSize = args->chunkSize/sizeof(T);
int const peer = args->peer;
Primitives<T, RedOp, FanAsymmetric<0, 1>, 1, Proto, 1> prims
(tid, nthreads, nullptr, &peer, args->buff, nullptr, /*redOpArg(ignored)=*/0, group);
ssize_t offset = 0;
do {
int nelem = min(chunkSize, count-offset);
prims.directSend(offset, offset, nelem);
offset += nelem;
} while(offset < count);
}
}
__device__ __forceinline__ void runRecv(const int tid, const int nthreads, const int group, struct ncclWorkElemP2p* args) {
if (args->peer != ncclShmem->comm.rank) {
using Proto = ProtoSimple<1, 1>;
ssize_t const count = args->count;
int const chunkSize = args->chunkSize/sizeof(T);
int const peer = args->peer;
Primitives<T, RedOp, FanAsymmetric<1, 0>, 1, Proto, 1> prims
(tid, nthreads, &peer, nullptr, nullptr, args->buff, /*redOpArg(ignored)=*/0, group);
ssize_t offset = 0;
do {
int nelem = min(chunkSize, count-offset);
prims.directRecv(offset, nelem);
offset += nelem;
} while(offset < count);
}
}
__device__ __forceinline__ void run(ncclWork *work) {
struct ncclWorkElemP2p* args = work->p2pElems;
int ngroups = args->ngroups;
int tid = threadIdx.x;
int wid = tid / WARP_SIZE;
// This has to work even for groups of 2.5 warps (which is 8 groups, and means 3
// warps for send, 2 warps for recv).
// warpStarts were rounded thanks to int division, but for group number we need to round the other way around
// So we mirror wid then mirror again the group.
#define NWARPS (NCCL_MAX_NTHREADS/WARP_SIZE)
int group = ngroups-1- (NWARPS-1-wid) * ngroups / NWARPS;
args += group;
if (args->header.type == ncclWorkTypeUnused) return;
tid -= args->warpStart * WARP_SIZE;
int nthreads = args->nWarps * WARP_SIZE;
group |= (args->connIndex<<16); // Used to select connIndex 1
if (tid >= nthreads || args->peer == -1) return;
if ((group%2) == 0) {
runRecv(tid, nthreads, group, args);
} else {
runSend(tid, nthreads, group, args);
}
}
};
src/collectives/sendrecv_api.cc
+5 -5
View File
@@ -1,6 +1,6 @@
/*************************************************************************
* Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
* Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
* Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
@@ -14,8 +14,8 @@ NCCL_API(ncclResult_t, ncclSend, const void* sendbuff, size_t count, ncclDataTyp
ncclResult_t ncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer,
ncclComm_t comm, hipStream_t stream) {
NVTX3_FUNC_RANGE_IN(nccl_domain);
struct ncclInfo info = { ncclFuncSendRecv, "Send",
sendbuff, NULL, count, datatype, ncclSum, peer, comm, stream, /* Args */
struct ncclInfo info = { ncclFuncSend, "Send",
NULL, (void*)sendbuff, count, datatype, ncclSum, peer, comm, stream, /* Args */
1, 1 };
ncclResult_t ret;
NCCLCHECK(ncclGroupStart());
@@ -29,7 +29,7 @@ NCCL_API(ncclResult_t, ncclRecv, void* recvbuff, size_t count, ncclDataType_t da
ncclResult_t ncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer,
ncclComm_t comm, hipStream_t stream) {
NVTX3_FUNC_RANGE_IN(nccl_domain);
struct ncclInfo info = { ncclFuncSendRecv, "Recv",
struct ncclInfo info = { ncclFuncRecv, "Recv",
NULL, recvbuff, count, datatype, ncclSum, peer, comm, stream, /* Args */
1, 1 };
ncclResult_t ret;