[AMDGPU] Add cross-platform min_blocks_per_cu occupancy hint

[paveltc]

Summary

Adds an optional, portable per-kernel occupancy hint:
@qd.kernel(min_blocks_per_cu=N). It requests that the scheduler keep at
least N thread-blocks resident per compute unit, and lowers to each backend's
native occupancy control.

API impact

Adds one new standard @qd.kernel parameter (min_blocks_per_cu: positive
int or None; default None). It is cross-platform by construction:

• CUDA parity: lowers to the minctasm launch bound, and parameterizes
  the previously hardcoded default of 2 (unset preserves that default).
• Platform-independent naming: min_blocks_per_cu (blocks per compute
  unit) — never amd_*.
• Ignored on CPU/Metal (no occupancy concept).

None is a true no-op, so existing code is unaffected.

Lowering

Backend	Lowers to	Conversion
CUDA	minctasm launch bound	1:1 (a CTA is a thread-block)
AMDGPU	amdgpu-waves-per-eu	ceil(min_blocks_per_cu * ceil(block_dim/64) / 4) (wave64, 4 SIMDs/CU)
CPU / Metal	—	ignored

Behavior & caching

• Validated at decoration time (positive int or None).
• Applied to both the primal and adjoint (autodiff) kernels.
• Participates in both the C++ offline cache key and the Python fastcache
  key, so kernels differing only in min_blocks_per_cu never collide or
  reuse stale code.

Due diligence

Experiments have been done to implement this without a user-facing knob.
The only real way to do it would be to implement a tuning loop where different min_blocks_per_cu are set and performance tests are run on each value set.

This is borne out by measuring the same hint on two platforms (details in
the two sections below): the optimal setting differs by platform and
hardware. On AMD MI308X the CG workloads want min_blocks_per_cu=4
(+1–3%), whereas on an NVIDIA B300 the same scenes — running 10–70× faster
and no longer latency-bound — are best at the default and slightly regress at
higher occupancy. There is therefore no single static default (or backend
heuristic) that is correct everywhere; the only automatic alternative is a
per-workload autotuning loop. That is precisely why the feature is an
opt-in per-kernel hint (default None = a no-op that preserves current
behavior on every backend) rather than a forced default.

Testing — tests/benchmarks/test_rigid.py (AMD MI308X)

Validated with the Genesis rigid-body benchmark suite (test_rigid.py),
runtime_fps metric, on the memory-latency-bound CG constraint-solver
path. Three arms of the same benchmark differ only by the hint applied to the
rigid-solver kernels (kernel_step_1, kernel_step_2 in
rigid_solver.py; func_solve_init, _kernel_solve_monolith in
constraint/solver.py), each gated on an env var so an arm is selected purely
by the environment of a fresh process:

• A — off: hint absent (compiler default occupancy)
• B: min_blocks_per_cu=2
• C: min_blocks_per_cu=4

Arms are interleaved per round across 5 rounds to cancel thermal/clock drift;
deltas are the geomean of per-round paired ratios vs the off arm.

Workload (CG)	n_envs	off (fps)	mbpc=2	mbpc=4
anymal	8,192	218,788	+1.74%	+2.81%
go2 (warm rounds)	4,096	879,497	+1.77%	+1.22%
franka	30,000	1,791,167	−0.21%	−0.16%

• anymal — monotonic wins every round (mbpc=2 → 4/5, mbpc=4 → 5/5).
• go2 — rounds 1–2 are cold-start (first per-arm JIT) artifacts; warm
  rounds (r3–5) are the true signal. No regression once cold start is excluded.
• franka — both arms within ±0.2% (2/5 wins): statistically
  indistinguishable from noise. A low-DOF arm at ~1.79M fps is not
  latency-bound, so extra occupancy hides nothing → flat, as expected.

Bottom line: a small, consistent, workload-dependent speedup on the
latency-bound CG path, neutral where the workload isn't latency-bound, and
no regression observed anywhere. Because the benefit is workload-dependent
with no static rule to decide it automatically, the hint is an opt-in knob
rather than a default.

Environment (AMD)

