Fix HIP grid overflow in expand_into_jagged_permute_cuda (#5906)

Summary:
Pull Request resolved: #5906

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

Apply the same `#ifdef USE_ROCM` cap pattern used in D104903707 /
D104937969 / parent diffs to the launch site in
`expand_into_jagged_permute_cuda` in `sparse_expand_into_jagged_permute.cu`.

The launch uses `dim3(kWarpSize=32, T_blocks=32)` (block size 1024) and
grid `cuda_calc_xblock_count(permute_size, 32)`. Once
`permute_size > 2^27 ≈ 134M`, total threads exceed the HIP `2^32`
limit and the launch is rejected on ROCm.
`expand_into_jagged_permute_kernel` already grid-strides over `t`
(line 27), so capping the grid is correctness-preserving.

The `#ifdef USE_ROCM / #else / #endif` selector keeps NVIDIA codegen
bit-identical and unconditionally caps on ROCm.

Same family of fix as:
- D104903707 (permute_1D_sparse_data — parent diff)
- D104937969 (permute_2D_sparse_data — parent diff)

Reviewed By: henrylhtsang

Differential Revision: D104951014

fbshipit-source-id: c8cd6ec332b8e0284d72e6da7f869def132a3583

Author: q10

fbgemm_gpu/src/sparse_ops/sparse_expand_into_jagged_permute.cu

@@ -55,7 +55,14 @@ DLL_PUBLIC Tensor expand_into_jagged_permute_cuda(
   // number of table per block
   constexpr int32_t T_blocks = kMaxThreads / kWarpSize;
   dim3 threads(kWarpSize, T_blocks);
-  const auto blocks = cuda_calc_xblock_count(permute_size, T_blocks);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). expand_into_jagged_permute_kernel grid-strides
+  // over t (line 27), so capping is correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks = utils::cuda::cap_grid_dim_x(
+      cuda_calc_xblock_count(permute_size, T_blocks),
+      kMaxThreads,
+      at::cuda::getCurrentCUDAStream());
   AT_DISPATCH_INDEX_TYPES(
       permute.scalar_type(), "expand_into_jagged_permute_kernel", [&] {
         using offsets_t = index_t;

fbgemm_gpu/test/jagged/expand_into_jagged_permute_test.py

@@ -21,9 +21,21 @@
 
 if open_source:
     # pyre-ignore[21]
-    from test_utils import gpu_available, on_oss_clang, optests
+    from test_utils import (
+        gpu_available,
+        gpu_memory_lt_gb,
+        gpu_unavailable,
+        on_oss_clang,
+        optests,
+    )
 else:
-    from fbgemm_gpu.test.test_utils import gpu_available, on_oss_clang, optests
+    from fbgemm_gpu.test.test_utils import (
+        gpu_available,
+        gpu_memory_lt_gb,
+        gpu_unavailable,
+        on_oss_clang,
+        optests,
+    )
 
 
 @optests.generate_opcheck_tests(additional_decorators=additional_decorators)
@@ -99,6 +111,81 @@ def test_expand_into_jagged_permute(
                 output_permute_gpu.cpu(), output_permute_ref_tensor
             )
 
+    @unittest.skipIf(*gpu_unavailable)
+    # Skip on GPUs with insufficient HBM. Inputs are int32 of size
+    # ~permute_size, totaling ~1.5 GiB at the chosen permute_size.
+    @unittest.skipIf(*gpu_memory_lt_gb(4))
+    def test_expand_into_jagged_permute_large_grid(self) -> None:
+        """
+        Reproduces the HIP grid-overflow bug in expand_into_jagged_permute_cuda
+        and verifies output correctness at the same scale.
+
+        With T_blocks = kMaxThreads/kWarpSize = 32 and dim3(32, 32)
+        (block size 1024), the launch grid is
+        cuda_calc_xblock_count(permute_size, 32). For
+        permute_size > 2**27, total threads exceed the HIP 2**32
+        limit, causing FBGEMM_LAUNCH_KERNEL ->
+        KernelLauncher::checkThreadCountNotExceeded to TORCH_CHECK-fail
+        on ROCm pre-fix. With the production fix in place, this test
+        additionally validates output correctness against the CPU
+        dispatch of the same op — the GPU output must match the CPU
+        reference element-for-element.
+
+        ``length`` (per segment) is sparse: all zero except for three
+        known non-zero positions (start / middle / end of the logical
+        range), so HBM usage stays bounded (~1.6 GiB int32) while the
+        permutation expansion logic is still exercised. ``permute`` is
+        a deterministic non-identity circular shift (``perm[i] != i``
+        everywhere), so any "kernel computed identity instead of
+        permutation" bug surfaces in the assertion below.
+        """
+
+        # Choose permute_size so that total threads strictly exceeds 2**32:
+        # cuda_calc_xblock_count(permute_size, 32) * 1024 ~= permute_size * 32;
+        # need permute_size > 2**27.
+        permute_size = (1 << 27) + 1
+
+        device = torch.device(torch.accelerator.current_accelerator() or "cuda")
+
+        # Sparse non-zero lengths at start / middle / end. Total = 10.
+        lengths_cpu = torch.zeros(permute_size, dtype=torch.int32)
+        lengths_cpu[0] = 3
+        lengths_cpu[permute_size // 2] = 5
+        lengths_cpu[permute_size - 1] = 2
+
+        # Deterministic non-identity permute: circular shift by +1.
+        # perm_cpu[0] == permute_size - 1 and perm_cpu[i] == i - 1 for
+        # i >= 1, so perm_cpu[i] != i for every i.
+        perm_cpu = torch.roll(torch.arange(permute_size, dtype=torch.int32), 1)
+
+        # Build offsets (size permute_size + 1) on CPU.
+        input_offsets_cpu = torch.zeros(permute_size + 1, dtype=torch.int32)
+        input_offsets_cpu[1:] = torch.cumsum(lengths_cpu, dim=0).to(torch.int32)
+
+        permuted_lengths_cpu = lengths_cpu[perm_cpu.long()]
+        output_offsets_cpu = torch.zeros(permute_size + 1, dtype=torch.int32)
+        output_offsets_cpu[1:] = torch.cumsum(permuted_lengths_cpu, dim=0).to(
+            torch.int32
+        )
+
+        total_length = int(input_offsets_cpu[-1].item())
+
+        # CPU reference oracle — same op, different dispatch.
+        output_permute_cpu = torch.ops.fbgemm.expand_into_jagged_permute(
+            perm_cpu, input_offsets_cpu, output_offsets_cpu, total_length
+        )
+
+        # GPU op under test. Pre-fix, this launch trips
+        # KernelLauncher::checkThreadCountNotExceeded on ROCm.
+        output_permute_gpu = torch.ops.fbgemm.expand_into_jagged_permute(
+            perm_cpu.to(device),
+            input_offsets_cpu.to(device),
+            output_offsets_cpu.to(device),
+            total_length,
+        )
+
+        torch.testing.assert_close(output_permute_gpu.cpu(), output_permute_cpu)
+
 
 if __name__ == "__main__":
     unittest.main()