Fix the grid problem with sparse permute 2d (V2: threshold-guarded) (#5898)

Summary:
Pull Request resolved: #5898

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

Migrates the four host-side cap sites in `sparse_permute_2d.cu` from
the unconditional ROCm cap pattern (V1) to the new
`fbgemm_gpu::utils::cuda::determine_grid_blocks` helper (introduced in
D106267802) with the default `BlockCapPolicy::OverflowOnly`.

Net behaviour change vs V1:
- ROCm small/medium grid (`blocks * threads_per_block <= UINT32_MAX`):
  uncapped grid restored. Was: capped to
  `MAX_THREAD_BLOCKS_FACTOR * #SMs` unconditionally. Now: passes
  through unchanged (matches pre-D104903707 behaviour).
- ROCm large grid (`blocks * threads_per_block > UINT32_MAX`):
  cap still applied. The kernel grid-strides over `b_t` so capping is
  correctness-preserving, and this is the regime where the HIP
  `2^32` thread-per-launch limit would otherwise fire.
- NVIDIA: bit-identical to V1 (the threshold check lives entirely
  under `#ifdef USE_ROCM` inside the helper).

The four sites:
- `permute_2D_sparse_data_cuda` `blocks_1` (line ~223), `threads_1=256`.
- `permute_2D_sparse_data_cuda` `blocks_2` (line ~262), block size
  `dim3(32, BT_blocks=32)` so threads_per_block = `32 * BT_blocks = 1024`.
- `permute_sparse_features_cuda` `blocks_1` (line ~439), `threads_1=256`.
- `permute_sparse_features_cuda` `blocks_2` (line ~486), block size
  same as above (1024).

Drive-by: switches `cuda_calc_block_count` (y/z-dim 65535 cap) to
`determine_grid_blocks` (which uses `cuda_calc_xblock_count` internally
with the 2^31-1 x-dim cap). These launches are 1-D x-dim launches so
`cuda_calc_xblock_count` is the correct primitive; the kernels
grid-stride so any change in grid size is correctness-preserving.

Reviewed By: spcyppt

Differential Revision: D104937969

fbshipit-source-id: e7c89c0942c8665595772b58ca72c0bf88177a0b

Author: q10

fbgemm_gpu/src/sparse_ops/sparse_permute_2d.cu

@@ -221,7 +221,12 @@ permute_2D_sparse_data_cuda(
   permuted_lengths = at::empty({T, B}, lengths.options());
 
   constexpr int32_t threads_1 = 256;
-  const auto blocks_1 = cuda_calc_block_count(B * T, threads_1);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). permute_2D_lengths_kernel uses CUDA_KERNEL_LOOP,
+  // which already grid-strides, so capping is correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_1 = utils::cuda::cap_grid_dim_x_from_workload(
+      B * T, threads_1, at::cuda::getCurrentCUDAStream());
   AT_DISPATCH_INDEX_TYPES(
       lengths.scalar_type(), "permute_2D_lengths_kernel", [&] {
         FBGEMM_LAUNCH_KERNEL(
@@ -250,7 +255,15 @@ permute_2D_sparse_data_cuda(
 
   constexpr int32_t BT_blocks = 32;
   dim3 threads_2(32, BT_blocks);
-  const auto blocks_2 = cuda_calc_block_count(B * T, BT_blocks);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). Both permute_2D_data_kernel and
+  // permute_2D_data_kernel_vec grid-stride over b_t, so capping is
+  // correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_2 = utils::cuda::cap_grid_dim_x(
+      cuda_calc_xblock_count(B * T, BT_blocks),
+      BT_blocks * 32,
+      at::cuda::getCurrentCUDAStream());
   permuted_indices = at::empty(permuted_indices_size, indices.options());
 
   AT_DISPATCH_INDEX_TYPES(
@@ -415,8 +428,12 @@ permute_sparse_features_cuda(
   permuted_lengths = at::empty({num_output_features, B}, lengths.options());
 
   constexpr int32_t threads_1 = 256;
-  const auto blocks_1 =
-      cuda_calc_block_count(B * num_output_features, threads_1);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). permute_2D_lengths_kernel uses CUDA_KERNEL_LOOP,
+  // which already grid-strides, so capping is correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_1 = utils::cuda::cap_grid_dim_x_from_workload(
+      B * num_output_features, threads_1, at::cuda::getCurrentCUDAStream());
   AT_DISPATCH_INDEX_TYPES(
       lengths.scalar_type(), "permute_2D_lengths_kernel", [&] {
         FBGEMM_LAUNCH_KERNEL(
@@ -452,8 +469,14 @@ permute_sparse_features_cuda(
 
   constexpr int32_t BT_blocks = 32;
   dim3 threads_2(32, BT_blocks);
-  const auto blocks_2 =
-      cuda_calc_block_count(B * num_output_features, BT_blocks);
+  // HIP enforces a hard limit of 2^32 total threads per launch (unlike CUDA,
+  // which silently wraps). permute_indices_weights_kernel grid-strides over
+  // b_t, so capping is correctness-preserving.
+  // See: https://github.com/ROCm/hip/issues/2253
+  const auto blocks_2 = utils::cuda::cap_grid_dim_x(
+      cuda_calc_xblock_count(B * num_output_features, BT_blocks),
+      BT_blocks * 32,
+      at::cuda::getCurrentCUDAStream());
   permuted_indices = at::empty(permuted_lengths_sum, indices.options());
   if (weights.has_value()) {
     const Tensor weights_value = weights.value();

fbgemm_gpu/test/sparse/permute_indices_test.py

@@ -866,6 +866,156 @@ def test_permute_1D_sparse_data_large_grid(self) -> None:
         torch.testing.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
         self.assertIsNone(permuted_weights_gpu)
 
+    @unittest.skipIf(*gpu_unavailable)
+    # Skip on GPUs with insufficient HBM (need ~512 MB for the int32
+    # lengths tensor at the chosen B).
+    @unittest.skipIf(*gpu_memory_lt_gb(4))
+    def test_permute_2D_sparse_data_large_grid(self) -> None:
+        """
+        Reproduces the HIP grid-overflow bug in permute_2D_sparse_data_cuda
+        and verifies output correctness at the same scale.
+
+        With BT_blocks=32 and dim3(32, 32) (block size 1024), the
+        permute_2D_data_kernel_vec launch grid is
+        cuda_calc_block_count(B*T, 32). For B*T > 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.
+
+        Uses ``T=2, B=2**26+1`` so ``B*T = 2**27 + 2`` strictly trips the
+        threshold. ``lengths`` is sparse: all zero except for four known
+        non-zero positions (one per row, plus one mid-row), so HBM usage
+        stays bounded (~537 MB int32) while the permutation logic is
+        still exercised. ``permute = [1, 0]`` is a deterministic row swap
+        on the T axis with ``perm[i] != i`` for every i, so any
+        "kernel computed identity instead of permutation" or wrong-``b_t``
+        decoding bug surfaces in the assertion below.
+        """
+
+        # Choose B*T so that total threads strictly exceeds 2**32:
+        # cuda_calc_block_count(B*T, 32) * 1024 ~= B*T * 32; need B*T > 2**27.
+        T = 2
+        B = (1 << 26) + 1
+
+        device = torch.device(torch.accelerator.current_accelerator() or "cuda")
+
+        # Deterministic non-identity permute: row swap on the T axis.
+        # perm[0] == 1 and perm[1] == 0, so perm[i] != i for every i.
+        perm_cpu = torch.tensor([1, 0], dtype=torch.int32)
+        permute = perm_cpu.to(device)
+
+        # Sparse non-zero lengths at four known positions. Total = 11.
+        lengths_cpu = torch.zeros((T, B), dtype=torch.int32)
+        lengths_cpu[0, 0] = 3
+        lengths_cpu[0, B // 2] = 5
+        lengths_cpu[1, 0] = 2
+        lengths_cpu[1, B - 1] = 1
+        lengths = lengths_cpu.to(device)
+
+        # Distinct indices per segment so the permutation is fully observable.
+        indices_cpu = torch.arange(11, dtype=torch.int32)
+        indices = indices_cpu.to(device)
+
+        # CPU reference oracle — same op, different dispatch.
+        (
+            permuted_lengths_cpu,
+            permuted_indices_cpu,
+            _permuted_weights_cpu,
+        ) = torch.ops.fbgemm.permute_2D_sparse_data(
+            perm_cpu, lengths_cpu, indices_cpu, None, None
+        )
+
+        # GPU op under test. Pre-fix, this launch trips
+        # KernelLauncher::checkThreadCountNotExceeded on ROCm.
+        (
+            permuted_lengths_gpu,
+            permuted_indices_gpu,
+            permuted_weights_gpu,
+        ) = torch.ops.fbgemm.permute_2D_sparse_data(
+            permute, lengths, indices, None, None
+        )
+
+        torch.testing.assert_close(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
+        torch.testing.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
+        self.assertIsNone(permuted_weights_gpu)
+
+    @unittest.skipIf(*gpu_unavailable)
+    # Skip on GPUs with insufficient HBM (need ~512 MB for the int32
+    # lengths tensor at the chosen B).
+    @unittest.skipIf(*gpu_memory_lt_gb(4))
+    def test_permute_sparse_features_large_grid(self) -> None:
+        """
+        Reproduces the HIP grid-overflow bug in permute_sparse_features_cuda
+        and verifies output correctness at the same scale.
+
+        With BT_blocks=32 and dim3(32, 32) (block size 1024), the
+        permute_indices_weights_kernel launch grid is
+        cuda_calc_block_count(B*T, 32). For B*T > 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.
+
+        Uses ``T=2, B=2**26+1`` so ``B*T = 2**27 + 2`` strictly trips the
+        threshold. ``lengths`` is sparse: all zero except for four known
+        non-zero positions (one per row, plus one mid-row), so HBM usage
+        stays bounded (~537 MB int32) while the permutation logic is
+        still exercised. ``permute = [1, 0]`` is a deterministic row swap
+        on the T axis with ``perm[i] != i`` for every i, so any
+        "kernel computed identity instead of permutation" or wrong-``b_t``
+        decoding bug surfaces in the assertion below.
+        """
+
+        # Choose B*T so that total threads strictly exceeds 2**32:
+        # cuda_calc_block_count(B*T, 32) * 1024 ~= B*T * 32; need B*T > 2**27.
+        T = 2
+        B = (1 << 26) + 1
+
+        device = torch.device(torch.accelerator.current_accelerator() or "cuda")
+
+        # Deterministic non-identity permute: row swap on the T axis.
+        # perm[0] == 1 and perm[1] == 0, so perm[i] != i for every i.
+        perm_cpu = torch.tensor([1, 0], dtype=torch.int32)
+        permute = perm_cpu.to(device)
+
+        # Sparse non-zero lengths at four known positions. Total = 11.
+        lengths_cpu = torch.zeros((T, B), dtype=torch.int32)
+        lengths_cpu[0, 0] = 3
+        lengths_cpu[0, B // 2] = 5
+        lengths_cpu[1, 0] = 2
+        lengths_cpu[1, B - 1] = 1
+        lengths = lengths_cpu.to(device)
+
+        # Distinct indices per segment so the permutation is fully observable.
+        indices_cpu = torch.arange(11, dtype=torch.int32)
+        indices = indices_cpu.to(device)
+
+        # CPU reference oracle — same op, different dispatch.
+        (
+            permuted_lengths_cpu,
+            permuted_indices_cpu,
+            _permuted_weights_cpu,
+        ) = torch.ops.fbgemm.permute_sparse_features(
+            perm_cpu, lengths_cpu, indices_cpu, None
+        )
+
+        # GPU op under test. Pre-fix, this launch trips
+        # KernelLauncher::checkThreadCountNotExceeded on ROCm.
+        (
+            permuted_lengths_gpu,
+            permuted_indices_gpu,
+            permuted_weights_gpu,
+        ) = torch.ops.fbgemm.permute_sparse_features(permute, lengths, indices, None)
+
+        torch.testing.assert_close(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
+        torch.testing.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
+        self.assertIsNone(permuted_weights_gpu)
+
 
 extend_test_class(PermuteIndicesTest)