Add grid-stride loop and ROCm cap to batched_unary_embeddings_{forward,backward}_kernel (#5933)

Summary:
Pull Request resolved: #5933

Tier-2 fix for HIP grid-overflow in `sparse_ops/sparse_batched_unary_embeddings.cu`.

Both forward and backward kernels previously used `blockIdx.x * blockDim.x + threadIdx.x` directly with early-returns. Capping the host-side x-dim grid without first adding a grid-stride loop would silently drop work.

Changes:
- `batched_unary_embeddings_forward_kernel`: convert to a grid-stride loop over `b = blockIdx.x * blockDim.x + threadIdx.x; b < B; b += gridDim.x * blockDim.x` (Pattern A on the saturating x-dim). Hoisted `t = blockIdx.y`, `n = blockIdx.z`, `sum_E = table_offsets[T]`, and `table_offset = table_offsets[t]` outside the loop since y/z dims are bounded by T/N <= 65535 and don't need grid-striding.
- `batched_unary_embeddings_backward_kernel`: convert to a grid-stride loop over `run_id`. Hoisted `n = blockIdx.y` and the single-element device read of `sorted_linear_indices_num_runs[0]` outside the loop. The `min(num_runs_size_0, num_runs_actual)` becomes the loop bound. The `SL == 0` early-return is replaced with `continue`.
- Apply standard `#ifdef USE_ROCM min(blocks_x_uncapped, get_max_thread_blocks(stream)) #else blocks_x_uncapped #endif` cap to the x-dim of the launch grid for both kernels. y/z dims are bounded by T/N/N <= 65535 (enforced by host-side `TORCH_CHECK`s) and need no cap.

Stacked on top of D105028336 (Tier-2 Diff 4/7). Plan:
`/home/bensonma415/.llms/plans/sparse_ops_rocm_grid_overflow_tier2_fix.plan.md` (Diff 5/7).

Reviewed By: henrylhtsang

Differential Revision: D105029028

fbshipit-source-id: 5b8af5cd8395334be9bff66ea077aaec6568eefa

Author: q10

fbgemm_gpu/src/sparse_ops/sparse_batched_unary_embeddings.cu

@@ -26,23 +26,25 @@ __launch_bounds__(kMaxThreads) void batched_unary_embeddings_forward_kernel(
     const index_t* __restrict__ indices,
     scalar_t* __restrict__ output // N * B * T
 ) {
-  index_t sum_E = table_offsets[T];
-  auto b = blockIdx.x * blockDim.x + threadIdx.x;
-  if (b >= B) {
-    return;
-  }
-  auto t = blockIdx.y;
-  auto n = blockIdx.z;
-  index_t table_offset = table_offsets[t];
-  index_t indices_start = offsets[t * B + b];
-  index_t indices_end = offsets[t * B + b + 1];
-  int32_t L = indices_end - indices_start;
-  at::acc_type<scalar_t, true> sum = 0.0;
-  for (int32_t l = 0; l < L; ++l) {
-    auto idx = LDG(&indices[indices_start + l]);
-    sum += weight[n * sum_E + table_offset + idx + 0];
+  const index_t sum_E = table_offsets[T];
+  const auto t = blockIdx.y;
+  const auto n = blockIdx.z;
+  const index_t table_offset = table_offsets[t];
+  // Grid-stride over the b dimension so a capped grid (used on ROCm to avoid
+  // the 2^32 launch-side limit) still covers all B elements. blockIdx.y/z are
+  // already bounded by T/N <= 65535.
+  for (auto b = blockIdx.x * blockDim.x + threadIdx.x; b < B;
+       b += gridDim.x * blockDim.x) {
+    index_t indices_start = offsets[t * B + b];
+    index_t indices_end = offsets[t * B + b + 1];
+    int32_t L = indices_end - indices_start;
+    at::acc_type<scalar_t, true> sum = 0.0;
+    for (int32_t l = 0; l < L; ++l) {
+      auto idx = LDG(&indices[indices_start + l]);
+      sum += weight[n * sum_E + table_offset + idx + 0];
+    }
+    output[(n * B + b) * T + t] = sum;
   }
-  output[(n * B + b) * T + t] = sum;
 }
 
 Tensor batched_unary_embeddings_forward_cuda(
@@ -66,7 +68,14 @@ Tensor batched_unary_embeddings_forward_cuda(
   TORCH_CHECK(T <= 65535);
   TORCH_CHECK(N <= 65535);
   int32_t threads = std::min<int32_t>(B, 512);
-  dim3 blocks(cuda_calc_xblock_count(B, threads), T, N);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). The forward kernel grid-strides over b, so capping
+  // is correctness-preserving. y/z dims are bounded by T/N <= 65535 and need
+  // no cap.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_x = utils::cuda::cap_grid_dim_x_from_workload(
+      B, threads, at::cuda::getCurrentCUDAStream());
+  dim3 blocks(blocks_x, T, N);
   auto output = at::empty({N, B, T}, weight.options());
   AT_DISPATCH_INDEX_TYPES(
       indices.scalar_type(), "batched_unary_embeddings_forward_kernel", [&] {
@@ -127,45 +136,49 @@ __launch_bounds__(kMaxThreads) void batched_unary_embeddings_backward_kernel(
     const int32_t* __restrict__ sorted_linear_indices_num_runs,
     const int32_t info_B_num_bits,
     const uint32_t info_B_mask) {
-  auto run_id = blockIdx.x * blockDim.x + threadIdx.x;
-  auto n = blockIdx.y;
+  const auto n = blockIdx.y;
   if (n >= N) {
     return;
   }
-  if (run_id >= sorted_linear_indices_run.size(0)) {
-    return;
-  }
-  if (run_id >= sorted_linear_indices_num_runs[0]) {
-    return;
-  }
-  int64_t linear_index = sorted_linear_indices_run[run_id];
-  int32_t segment_start = sorted_linear_indices_cumulative_run_lengths[run_id];
-  int32_t segment_end =
-      sorted_linear_indices_cumulative_run_lengths[run_id + 1];
-  int32_t SL = segment_end - segment_start;
+  // Hoist single-element device tensor read out of the stride loop.
+  const auto num_runs_actual = sorted_linear_indices_num_runs[0];
+  const auto num_runs_size_0 = sorted_linear_indices_run.size(0);
+  const auto num_runs = ::min(num_runs_size_0, num_runs_actual);
+  // Grid-stride over run_id so a capped grid (used on ROCm to avoid the
+  // 2^32 launch-side limit) still covers all runs. blockIdx.y is bounded
+  // by N <= 65535 (enforced by the host).
+  for (auto run_id = blockIdx.x * blockDim.x + threadIdx.x; run_id < num_runs;
+       run_id += gridDim.x * blockDim.x) {
+    int64_t linear_index = sorted_linear_indices_run[run_id];
+    int32_t segment_start =
+        sorted_linear_indices_cumulative_run_lengths[run_id];
+    int32_t segment_end =
+        sorted_linear_indices_cumulative_run_lengths[run_id + 1];
+    int32_t SL = segment_end - segment_start;
 
-  if (SL == 0) {
-    return;
-  }
+    if (SL == 0) {
+      continue;
+    }
 
-  // now, each segment corresponds to exactly one table `t` and row in
-  // that table (`idx`). Thus, we can hoist out some of the book-keeping.
-  const auto info =
-      reinterpret_cast<const uint32_t*>(sorted_infos)[segment_start];
-  int t = info >> info_B_num_bits;
-
-  at::acc_type<scalar_t, true> grad_sum = 0.0;
-  for (int32_t sl = 0; sl < SL; ++sl) {
-    const auto b =
-        reinterpret_cast<const uint32_t*>(sorted_infos)[segment_start + sl] &
-        info_B_mask;
-    grad_sum += grad_output[(n * B + b) * T + t];
-  }
+    // now, each segment corresponds to exactly one table `t` and row in
+    // that table (`idx`). Thus, we can hoist out some of the book-keeping.
+    const auto info =
+        reinterpret_cast<const uint32_t*>(sorted_infos)[segment_start];
+    int t = info >> info_B_num_bits;
 
-  index_t table_offset = table_offsets[t];
-  index_t sum_E = table_offsets[T];
-  int64_t idx = linear_index - table_offset;
-  grad_weight[n * sum_E + table_offset + idx] = grad_sum;
+    at::acc_type<scalar_t, true> grad_sum = 0.0;
+    for (int32_t sl = 0; sl < SL; ++sl) {
+      const auto b =
+          reinterpret_cast<const uint32_t*>(sorted_infos)[segment_start + sl] &
+          info_B_mask;
+      grad_sum += grad_output[(n * B + b) * T + t];
+    }
+
+    index_t table_offset = table_offsets[t];
+    index_t sum_E = table_offsets[T];
+    int64_t idx = linear_index - table_offset;
+    grad_weight[n * sum_E + table_offset + idx] = grad_sum;
+  }
 }
 
 DLL_PUBLIC Tensor batched_unary_embeddings_backward_cuda(
@@ -219,8 +232,16 @@ DLL_PUBLIC Tensor batched_unary_embeddings_backward_cuda(
           info_B_mask);
 
   int threads = std::min<int32_t>(sorted_linear_indices_run.numel(), 512);
-  dim3 blocks(
-      cuda_calc_xblock_count(sorted_linear_indices_run.numel(), threads), N);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). The backward kernel grid-strides over run_id, so
+  // capping is correctness-preserving. y dim is bounded by N <= 65535 and
+  // needs no cap.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_x = utils::cuda::cap_grid_dim_x_from_workload(
+      sorted_linear_indices_run.numel(),
+      threads,
+      at::cuda::getCurrentCUDAStream());
+  dim3 blocks(blocks_x, N);
   auto grad_weight = at::zeros_like(weight);
   AT_DISPATCH_INDEX_TYPES(
       indices.scalar_type(), "batched_unary_embeddings_backward_kernel", [&] {

fbgemm_gpu/test/batched_unary_embeddings_test.py

@@ -7,6 +7,8 @@
 
 # pyre-strict
 
+# pyre-ignore-all-errors[56]
+
 
 import random
 import unittest
@@ -21,10 +23,10 @@
     from fbgemm_gpu import open_source  # noqa: F401
 
     # pyre-ignore[21]
-    from test_utils import gpu_unavailable
+    from test_utils import gpu_memory_lt_gb, gpu_unavailable
 
 except Exception:
-    from fbgemm_gpu.test.test_utils import gpu_unavailable
+    from fbgemm_gpu.test.test_utils import gpu_memory_lt_gb, gpu_unavailable
 
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
 
@@ -236,6 +238,108 @@ def test_gpu(self) -> None:
     def test_cpu(self) -> None:
         self._test_main(gpu_infer=False)
 
+    @unittest.skipIf(*gpu_unavailable)
+    # This test exercises the HIP launch-side limit and requires a large
+    # output tensor (~17 GiB) plus offsets (~4 GiB) — total ~22 GiB GPU
+    # peak. Most CI machines without 24+ GiB of HBM will skip this. ROCm
+    # developer machines (MI300/MI350, 192/256 GiB HBM) run it.
+    @unittest.skipIf(*gpu_memory_lt_gb(24))
+    def test_batched_unary_embeddings_forward_large_grid(self) -> None:
+        """
+        Reproduces the HIP grid-overflow bug in
+        batched_unary_embeddings_forward_kernel and verifies output
+        correctness via a downsampled CPU oracle.
+
+        With block size up to 512 and grid (cuda_calc_xblock_count(B, 512),
+        T, N), total threads ~= B * T * N. For B*T*N > 2**32, total threads
+        exceed the HIP 2**32 limit, causing FBGEMM_LAUNCH_KERNEL ->
+        KernelLauncher::checkThreadCountNotExceeded to TORCH_CHECK-fail on
+        ROCm pre-fix.
+
+        Verification strategy (per master plan's downsampled-oracle
+        guidance for ops where the full-scale CPU oracle is impractical):
+
+        1. Full-scale invocation to verify launch survival at the
+           cap-trip scale. Uses zero-length offsets so the kernel does no
+           per-segment work; output shape and zero-fill are asserted.
+        2. Small-scale invocation vs CPU dispatch to validate kernel
+           correctness end-to-end at a scale where the CPU oracle output
+           (~MB range) is cheap to compute and compare element-wise.
+        """
+
+        # The production cap is `blocks_x_capped = min(blocks_x_uncapped,
+        # get_max_thread_blocks(stream))` where `get_max_thread_blocks =
+        # 64 * #SMs ~= 16384` on MI300/MI350. For the cap to actually help,
+        # we need:
+        #   (a) blocks_x_uncapped > 16384 (so cap applies) -> B > 16384*512.
+        #   (b) blocks_x_uncapped * T * N * 512 > 2**32 (pre-fix trips).
+        #   (c) 16384 * T * N * 512 < 2**32 (post-fix passes), i.e. T*N < 512.
+        # Choose T=4, N=8 (T*N=32) and B = (1<<27)+1 = 2**27+1 so blocks_x =
+        # ceil(B/512) = 2**18 + 1 (>> 16384), pre-fix total ~= 2**32 + 2**14
+        # (trips), post-fix total = 16384*32*512 = 2**28 (well under 2**32).
+        B = (1 << 27) + 1
+        T = 4
+        N = 8
+
+        device = torch.device(torch.accelerator.current_accelerator() or "cuda")
+
+        # ---- Step 1: full-scale launch survival (cap-trip detection). ----
+        # weight[N, sum_E] where sum_E = T (each table has E=1). ~16 MB.
+        weight_large = torch.zeros((N, T), dtype=torch.float32, device=device)
+        table_offsets_large = torch.arange(T + 1, dtype=torch.int64, device=device)
+        # offsets[T*B+1] = all-zero so each (t,b) has L=0; no per-segment work.
+        offsets_large = torch.zeros(T * B + 1, dtype=torch.int64, device=device)
+        indices_large = torch.empty(0, dtype=torch.int64, device=device)
+
+        output_large = torch.ops.fbgemm.batched_unary_embeddings(
+            weight_large, table_offsets_large, offsets_large, indices_large
+        )
+
+        self.assertEqual(output_large.shape, (N, B, T))
+        self.assertTrue(torch.all(output_large == 0).item())
+        del output_large
+
+        # ---- Step 2: downsampled CPU-oracle correctness check. ----
+        # Same kernel code path, smaller scale to keep the CPU oracle cheap.
+        small_B = 4
+        small_T = 3
+
+        # Two embedding tables, each with E=2 rows. sum_E = 2*small_T = 6.
+        # N=2 (rows of weight). Distinct values per (n, row) so any "kernel
+        # addressed wrong row" bug surfaces.
+        small_weight_cpu = torch.tensor(
+            [
+                [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+                [10.0, 20.0, 30.0, 40.0, 50.0, 60.0],
+            ],
+            dtype=torch.float32,
+        )
+        # table_offsets[t]: start row of table t in weight columns.
+        small_table_offsets_cpu = torch.tensor([0, 2, 4, 6], dtype=torch.int64)
+        # offsets[t*B+b+1] - offsets[t*B+b] = length of (t,b) segment.
+        # Use 1 index per (t,b) cell. Total indices = T*B = 12.
+        small_offsets_cpu = torch.arange(small_T * small_B + 1, dtype=torch.int64)
+        # Each (t,b) picks row 0 of its table.
+        small_indices_cpu = torch.zeros(small_T * small_B, dtype=torch.int64)
+
+        # CPU reference oracle — same op, different dispatch.
+        small_output_cpu = torch.ops.fbgemm.batched_unary_embeddings(
+            small_weight_cpu,
+            small_table_offsets_cpu,
+            small_offsets_cpu,
+            small_indices_cpu,
+        )
+
+        # GPU under test.
+        small_output_gpu = torch.ops.fbgemm.batched_unary_embeddings(
+            small_weight_cpu.to(device),
+            small_table_offsets_cpu.to(device),
+            small_offsets_cpu.to(device),
+            small_indices_cpu.to(device),
+        )
+
+        torch.testing.assert_close(small_output_gpu.cpu(), small_output_cpu)
+
 
 if __name__ == "__main__":
     unittest.main()