Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[release/2.5][ROCm] Fix largeIndexBlockSize #1659

Open
wants to merge 1 commit into
base: release/2.5
Choose a base branch
from
+15 −10
Open

[release/2.5][ROCm] Fix largeIndexBlockSize #1659

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 15 additions & 10 deletions aten/src/ATen/native/cuda/Indexing.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,15 @@
#include <c10/macros/Macros.h>

namespace {
constexpr uint64_t getDefaultMaxThreadsPerBlock() {
#ifndef USE_ROCM
return 128;
#else
// bigger default
return 512;
#endif
}

template <typename scalar_t, int SZ>
__global__ void indexing_backward_kernel(
const int64_t* sorted_indices, const int64_t* indices, const scalar_t* grad_output, scalar_t* grad_weight,
Expand Down Expand Up @@ -933,11 +942,13 @@ void index_add_cuda_impl(const Tensor& self, int64_t dim, const Tensor& index, c
selfAddDimSize, selfNumel, reduce_add, alpha_value); \
C10_CUDA_KERNEL_LAUNCH_CHECK();

uint64_t defaultMaxBlockThreads = getDefaultMaxThreadsPerBlock();
const dim3 smallIndexGrid(std::min(ceil_div(sliceSize, (uint64_t)128), (uint64_t)(mpc * 8)));
const dim3 smallIndexBlock(std::min(sliceSize, (uint64_t)128));

const dim3 largeIndexGrid(std::min(ceil_div(sourceTotalSize, (uint64_t)128), (uint64_t)(mpc * 8)));
const dim3 largeIndexBlock(std::min(sourceTotalSize, (uint64_t)128));
//On ROCm, std::min -> ::min did not work as expected on when outTotalSize>=2147483648
dim3 largeIndexBlock( (sourceTotalSize < defaultMaxBlockThreads) ? sourceTotalSize : defaultMaxBlockThreads );

if (cuda::detail::canUse32BitIndexMath(result) &&
cuda::detail::canUse32BitIndexMath(source) &&
Expand Down Expand Up @@ -1106,11 +1117,13 @@ void index_reduce_func_cuda_impl(
selfReduceDimSize, selfNumel, reduce_func, alpha_value); \
C10_CUDA_KERNEL_LAUNCH_CHECK();

uint64_t defaultMaxBlockThreads = getDefaultMaxThreadsPerBlock();
dim3 smallIndexGrid(std::min(ceil_div(sliceSize, (uint64_t)128), (uint64_t)(mpc * 8)));
dim3 smallIndexBlock(std::min(sliceSize, (uint64_t)128));

dim3 largeIndexGrid(std::min(ceil_div(sourceTotalSize, (uint64_t)128), (uint64_t)(mpc * 8)));
dim3 largeIndexBlock(std::min(sourceTotalSize, (uint64_t)128));
//On ROCm, std::min -> ::min did not work as expected on when outTotalSize>=2147483648
dim3 largeIndexBlock( (sourceTotalSize < defaultMaxBlockThreads) ? sourceTotalSize : defaultMaxBlockThreads );

if (cuda::detail::canUse32BitIndexMath(result) &&
cuda::detail::canUse32BitIndexMath(source) &&
Expand Down Expand Up @@ -1334,14 +1347,6 @@ tensorInfoLegacyIfScalar(cuda::detail::TensorInfo<T, IndexType> ti) {
return ti;
}

constexpr uint64_t getDefaultMaxThreadsPerBlock() {
#ifndef USE_ROCM
return 128;
#else
// bigger default
return 512;
#endif
}

}

Expand Down