add kWarpSizeHost() for host-side launch configs (#6024)

Summary:
X-link: facebookresearch/FBGEMM#2934

Host-side launch configurations compute block/grid dims from `kWarpSize`. On ROCm, `kWarpSize` is a device-pass constant: during the host compilation pass it is a placeholder 64 (so `__global__` kernel bodies parse), which is wrong for sizing host-side `dim3` / grid quotients on a multi-arch wheel that runs on a wave32 arch (e.g. gfx1100), where the launch would use 64 lanes per warp instead of 32.

Add `kWarpSizeHost()` in `cuda_prelude.cuh`:
- CUDA: constexpr function returning 32 (usable as a template argument).
- ROCm: inline function returning `at::cuda::warp_size()`, a cached device-properties lookup of the active device's warp size.

Switch host-side `dim3` / quotient sites in the non-cache, non-TBE-codegen kernels (sparse, jagged, permute, quantize, input_combine, layout_transform, embedding_inplace, bounds_check) to `kWarpSizeHost()`. Device code keeps `kWarpSize`, which is per-arch correct via the device pass. Cache-geometry and TBE codegen sites are converted in later PRs in this chain.

Second in the chain splitting #5804 into reviewable pieces; stacked on #6014.

Authored with assistance from Claude (Anthropic).


Reviewed By: henrylhtsang

Differential Revision: D112414352

Pulled By: q10

Author: jeffdaily

fbgemm_gpu/codegen/utils/embedding_bounds_check_v1.cu

@@ -202,8 +202,10 @@ void _bounds_check_indices_cuda_v1(
                  : bounds_check_indices_kernel_v1<index_t, false>);
         FBGEMM_LAUNCH_DSA_KERNEL(
             bounds_check_kernel,
-            div_round_up(max_B_ * T, kNumThreads / fbgemm_gpu::kWarpSize),
-            dim3(fbgemm_gpu::kWarpSize, kNumThreads / fbgemm_gpu::kWarpSize),
+            div_round_up(max_B_ * T, kNumThreads / fbgemm_gpu::kWarpSizeHost()),
+            dim3(
+                fbgemm_gpu::kWarpSizeHost(),
+                kNumThreads / fbgemm_gpu::kWarpSizeHost()),
             0,
             at::cuda::getCurrentCUDAStream(),
             PTA_B(rows_per_table, int64_t, 1, 32),

fbgemm_gpu/codegen/utils/embedding_bounds_check_v2.cu

@@ -240,7 +240,8 @@ void _bounds_check_indices_cuda_v2(
 
   constexpr size_t kNumThreads = 1024;
   auto grid_dim = fbgemm_gpu::utils::cuda::cap_grid_dim_x(
-      cuda_calc_xblock_count(total_B, kNumThreads / fbgemm_gpu::kWarpSize),
+      cuda_calc_xblock_count(
+          total_B, kNumThreads / fbgemm_gpu::kWarpSizeHost()),
       kNumThreads,
       at::cuda::getCurrentCUDAStream(),
       fbgemm_gpu::utils::cuda::BlockCapPolicy::Always);
@@ -251,33 +252,34 @@ void _bounds_check_indices_cuda_v2(
     grid_dim = std::min<uint32_t>(grid_dim, PREFETCH_KERNEL_MAX_BLOCKS);
   }
 
-#define INVOKE_BOUNDS_CHECK_INDICES(MODE)                                      \
-  if (bounds_check_mode == MODE) {                                             \
-    AT_DISPATCH_INDEX_TYPES(                                                   \
-        indices.scalar_type(), "bounds_check_indices_cuda", [&] {              \
-          [[maybe_unused]] const auto func_name =                              \
-              "bounds_check_indices_cuda_v2";                                  \
-          const auto bounds_check_kernel =                                     \
-              (vbe ? bounds_check_indices_kernel_v2<index_t, true, MODE>       \
-                   : bounds_check_indices_kernel_v2<index_t, false, MODE>);    \
-          FBGEMM_LAUNCH_DSA_KERNEL(                                            \
-              bounds_check_kernel,                                             \
-              grid_dim,                                                        \
-              dim3(                                                            \
-                  fbgemm_gpu::kWarpSize, kNumThreads / fbgemm_gpu::kWarpSize), \
-              0,                                                               \
-              at::cuda::getCurrentCUDAStream(),                                \
-              PTA_B(rows_per_table, int64_t, 1, 32),                           \
-              PTA_B(indices, index_t, 1, 32),                                  \
-              PTA_B(offsets, index_t, 1, 32),                                  \
-              vbe ? B_offsets.value().data_ptr<int32_t>() : nullptr,           \
-              PTA_B(warning, int64_t, 1, 32),                                  \
-              FixedDivisor(B),                                                 \
-              vbe ? b_t_map.value().data_ptr<int32_t>() : nullptr,             \
-              info_B_num_bits,                                                 \
-              info_B_mask,                                                     \
-              disable_offsets_adjustment);                                     \
-        });                                                                    \
+#define INVOKE_BOUNDS_CHECK_INDICES(MODE)                                   \
+  if (bounds_check_mode == MODE) {                                          \
+    AT_DISPATCH_INDEX_TYPES(                                                \
+        indices.scalar_type(), "bounds_check_indices_cuda", [&] {           \
+          [[maybe_unused]] const auto func_name =                           \
+              "bounds_check_indices_cuda_v2";                               \
+          const auto bounds_check_kernel =                                  \
+              (vbe ? bounds_check_indices_kernel_v2<index_t, true, MODE>    \
+                   : bounds_check_indices_kernel_v2<index_t, false, MODE>); \
+          FBGEMM_LAUNCH_DSA_KERNEL(                                         \
+              bounds_check_kernel,                                          \
+              grid_dim,                                                     \
+              dim3(                                                         \
+                  fbgemm_gpu::kWarpSizeHost(),                              \
+                  kNumThreads / fbgemm_gpu::kWarpSizeHost()),               \
+              0,                                                            \
+              at::cuda::getCurrentCUDAStream(),                             \
+              PTA_B(rows_per_table, int64_t, 1, 32),                        \
+              PTA_B(indices, index_t, 1, 32),                               \
+              PTA_B(offsets, index_t, 1, 32),                               \
+              vbe ? B_offsets.value().data_ptr<int32_t>() : nullptr,        \
+              PTA_B(warning, int64_t, 1, 32),                               \
+              FixedDivisor(B),                                              \
+              vbe ? b_t_map.value().data_ptr<int32_t>() : nullptr,          \
+              info_B_num_bits,                                              \
+              info_B_mask,                                                  \
+              disable_offsets_adjustment);                                  \
+        });                                                                 \
   }
 
   INVOKE_BOUNDS_CHECK_INDICES(BoundsCheckMode::FATAL)

fbgemm_gpu/include/fbgemm_gpu/utils/cuda_prelude.cuh

@@ -60,16 +60,16 @@ namespace fbgemm_gpu {
 
 #define DIV_ROUND_UP(a, b) (a + b - 1) / b
 
-// Warp size
+// Warp size -- device-pass constant
 //
 // Device code: per-arch constexpr. __GFX9__ is defined by HIP-Clang during
 // device compilation for gfx9xx targets (warpSize 64); it is undefined for
 // warpSize 32 targets. HIP-Clang runs the device-code backend once per
 // --offload-arch, so the same source produces correct per-arch device code.
 //
-// Host code on ROCm: warpSize is only known at runtime. This 64 is a
-// stop-gap constexpr. Host-side launches that must size thread blocks for
-// the active device should call at::cuda::warp_size() directly.
+// Host code on ROCm: warpSize is only known at runtime. The 64 below is a
+// stop-gap so __global__ kernel bodies parse on the host pass; never rely
+// on it in host-side computation. Use kWarpSizeHost (below) instead.
 #if !defined(USE_ROCM)
 static constexpr int32_t kWarpSize = 32;
 #elif defined(__GFX9__)
@@ -80,6 +80,32 @@ static constexpr int32_t kWarpSize = 32;
 static constexpr int32_t kWarpSize = 64;
 #endif
 
+// Host-side warp size
+//
+// Use this in host code anywhere kWarpSize would be wrong on a ROCm
+// multi-arch build (block-dim computations, grid sizing, etc.). It is:
+//   * CUDA: a constexpr function returning 32. Usable as a template
+//     argument and in static_assert.
+//   * ROCm: an inline function returning at::cuda::warp_size() -- a runtime
+//     query of the active device. NOT constexpr; do not use as a template
+//     argument. Cheap (a cached device-properties array lookup).
+//
+// Always invoke with parentheses (kWarpSizeHost()) so the same call shape
+// works under both back-ends and inside namespace-qualified call sites.
+//
+// Do not use in __device__ code: on ROCm the at::cuda::warp_size() call
+// is not callable from device. Use kWarpSize there, which is per-arch
+// correct via the device pass.
+#if defined(USE_ROCM)
+inline int32_t kWarpSizeHost() {
+  return at::cuda::warp_size();
+}
+#else
+inline constexpr int32_t kWarpSizeHost() {
+  return 32;
+}
+#endif
+
 // Max thread num in one thread block
 static constexpr int32_t kMaxThreads = 1024;
 

fbgemm_gpu/src/embedding_inplace_ops/embedding_inplace_update.cu

@@ -240,7 +240,7 @@ void embedding_inplace_update_cuda(
   }
   TORCH_CHECK_EQ(N, update_table_idx.numel());
 
-  const int32_t warpsPerBlock = kMaxThreads / kWarpSize;
+  const int32_t warpsPerBlock = kMaxThreads / kWarpSizeHost();
 
   auto lxu_cache_weights_value = lxu_cache_weights.value_or(
       at::empty({0, 0}, dev_weights.options().dtype(at::kByte)));
@@ -261,7 +261,7 @@ void embedding_inplace_update_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (embedding_inplace_update_kernel_1<index_t>),
             blocks, // number of blocks needed
-            dim3(kWarpSize, warpsPerBlock), // shape of each block
+            dim3(kWarpSizeHost(), warpsPerBlock), // shape of each block
             0,
             at::cuda::getCurrentCUDAStream(),
 
@@ -316,7 +316,7 @@ void embedding_inplace_update_single_placement_cuda(
   }
   TORCH_CHECK_EQ(N, update_table_idx.numel());
 
-  const int32_t warpsPerBlock = kMaxThreads / kWarpSize;
+  const int32_t warpsPerBlock = kMaxThreads / kWarpSizeHost();
 
   auto lxu_cache_weights_value = lxu_cache_weights.value_or(
       at::empty({0, 0}, dev_weights.options().dtype(at::kByte)));
@@ -336,7 +336,7 @@ void embedding_inplace_update_single_placement_cuda(
         FBGEMM_LAUNCH_KERNEL(
             (embedding_inplace_update_kernel_2<index_t>),
             blocks, // number of blocks needed
-            dim3(kWarpSize, warpsPerBlock), // shape of each block
+            dim3(kWarpSizeHost(), warpsPerBlock), // shape of each block
             0,
             at::cuda::getCurrentCUDAStream(),
 

fbgemm_gpu/src/input_combine_ops/input_combine.cu

@@ -149,9 +149,9 @@ std::tuple<Tensor, Tensor, Tensor> tbe_input_combine_with_length_cuda(
 #else
   constexpr uint32_t VEC_WIDTH = 8;
 #endif
-  constexpr uint32_t NUM_WARPS_PER_BLOCK = kMaxThreads / kWarpSize;
+  const uint32_t NUM_WARPS_PER_BLOCK = kMaxThreads / kWarpSizeHost();
   const auto num_warps_per_list =
-      div_round_up(max_list_size, kWarpSize * VEC_WIDTH);
+      div_round_up(max_list_size, kWarpSizeHost() * VEC_WIDTH);
   const auto total_warp_work = num_warps_per_list * num_lists;
   // HIP enforces a hard limit of 2^32 total threads per launch.
   // tbe_input_combine_with_length_kernel grid-strides over warps, so capping
@@ -170,7 +170,7 @@ std::tuple<Tensor, Tensor, Tensor> tbe_input_combine_with_length_cuda(
   FBGEMM_LAUNCH_KERNEL(
       (tbe_input_combine_with_length_kernel<VEC_WIDTH, IS_LONG_NUM_BITS>),
       num_blocks,
-      dim3(kWarpSize, NUM_WARPS_PER_BLOCK),
+      dim3(kWarpSizeHost(), NUM_WARPS_PER_BLOCK),
       0,
       at::cuda::getCurrentCUDAStream(),
       combined_indices.data_ptr<int32_t>(),

fbgemm_gpu/src/jagged_tensor_ops/batched_dense_vec_jagged_2d_mul_backward.cu

@@ -114,7 +114,7 @@ std::tuple<Tensor, Tensor> batched_dense_vec_jagged_2d_mul_backward(
               "dense_vec_jagged_2d_bmm_backward_kernel_2",
               [&] {
                 int block_dim_x = std::min(
-                    div_round_up(max_L, kWarpSize) * kWarpSize, kMaxThreads);
+                    div_round_up(max_L, kWarpSizeHost()) * kWarpSizeHost(), kMaxThreads);
                 int block_dim_y = kMaxThreads / block_dim_x;
 
                 // HIP enforces a hard limit of 2^32 total threads per launch
@@ -141,7 +141,7 @@ std::tuple<Tensor, Tensor> batched_dense_vec_jagged_2d_mul_backward(
                     PTA_B(v_grad, scalar_t, 2, 32));
 
                 block_dim_x = std::min(
-                    div_round_up(D, kWarpSize) * kWarpSize, kMaxThreads);
+                    div_round_up(D, kWarpSizeHost()) * kWarpSizeHost(), kMaxThreads);
                 block_dim_y = kMaxThreads / block_dim_x;
 
                 // HIP 2^32 cap. outer_prod_jagged_2d_output grid-strides

fbgemm_gpu/src/jagged_tensor_ops/batched_dense_vec_jagged_2d_mul_forward.cu

@@ -71,7 +71,7 @@ Tensor batched_dense_vec_jagged_2d_mul_forward(
 
   if (B > 0 && D > 0) {
     const int block_dim_x =
-        std::min(div_round_up(D, kWarpSize) * kWarpSize, kMaxThreads);
+        std::min(div_round_up(D, kWarpSizeHost()) * kWarpSizeHost(), kMaxThreads);
     const int block_dim_y = kMaxThreads / block_dim_x;
 
     // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,

fbgemm_gpu/src/jagged_tensor_ops/common.cuh

@@ -222,9 +222,10 @@ inline std::tuple<dim3, dim3, StackArray<int64_t>> check_shape_and_partition_(
   const int jagged_folded_size =
       dense_tensor.numel() / (outer_dense_size * inner_dense_size);
 
-  const int threads_x =
-      inner_dense_size >= kWarpSize / 2 ? kWarpSize : inner_dense_size;
-  const int threads_y = kMaxThreads / kWarpSize;
+  const int threads_x = inner_dense_size >= kWarpSizeHost() / 2
+      ? kWarpSizeHost()
+      : inner_dense_size;
+  const int threads_y = kMaxThreads / kWarpSizeHost();
   const dim3 blocks(
       div_round_up(outer_dense_size * jagged_folded_size, threads_y));
 

fbgemm_gpu/src/layout_transform_ops/layout_transform_ops.cu

@@ -138,11 +138,12 @@ Tensor recat_embedding_grad_output_mixed_D_batch_cuda(
   const auto dim_sum = grad_output.size(1);
 
   const dim3 threads(
-      fbgemm_gpu::kWarpSize, fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSize);
+      fbgemm_gpu::kWarpSizeHost(),
+      fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSizeHost());
   const dim3 blocks(
       fbgemm_gpu::div_round_up(
           (B_local * dim_num),
-          fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSize));
+          fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSizeHost()));
 
   FBGEMM_DISPATCH_FLOAT_AND_HALF(
       grad_output.scalar_type(), "recat_embedding_gradients", [&] {

fbgemm_gpu/src/permute_multi_embedding_ops/permute_multi_embedding_ops.cu

@@ -272,9 +272,9 @@ std::vector<Tensor> permute_multi_embedding_function_gpu(
   // blocks. The grid z dimension is also used by batch_size in case it's
   // greater than 65535.
   const int32_t warp_per_block =
-      fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSize;
+      fbgemm_gpu::kMaxThreads / fbgemm_gpu::kWarpSizeHost();
   const int32_t max_grid_dim = 32768; // The CUDA maximum is 65535, not 1<<N.
-  const dim3 block_dim(fbgemm_gpu::kWarpSize, warp_per_block);
+  const dim3 block_dim(fbgemm_gpu::kWarpSizeHost(), warp_per_block);
   const dim3 grid_dim(
       fbgemm_gpu::div_round_up(permute_size, warp_per_block),
       std::min(static_cast<int32_t>(batch_size), max_grid_dim),