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

Merge pull request #134 from changpeng/master

Browse Source 
Tuning the inline and unroll to reduce the scratch usage

[ROCm/rccl commit: 662281e599]
This commit is contained in:
Wenkai Du
2019-10-09 10:58:38 -07:00
committed by GitHub
parent de13a48f7b d8a06589c9
commit c4ed3d2e08
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rccl/src/collectives/device/all_reduce.h
+3 -4
View File
@@ -102,8 +102,7 @@ __device__ void ncclAllReduceRingKernel(struct CollectiveArgs* args) {
#endif #endif
} }
template<int UNROLL, class FUNC, typename T> template<int UNUSED, class FUNC, typename T>
__attribute__((noinline))
__device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) { __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
const int nthreads = blockDim.x; const int nthreads = blockDim.x;
@@ -122,7 +121,7 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
do { do {
// Reduce : max number of recv is 3, max number of send is 1 (binary tree + local) // Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, nthreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount); ncclPrimitives<1, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, nthreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) { for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Up // Up
ssize_t offset = gridOffset + bid*chunkSize; ssize_t offset = gridOffset + bid*chunkSize;
@@ -139,7 +138,7 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
do { do {
// Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local) // Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
ncclPrimitives<UNROLL, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, nthreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount); ncclPrimitives<1, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, nthreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount);
for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) { for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
// Down // Down
ssize_t offset = gridOffset + bid*chunkSize; ssize_t offset = gridOffset + bid*chunkSize;