[UP][release/2.13] Guard native BMM outer-product launches against the HIP work-item limit

[naromero77amd]

Resolves: https://amd-hub.atlassian.net/browse/ROCM-29320. Request is from a high-priority customer.

Backport of landed pytorch/pytorch#194131 to release/2.13 as a single cherry-pick -x of upstream commit 58a938a.

Summary

• decline the eager native outer-product BMM specialization when its launch would
  exceed HIP's work-item limit, and fall back to ATen
• share the launch configuration between the guard and the launch, and pin
  num_warps, so the checked and launched grids cannot drift
• add a ROCm regression that pins the boundary from both sides

HIP rejects a launch once gridDim.x * blockDim.x reaches 2^32. The eager
outer-product kernel launches a 1D grid with one program per (batch, M tile, N
tile) using a fixed configuration and, unlike Inductor, cannot grow its tiles to
shrink the grid, so it declines the specialization and falls back to ATen.

The limit is on total work items, so the boundary moves with the wave size and
the guard reads it from the device. For M = N = 1 that is B = 16,777,216 on
wave64 (CDNA) and B = 33,554,432 on wave32 (RDNA).

At these sizes ATen can also fail on some architectures, but it raises a clean,
catchable AcceleratorError; the rejected Triton launch can leave the HIP
context unusable.

Why this needs a guard

On gfx90a the rejected launch does not fail cleanly. Triton reports
RuntimeError: Triton Error [HIP]: Code: 1, Messsage: invalid argument, which
its knobs machinery then mangles into a SystemError, and the HIP context is
left unusable—the next allocation fails with
AcceleratorError: CUDA error: invalid argument.

Boundary verification (gfx90a, wave64)

Measured by calling the launcher directly, bypassing the guard:

B (M = N = 1)	work items	guard	unguarded launch
16,777,215	4,294,967,040	accept	succeeds, results correct
16,777,216	4,294,967,296 (2^32)	decline	fails, HIP context left unusable

The guard declines exactly the launches that fail and accepts every launch that
works.

Test Plan

The upstream PR passed its full BMM outer-product suite and its ROCm regression
ran successfully on MI350 (gfx950, wave64).

The release backport's adapted files pass targeted lint:

lintrunner --take FLAKE8,PYFMT,RUFF,CODESPELL,NEWLINE,SPACES,TABS,TESTOWNERS,TEST_HAS_MAIN \
  test/test_bmm_outer_product.py \
  torch/_native/ops/bmm_outer_product/triton_impl.py \
  torch/_native/ops/bmm_outer_product/triton_kernels.py

release/2.13 predates the upstream branch's device-generic test harness and
COW/instrumentation refactors. Conflict resolution preserves the release
branch's existing behavior while applying the landed launch guard, shared
configuration, explicit warp count, and regression test.

Made with Cursor

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @jataylo @hongxiayang @pragupta @jerrymannil @xinyazhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben