diff --git a/fbgemm_gpu/src/split_embeddings_cache/linearize_cache_indices.cu b/fbgemm_gpu/src/split_embeddings_cache/linearize_cache_indices.cu index 04647ba147..70708a3004 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/linearize_cache_indices.cu +++ b/fbgemm_gpu/src/split_embeddings_cache/linearize_cache_indices.cu @@ -7,6 +7,7 @@ */ #include "common.cuh" +#include "fbgemm_gpu/utils/cuda_utilities.cuh" using Tensor = at::Tensor; using namespace fbgemm_gpu; @@ -24,36 +25,45 @@ __global__ __launch_bounds__(kMaxThreads) void linearize_cache_indices_kernel( pta::PackedTensorAccessor32 linear_cache_indices, const int64_t indices_base_offset) { +#ifdef USE_ROCM + for (index_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < indices.size(0); + index += blockDim.x * gridDim.x) { +#else const index_t index = blockIdx.x * blockDim.x + threadIdx.x; if (index >= indices.size(0)) { return; } +#endif + + // Perform binary search. + int left = 0; + int right = table_offsets.size(0); + const auto index_with_offset = index + indices_base_offset; + while (left != right) { + const int middle = + left + (right - left) / 2; // Avoid overflow in midpoint calculation + if (table_offsets[middle] <= index_with_offset) { + left = middle + 1; + } else { + right = middle; + } + } + const int table_index = left; - // Perform binary search. - int left = 0; - int right = table_offsets.size(0); - const auto index_with_offset = index + indices_base_offset; - while (left != right) { - const int middle = - left + (right - left) / 2; // Avoid overflow in midpoint calculation - if (table_offsets[middle] <= index_with_offset) { - left = middle + 1; + const auto max_offset = + ::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]); + const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]); + if (curr_offset >= 0 && indices[index] >= 0) { + linear_cache_indices[index] = indices[index] + curr_offset; } else { - right = middle; + // Either table index is wrong, or index value is negative (due to + // pruning): set it to invalid value. + linear_cache_indices[index] = max_offset; } - } - const int table_index = left; - - const auto max_offset = - ::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]); - const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]); - if (curr_offset >= 0 && indices[index] >= 0) { - linear_cache_indices[index] = indices[index] + curr_offset; - } else { - // Either table index is wrong, or index value is negative (due to pruning): - // set it to invalid value. - linear_cache_indices[index] = max_offset; - } +#ifdef USE_ROCM + } // for index (grid-stride loop, ROCm only) +#endif } } // namespace @@ -107,7 +117,10 @@ DLL_PUBLIC Tensor linearize_cache_indices_cuda( indices.scalar_type(), "linearize_cache_indices_kernel_2", [&] { FBGEMM_LAUNCH_KERNEL( (linearize_cache_indices_kernel), - div_round_up(num_indices, kMaxThreads), + utils::cuda::cap_grid_dim_x_from_workload( + num_indices, + kMaxThreads, + at::cuda::getCurrentCUDAStream()), kMaxThreads, 0, at::cuda::getCurrentCUDAStream(), @@ -134,22 +147,31 @@ __launch_bounds__(kMaxThreads) void linearize_cache_indices_from_row_idx_kernel( update_row_indices, pta::PackedTensorAccessor32 linear_cache_indices) { +#ifdef USE_ROCM + for (index_t index = blockIdx.x * blockDim.x + threadIdx.x; + index < update_row_indices.size(0); + index += blockDim.x * gridDim.x) { +#else const index_t index = blockIdx.x * blockDim.x + threadIdx.x; if (index >= update_row_indices.size(0)) { return; } - const int table_index = update_table_indices[index]; - - const auto max_offset = - ::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]); - const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]); - if (curr_offset >= 0 && update_row_indices[index] >= 0) { - linear_cache_indices[index] = update_row_indices[index] + curr_offset; - } else { - // Either table index is wrong, or index value is negative (due to pruning): - // set it to invalid value. - linear_cache_indices[index] = max_offset; - } +#endif + const int table_index = update_table_indices[index]; + + const auto max_offset = + ::__ldg(&cache_hash_size_cumsum[cache_hash_size_cumsum.size(0) - 1]); + const auto curr_offset = ::__ldg(&cache_hash_size_cumsum[table_index]); + if (curr_offset >= 0 && update_row_indices[index] >= 0) { + linear_cache_indices[index] = update_row_indices[index] + curr_offset; + } else { + // Either table index is wrong, or index value is negative (due to + // pruning): set it to invalid value. + linear_cache_indices[index] = max_offset; + } +#ifdef USE_ROCM + } // for index (grid-stride loop, ROCm only) +#endif } } // namespace @@ -178,7 +200,8 @@ DLL_PUBLIC Tensor linearize_cache_indices_from_row_idx_cuda( [&] { FBGEMM_LAUNCH_KERNEL( (linearize_cache_indices_from_row_idx_kernel), - div_round_up(num_indices, kMaxThreads), + utils::cuda::cap_grid_dim_x_from_workload( + num_indices, kMaxThreads, at::cuda::getCurrentCUDAStream()), kMaxThreads, 0, at::cuda::getCurrentCUDAStream(), diff --git a/fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu b/fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu index 121be6d553..8633fca6d9 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu +++ b/fbgemm_gpu/src/split_embeddings_cache/lxu_cache.cu @@ -272,52 +272,66 @@ __global__ __launch_bounds__(kMaxThreads) void lxu_cache_lookup_kernel( const int32_t C = lxu_cache_state.size(0); const int32_t N = N_unique == nullptr ? linear_cache_indices.size(0) : *N_unique; + // On ROCm the launch caps the grid to stay within the HIP 2^32 + // threads-per-launch limit, so we grid-stride to cover the full workload. + // On CUDA the grid is not capped and the loop body runs once per warp. + const auto n_stride = gridDim.x * blockDim.y * blockDim.x; +#ifdef USE_ROCM + for (auto n0 = + blockIdx.x * blockDim.y * blockDim.x + threadIdx.y * blockDim.x; + n0 < N; + n0 += n_stride) { +#else const auto n0 = blockIdx.x * blockDim.y * blockDim.x + threadIdx.y * blockDim.x; if (n0 >= N) { return; } +#endif - int32_t cache_location = kCacheLocationMissing; - int32_t n_indices = 0; - int32_t n_hits = 0; - const auto slot = threadIdx.x; - for (int i = 0; i < blockDim.x; ++i) { - const auto n = n0 + i; - if (n >= N) { - continue; - } - const int64_t idx = linear_cache_indices[n]; - if (idx == invalid_index) { - continue; - } - const int32_t cache_set = cache_slot(idx, C); - n_indices++; - const bool found = - (::__ldg((&lxu_cache_state[cache_set][0]) + slot) == idx); - // fbgemm_gpu::ballot_sync wraps __ballot (ROCm, unsynchronized) vs - // __ballot_sync (CUDA). Return type is uint64_t on ROCm and uint32_t - // on CUDA, so __ffsll handles both via implicit promotion. - const auto bitmap = fbgemm_gpu::ballot_sync(found); - if (bitmap) { - // LSB == 1 hence we need to subtract one to get lane ID. - const auto way = __ffsll(static_cast(bitmap)) - 1; - if (i == threadIdx.x) { - cache_location = cache_set * kWarpSize + way; + int32_t cache_location = kCacheLocationMissing; + int32_t n_indices = 0; + int32_t n_hits = 0; + const auto slot = threadIdx.x; + for (int i = 0; i < blockDim.x; ++i) { + const auto n = n0 + i; + if (n >= N) { + continue; + } + const int64_t idx = linear_cache_indices[n]; + if (idx == invalid_index) { + continue; + } + const int32_t cache_set = cache_slot(idx, C); + n_indices++; + const bool found = + (::__ldg((&lxu_cache_state[cache_set][0]) + slot) == idx); + // fbgemm_gpu::ballot_sync wraps __ballot (ROCm, unsynchronized) vs + // __ballot_sync (CUDA). Return type is uint64_t on ROCm and uint32_t + // on CUDA, so __ffsll handles both via implicit promotion. + const auto bitmap = fbgemm_gpu::ballot_sync(found); + if (bitmap) { + // LSB == 1 hence we need to subtract one to get lane ID. + const auto way = __ffsll(static_cast(bitmap)) - 1; + if (i == threadIdx.x) { + cache_location = cache_set * kWarpSize + way; + } + n_hits++; } - n_hits++; } - } - const auto n = n0 + threadIdx.x; - if (n < N) { - lxu_cache_locations[n] = cache_location; - } - if (gather_cache_stats && threadIdx.x == 0 && n_indices > n_hits) { - atomicAdd( - &uvm_cache_stats[uvm_cache_stats_index::num_conflict_misses], - (n_indices - n_hits)); - } + const auto n = n0 + threadIdx.x; + if (n < N) { + lxu_cache_locations[n] = cache_location; + } + if (gather_cache_stats && threadIdx.x == 0 && n_indices > n_hits) { + atomicAdd( + &uvm_cache_stats[uvm_cache_stats_index::num_conflict_misses], + (n_indices - n_hits)); + } +#ifdef USE_ROCM + } // for n0 (grid-stride loop, ROCm only) +#endif } template @@ -447,7 +461,12 @@ DLL_PUBLIC Tensor lxu_cache_lookup_cuda( } const dim3 threads(kWarpSizeHost(), kMaxThreads / kWarpSizeHost()); - const dim3 blocks(div_round_up(N, kMaxThreads)); + // HIP enforces a hard limit of 2^32 total threads per launch; + // lxu_cache_lookup_kernel grid-strides over n0, so capping is + // correctness-preserving. See: https://github.com/ROCm/hip/issues/2253 + const dim3 blocks( + utils::cuda::cap_grid_dim_x_from_workload( + N, kMaxThreads, at::cuda::getCurrentCUDAStream())); AT_DISPATCH_INDEX_TYPES( linear_cache_indices.scalar_type(), "lxu_cache_lookup_cuda", [&] {