SIGN IN SIGN UP

Fix pdist and cdist launch failure past 2^24 outputs on ROCm (#188006)

Fixes https://github.com/pytorch/pytorch/issues/168868

`pdist_forward_kernel_impl` and `cdist_kernel_impl` launch one 256-thread block per output element. HSA expresses a grid dimension in work-items rather than workgroups, and `rocminfo` reports a per-dimension `Grid Max Size` of 4294967295, so `gridDim.x * blockDim.x` has to fit in uint32. With `blockDim.x` of 256 that caps a launch at 2^24 outputs. Bisected on gfx90a and gfx950 under HIP 7.14, pdist at n=5793 (16776528 blocks) succeeds and n=5794 (16782321 blocks) raises `hipErrorInvalidConfiguration`; cdist crosses the same line at exactly 16777216 outputs. `test_pdist_norm_large` uses n=50000, which needs 1249975000 blocks or roughly 74x the limit, which is why it was skipped on ROCm. CUDA expresses `maxGridSize[0]` in blocks, so the same launch is legal there.

The failure mode depends on the HIP runtime rather than the architecture. HIP 7.14 rejects the oversized dispatch outright, while HIP 7.2 on the same MI355X silently produces wrong results with outputs left unwritten. Both are the same oversized launch, so both are fixed by bounding it: each launch caps the grid at 32 blocks per CU through a shared `dist_grid` helper and grid-strides over the remaining outputs. `cdist` is fixed alongside `pdist` because it has the identical bug and fails at 2^24 outputs for the same reason.

The alternative to capping the grid is to keep one block per output and spread the launch across a 2-D grid, which would preserve the existing block-to-output mapping exactly. The cap plus grid-stride loop was chosen instead because it bounds the launch on CUDA as well, keeps a single index expression in the kernel, and measured within a few percent of the previous cost.

### Test Plan

On gfx950 (MI355X, HIP 7.14):

```
python -m pytest test/test_torch.py -k "pdist or cdist" -v
python -m pytest test/test_ops.py -k "cdist or pdist" -q
python -m pytest test/test_autograd.py test/test_nn.py -k "cdist or pdist or pairwise" -q
```

27 passed / 9 skipped, 81 passed / 15 skipped / 5 xfailed, and 31 passed / 2 xfailed respectively. `test_pdist_norm_large_cuda` passes in 4.7s once it is no longer skipped on ROCm. Note that it cannot be exercised through `PYTORCH_TEST_WITH_ROCM`, since `TEST_WITH_ROCM` is derived from `torch.version.hip` and `skipIfRocm` therefore always skips it.

Forward output was also compared against CPU for p in {0, 0.5, 1, 2, 3, inf} above the old boundary, for batched cdist, and for degenerate shapes, plus n=50000 end to end.

Performance, 10-iteration averages on MI355X, before versus after:

```
pdist n=4096 m=128     4.220 ms -> 4.283 ms
pdist n=4096 m=8       3.683 ms -> 3.928 ms
pdist n=2048 m=128     1.063 ms -> 1.082 ms
cdist 2048x4096 m=128  9.468 ms -> 9.384 ms
cdist 2048x4096 m=8    8.155 ms -> 8.769 ms
cdist 1024x4096 m=128  2.643 ms -> 2.774 ms
```

Original investigation and reproduction by @geozhai, who identified the oversized launch grid and the ROCm skip that was masking it. Portions of the analysis in this description were drafted with an AI assistant (Claude Code).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/188006
Approved by: https://github.com/jeffdaily

Co-authored-by: Jeff Daily <jeff.daily@amd.com>
G
geozhai committed
32ba93d1ef6f675561a99776fc415a22c9f64870
Parent: 45592f3
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/24/2026, 11:42:30 PM