Revert "Tuning the inline and unroll to reduce the scratch usage"

This reverts commit d8a06589c9.


[ROCm/rccl commit: ca493a6b51]
This commit is contained in:
Wenkai Du
2020-05-15 14:15:40 -07:00
parent 5780637fff
commit e7eff47be4
@@ -102,7 +102,7 @@ __device__ void ncclAllReduceRingKernel(struct CollectiveArgs* args) {
#endif #endif
} }
template<int UNUSED, class FUNC, typename T> template<int UNROLL, class FUNC, typename T>
__attribute__((noinline)) __attribute__((noinline))
__device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) { __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
const int tid = threadIdx.x; const int tid = threadIdx.x;
@@ -128,7 +128,7 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
struct ncclTree* tree = &channel->treeUp; struct ncclTree* tree = &channel->treeUp;
// 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)
ncclPrimitivesRecvData<T, NCCL_MAX_TREE_ARITY> recvData; ncclPrimitivesRecvData<T, NCCL_MAX_TREE_ARITY> recvData;
ncclPrimitives<1, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, args->nThreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount, recvData); ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, args->nThreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount, recvData);
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;
@@ -147,7 +147,7 @@ __device__ void ncclAllReduceTreeKernel(struct CollectiveArgs* args) {
struct ncclTree* tree = &channel->treeDn; struct ncclTree* tree = &channel->treeDn;
// 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)
ncclPrimitivesSendData<T, NCCL_MAX_TREE_ARITY> sendData; ncclPrimitivesSendData<T, NCCL_MAX_TREE_ARITY> sendData;
ncclPrimitives<1, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, args->nThreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount, sendData); ncclPrimitives<UNROLL, 1, 1, T, 1, NCCL_MAX_TREE_ARITY, FUNC> prims(tid, args->nThreads, &tree->up, tree->down, NULL, stepSize, channel, comm, args->opCount, sendData);
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;