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

Use a smaller grid size for TBE bwd cta_per_row #1353

Closed
wants to merge 1 commit into from
Closed

Use a smaller grid size for TBE bwd cta_per_row #1353

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
22 changes: 19 additions & 3 deletions fbgemm_gpu/codegen/embedding_backward_split_template.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,18 @@ constexpr size_t kBackwardMaxThreads = 512;
using Tensor = at::Tensor;
using namespace fbgemm_gpu;

namespace {

// Based on the empirical study, max grid size that is 64x larger than the
// number of SMs gives good performance across the board
constexpr int MAX_THREAD_BLOCKS_FACTOR = 64;

int get_max_thread_blocks_() {
return MAX_THREAD_BLOCKS_FACTOR * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
}

} // namespace

__global__ __launch_bounds__(kMaxThreads) void
split_embedding_backward_codegen_{{ optimizer }}_{{ wdesc }}_find_long_segments(
const at::PackedTensorAccessor32<int32_t, 1, at::RestrictPtrTraits>
Expand Down Expand Up @@ -1088,6 +1100,10 @@ split_embedding{{ "_nobag" if nobag else "" }}_backward_codegen_{{ optimizer }}_
{use_deterministic_algorithms ? 0 : grad_accum_counter.numel(), max_D},
grad_output.options().dtype(std::is_same<cache_t, double>::value ? at::kDouble : at::kFloat));

int32_t grid_size = std::min(
div_round_up(long_run_ids.numel(), kMaxThreads),
get_max_thread_blocks_());

// Check https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory-7-x
// "Compute capability 7.x devices allow a single thread block to
// address the full capacity of shared memory: 96 KB on Volta,
Expand Down Expand Up @@ -1127,7 +1143,7 @@ split_embedding{{ "_nobag" if nobag else "" }}_backward_codegen_{{ optimizer }}_
{% endif %}
kMaxVecsPerThread,
kThreadGroupSize>
<<<div_round_up(long_run_ids.numel(), kMaxThreads),
<<<grid_size,
dim3(kThreadGroupSize, BT_block_size),
BT_block_size * sizeof(at::acc_type<{{ "scalar_t" if dense else "cache_t" }}, true>) * 4 * kWarpSize *
kMaxVecsPerThread,
Expand Down Expand Up @@ -1182,9 +1198,9 @@ split_embedding{{ "_nobag" if nobag else "" }}_backward_codegen_{{ optimizer }}_
use_deterministic_algorithms,
{{ args.split_kernel_arg_constructors | join(", ") }});
C10_CUDA_KERNEL_LAUNCH_CHECK();
int32_t grid_size = std::min(
grid_size = std::min(
div_round_up(sorted_linear_indices_run.numel(), kBackwardMaxThreads / kThreadGroupSize),
64 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount);
get_max_thread_blocks_());
#ifndef __HIP_PLATFORM_HCC__
cudaFuncSetAttribute(
split_embedding{{ "_nobag" if nobag else "" }}_backward_codegen_{{ optimizer }}_{{ wdesc }}_kernel_warp_per_row_1<
Expand Down