• GPU: AMD Instinct MI308X — gfx942 (CDNA3), ROCm 7.2.4
• quadrants: feat/cross-platform-occupancy (base 96c594499), LLVM 22.1.0
• Genesis: genesis-world main @ eeb91f8
• PyTorch: 2.10.0+rocm7.0, Python 3.10.12 · run 2026-08-25

Cross-platform validation — CUDA (NVIDIA B300)

Same PR, exercised on CUDA to confirm the hint is portable (not AMD-only).

Functional — pass. PTX launch bounds (print_kernel_asm):

min_blocks_per_cu	PTX
unset	.minnctapersm 2
2	.minnctapersm 2
4	.minnctapersm 4
8	.minnctapersm 8

• Confirms the CUDA lowering and that unset ≡ 2 — i.e. no behavior
  change on CUDA unless the hint is set.
• Decorator rejects 0, -1, True, 1.5 (QuadrantsSyntaxError).
• Cache isolation: =4 vs =8 produce distinct fastcache keys.

Performance. Same CG protocol (5 interleaved rounds, 20 s warmup / 8 s
record, geomean vs off = minctasm 2):

Workload (CG)	n_envs	off (fps)	mbpc=4	mbpc=8
anymal	8,192	2,300,460	−1.26%	−5.60%
go2	4,096	4,598,749	−0.13%	−2.33%
franka	30,000	14,817,374	+0.18%	−0.27%

On B300 these CG scenes run ~10–70× the MI308X fps and are not
latency-bound, so extra occupancy doesn't help: =4 is within noise, =8
is a small, consistent loss (register pressure), franka is flat. Because the
default is a no-op (unset ≡ minctasm 2), there is no CUDA regression for
non-opted-in code.

Takeaway. The lowering is correct on both platforms; the occupancy
sweet-spot is workload- and hardware-dependent (helps AMD MI308X CG at 4;
neutral/slightly-worse on B300), which is exactly why this is an opt-in
per-kernel hint rather than a forced default.

Environment (CUDA)

• GPU: NVIDIA B300 SXM6 (sm_100), GPU 0
• quadrants: PR head 487002d
• Genesis: genesis-world @ eeb91f8 (four-kernel GS_MIN_BLOCKS_PER_CU patch)
• PyTorch: 2.13.0+cu130

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 487002d22f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Derive the SIMD count for supported RDNA targets

On AMD RDNA devices, which this repository explicitly supports in quadrants/runtime/amdgpu/jit_amdgpu.cpp:68-73 and llvm_context.cpp:583-588, a CU has two SIMD execution units rather than the four assumed here. For example, a 256-thread wave64 block with min_blocks_per_cu=4 is lowered to amdgpu-waves-per-eu=4 instead of the required 8, so register allocation may permit only half the requested blocks per CU. Select the execution-unit count from the actual mcpu/AMDGPU subtarget rather than fixing it to the CDNA3 value.

Useful? React with 👍 / 👎.

[paveltc]

Good catch — this is a real CDNA-vs-RDNA asymmetry. The 4 SIMDs/CU is CDNA3-specific; RDNA has 2/CU, so as written the hint under-provisions occupancy on RDNA (~half the requested blocks/CU).

Two reasons I'm deferring the mcpu-derived fix to a follow-up rather than taking it here:

1. It's opt-in and never affects correctness — amdgpu-waves-per-eu only bounds the register allocator, so a low value just means less aggressive occupancy, not wrong results. Default (min_blocks_per_cu=None) is a no-op on every target.
2. RDNA is genuinely untested for this. jit_amdgpu.cpp forces wave64 on RDNA hosts (+wavefrontsize64), so the correct waves-per-EU there depends on LLVM's occupancy model in forced-wave64 mode, which I can't confirm without RDNA hardware. This PR is tuned/validated only on CDNA3 (MI308X) and CUDA, so I'd rather not swap a CDNA-correct constant for an unvalidated RDNA formula in the same PR.

I've pushed a comment clarifying that 4 SIMDs/CU is CDNA-specific and that mcpu-based SIMD selection is deferred to a follow-up with RDNA hardware to validate the wave64-on-RDNA occupancy model. Thanks!