[Kernel][MI350] Add cache-aware 8-wave BF16/FP16 GEMM

[zhanglx13]

Summary

Add a gfx950 8-wave GEMM specialized for BF16 and FP16, plus reusable XCD-aware and grouped-M workgroup mappings. The existing FP8 8-wave GEMM now uses the same composable locality mapping.

The half-precision kernel keeps BF16 and FP16 in one compile-time-specialized implementation because they share the 2-byte operand geometry, 16x16x32 MFMA tiling, LDS layout, pipeline, and epilogue. FP8 remains a separate hot loop.

Motivation

The current generic half GEMM does not implement the slice-MN inter-wave pipeline used by the gfx950 Gluon tutorial. Splitting the 256x256 tile into four independently addressable LDS regions lets two four-wave groups overlap MFMA work with direct-to-LDS loads without overwriting operands that peer waves still consume.

PID locality is independent of wave count, so this PR places it in a general helper rather than an 8-wave-specific utility.

Changes

• Add reusable workgroup mapping operations:
  • round-robin physical XCD IDs to contiguous logical ranges, including uneven grids
  • linear IDs to grouped-M tile traversal with a shortened final group
  • a composed XCD plus grouped-M mapping
• Apply configurable num_xcds and group_size_m mapping to the FP8 8-wave GEMM.
• Add a dedicated gfx950 256x256x64 8-wave half GEMM with:
  • BF16 and FP16 compile-time specializations
  • 2x4 wave layout and four 128x128 output quadrants
  • double-buffered A-top, A-bottom, B-left, and B-right LDS regions
  • direct-to-LDS loads and compiler-visible waitcnt operations
  • four-wave-group phase rotation for inter-wave compute/load overlap
  • cached prepared launch wrapper
• Add BF16/FP16 correctness tests and all four locality ablations:
  • linear
  • XCD only
  • grouped-M only
  • XCD plus grouped-M
• Add prepared benchmark entry points with rotating tensors.

Pipeline design reference: https://github.com/ROCm/gfx950-gluon-tutorials/tree/main/kernels/gemm/inter_wave/a16w16

Performance

Hardware: gfx950 MI350X-profile system, 256 CUs, 1000 W socket power limit. The device market string reports MI355 OAM and the firmware reports MI350X.

Protocol: AMD_SERIALIZE_KERNEL=3, prepared compiled launch, three rotating tensor sets, 10 warmups, 1,000 measured dispatches, and the mean kernel duration of the final 100 matching rocprofv3 dispatches. TFLOP/s is 2MN*K/time.

Dtype / shape	Current 8-wave kernel	New / locality optimized	TFLOP/s gain
BF16, 4096x4096x8192	277.966 us / 988.89 TFLOP/s	239.700 us / 1146.76 TFLOP/s	+15.96%
FP16, 4096x4096x8192	296.150 us / 928.17 TFLOP/s	252.625 us / 1088.09 TFLOP/s	+17.23%
FP8 row-scale, 4096x4096x16384	244.486 us / 2248.66 TFLOP/s	236.239 us / 2327.20 TFLOP/s	+3.49%

BF16 locality ablation within the new pipeline:

num_xcds	group_size_m	Time	TFLOP/s
1	1	247.239 us	1111.79
8	1	247.858 us	1109.01
1	4	238.970 us	1150.26
8	4	239.700 us	1146.76

The generated BF16 kernel uses 128 KiB LDS, 128 VGPRs, 112 SGPRs, and no scratch.

Testing

• ruff check and ruff format --check on changed files
• 6 BF16/FP16 correctness and workgroup-mapping tests on gfx950
• Existing FP8 8-wave row-major and preshuffled-B tests on gfx950
• FP8 correctness for all four mapping ablations
• BF16 and FP16 final-100 rocprofv3 performance runs

Dependencies

• No new third-party dependencies added

Breaking Changes

None. Existing call sites continue to work; the FP8 8-wave kernel gains optional locality controls with gfx950 defaults.