Add grid-stride + ROCm cap to keyed_jagged_index_select_dim1_kernel and keyed_jagged_index_add_dim1_kernel; fast-fail TORCH_CHECK for index_select_scalar_cumsum_kernel (#6025)

Summary:
Pull Request resolved: #6025

X-link: https://github.com/facebookresearch/FBGEMM/pull/2930

Diff 10/10 (final) of the Tier-2 ROCm grid-overflow fix stack.

Three kernels in keyed_jagged_index_select_dim1.cu need attention:

- keyed_jagged_index_select_dim1_kernel (forward) and
  keyed_jagged_index_add_dim1_kernel (backward) both used flat
  if (tid < N) predication. This diff wraps each in an explicit int64_t
  grid-stride for-loop and adds the standard #ifdef USE_ROCM host-side cap
  on grid_size at both launch sites.

- index_select_scalar_cumsum_kernel is a cross-block prefix-sum with
  block_flags / block_sums coordination. Capping the grid would corrupt the
  cumsum. Instead, this diff adds a fast-fail TORCH_CHECK on the host side
  before the launch, requiring num_output_lengths < 2^32. The full
  algorithmic restructure (per-segment block tiling or two-pass scan via
  cub::DeviceScan) is tracked as a Tier-3 follow-up.

Reviewed By: henrylhtsang

Differential Revision: D105205634

fbshipit-source-id: 852630d2c4d3e9ad5e3843b5031443ca99723943

Author: q10

fbgemm_gpu/src/jagged_tensor_ops/keyed_jagged_index_select_dim1.cu

@@ -169,11 +169,15 @@ __global__ void keyed_jagged_index_select_dim1_kernel(
         output_offsets,
     const int num_batches,
     const int input_batch_size) {
-  const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
   const int output_batch_size = indices.size(0);
   const int64_t num_outputs = output.size(0);
 
-  if (tid < num_outputs) {
+  // Grid-stride over outputs so a capped grid (used on ROCm to avoid the
+  // 2^32 launch-side limit) still covers every output element.
+  for (int64_t tid =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       tid < num_outputs;
+       tid += static_cast<int64_t>(gridDim.x) * blockDim.x) {
     entry_t index_pos;
     const entry_t num_entries =
         static_cast<entry_t>(num_batches) * output_batch_size;
@@ -220,11 +224,15 @@ __global__ void keyed_jagged_index_add_dim1_kernel(
         output_offsets,
     const int num_batches,
     const int output_batch_size) {
-  const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
   const int input_batch_size = indices.size(0);
   const int64_t num_inputs = input.size(0);
 
-  if (tid < num_inputs) {
+  // Grid-stride over inputs so a capped grid (used on ROCm to avoid the
+  // 2^32 launch-side limit) still covers every input element.
+  for (int64_t tid =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       tid < num_inputs;
+       tid += static_cast<int64_t>(gridDim.x) * blockDim.x) {
     entry_t index_pos;
     const entry_t num_entries =
         static_cast<entry_t>(num_batches) * input_batch_size;
@@ -278,6 +286,20 @@ class KeyedJaggedIndexSelectDim1GPUOp
         " * indices.numel()=",
         indices.numel(),
         " causes overflow.");
+    // Special: index_select_scalar_cumsum_kernel uses a cross-block
+    // decoupled-look-back scan with block_flags / block_sums; capping the
+    // grid would corrupt the cumsum. As a Tier-2 stop-gap, fast-fail with
+    // TORCH_CHECK if the launch would trip the HIP 2^32 limit. Algorithmic
+    // restructure (cub::DeviceScan or per-segment block tiling) is tracked
+    // as a Tier-3 follow-up.
+    TORCH_CHECK(
+        static_cast<uint64_t>(num_output_lengths) <
+            (static_cast<uint64_t>(1) << 32),
+        "index_select_scalar_cumsum_kernel: num_output_lengths (",
+        num_output_lengths,
+        ") must be < 2^32 to avoid HIP grid-overflow. The kernel uses a "
+        "cross-block prefix scan that cannot be naively grid-strided. See "
+        "Tier-3 follow-up for restructure.");
     auto grid_size = cuda_calc_xblock_count(
         num_output_lengths, MAX_CUMSUM_ENTRIES_PER_BLOCK);
     TORCH_CHECK_VALUE(
@@ -453,7 +475,12 @@ class KeyedJaggedIndexSelectDim1GPUOp
     if (weights.has_value()) {
       output_weights = at::empty({num_outputs}, weights.value().options());
     }
-    grid_size = cuda_calc_xblock_count(num_outputs, kMaxThreads);
+    // HIP enforces a hard limit of 2^32 total threads per launch.
+    // keyed_jagged_index_select_dim1_kernel grid-strides over outputs, so
+    // capping is correctness-preserving.
+    // See: https://github.com/ROCm/hip/issues/2253
+    grid_size = utils::cuda::cap_grid_dim_x_from_workload(
+        num_outputs, kMaxThreads, at::cuda::getCurrentCUDAStream());
 
     // output_offsets has to be contiguous because it is passed to
     // binary_search_range which takes raw pointers as arguments
@@ -633,7 +660,12 @@ class KeyedJaggedIndexSelectDim1GPUOp
     device_guard.set_index(grad.get_device());
 
     Tensor grad_input = at::zeros({num_outputs}, grad.options());
-    auto grid_size = cuda_calc_xblock_count(grad.numel(), kMaxThreads);
+    // HIP enforces a hard limit of 2^32 total threads per launch.
+    // keyed_jagged_index_add_dim1_kernel grid-strides over inputs, so
+    // capping is correctness-preserving.
+    // See: https://github.com/ROCm/hip/issues/2253
+    auto grid_size = utils::cuda::cap_grid_dim_x_from_workload(
+        grad.numel(), kMaxThreads, at::cuda::getCurrentCUDAStream());
     // grad_offsets has to be contiguous because it is passed to
     // binary_search_range which takes raw pointers as arguments
     const auto grad_offsets_contig = grad_offsets.expect_contiguous();

fbgemm_gpu/test/jagged/keyed_jagged_index_select_test.py

@@ -347,6 +347,90 @@ def test_keyed_jagged_index_select_dim1_int32_overflow(
             atol=1e-2 if jagged_tensor_dtype in [torch.half, torch.bfloat16] else None,
         )
 
+    @optests.dontGenerateOpCheckTests(
+        "large-grid GPU-memory-gated stress repro; opcheck variants add no op "
+        "coverage and only skip on CPU samples (T191384137)"
+    )
+    @unittest.skipUnless(torch.cuda.is_available(), "GPU not available")
+    def test_keyed_jagged_index_select_dim1_large_grid(self) -> None:
+        """
+        Regression test for the HIP grid-overflow bug in
+        ``keyed_jagged_index_select_dim1_kernel`` and
+        ``keyed_jagged_index_add_dim1_kernel`` (D105205634 / Subplan D
+        Diff #30).
+
+        Both kernels grid-stride over outputs/inputs respectively. The
+        production fix caps grid_x to ``get_max_thread_blocks(stream)``
+        on ROCm; on CUDA the grid still covers every element in a single
+        stride. Either way the refactor must remain correctness-
+        preserving. ``keyed_jagged_index_select_dim1`` is a CUDA-only op
+        (no CPU kernel), so verification is end-to-end forward + backward
+        correctness on the accelerator against a per-key
+        ``jagged_index_select`` reference, using sentinel non-zero
+        lengths at a multi-key scale.
+
+        The third kernel touched by this diff —
+        ``index_select_scalar_cumsum_kernel`` — uses a cross-block
+        decoupled-look-back scan and was given a Tier-C fast-fail
+        TORCH_CHECK guard rather than a grid-stride retrofit. That
+        Tier-C branch is exercised by code review of the diff itself
+        (the production TORCH_CHECK line) and not unit-testable here
+        without allocating a ~16 GiB indices tensor.
+        """
+        device = torch.accelerator.current_accelerator()
+        num_batches = 2
+        input_batch_size = 32
+        # Sparse sentinel non-zero lengths spread across batches.
+        lengths = torch.zeros(
+            num_batches * input_batch_size, dtype=torch.int64, device=device
+        )
+        lengths[0] = 3
+        lengths[input_batch_size + input_batch_size // 2] = 2
+        lengths[-1] = 4
+        offsets = torch.concat(
+            [torch.zeros(1, dtype=torch.long, device=device), lengths.cumsum(0)]
+        )
+        total = int(offsets[-1].item())
+
+        # Distinct values across rows; indices select all batches.
+        values_init = torch.arange(total, dtype=torch.float32, device=device)
+        indices = torch.arange(input_batch_size, dtype=torch.int64, device=device)
+
+        # Op under test: forward + backward on the accelerator.
+        values = values_init.detach().clone().requires_grad_(True)
+        out = torch.ops.fbgemm.keyed_jagged_index_select_dim1(
+            values,
+            lengths,
+            offsets,
+            indices,
+            input_batch_size,
+        )[0]
+
+        # Reference: keyed_jagged_index_select_dim1 has no CPU kernel, so build
+        # the expected result on-device from the per-key jagged_index_select
+        # primitive (same equivalence the other tests in this file rely on).
+        values_ref = values_init.detach().clone().requires_grad_(True)
+        output_ref = []
+        for k in range(num_batches):
+            key_lengths = lengths[k * input_batch_size : (k + 1) * input_batch_size]
+            start_offset = offsets[k * input_batch_size]
+            end_offset = offsets[(k + 1) * input_batch_size]
+            key_values = values_ref[start_offset:end_offset].view(-1, 1)
+            output_ref.append(
+                torch.ops.fbgemm.jagged_index_select(key_values, key_lengths, indices)[
+                    0
+                ].view(-1)
+            )
+        output_ref = torch.concat(output_ref)
+
+        torch.testing.assert_close(out, output_ref)
+
+        out.sum().backward()
+        output_ref.sum().backward()
+        assert values.grad is not None
+        assert values_ref.grad is not None
+        torch.testing.assert_close(values.grad, values_ref.grad)
+
 
 if __name__ == "__main__":
     unittest.main()