i[gfx1250] Switch BF16 GEMM LDS loads to ds_read_b256, +0.4 TB/s

[amd-hhashemi]

Summary

The BF16 GEMM kernel on gfx1250 reads WMMA operand fragments from LDS using two ds_read_b128 instructions per fragment (16 bytes
each, 32 bytes apart), followed by a shuffle. This halves the effective latency-hiding window per K-step — there are only 2×KS_DS
issue slots between read issuance and the s_wait_dscnt stall, which is shorter than LDS round-trip latency at this tile size,
causing exposed wait cycles visible in ATT traces.

This PR switches to a single ds_read_b256 (32 bytes) per fragment by:

1. LDS_PAD: 16 → 32 bytes — increases row padding so every LDS row starts on a 32-byte boundary, required for aligned ds_read_b256
   accesses.
2. _frag: 2×ds_read_b128 + shuffle → single ds_read_b256 — halves KS_DS from 16 to 8 reads per WMMA K-step, doubling the time each
   batch has to complete before its s_wait_dscnt.
3. kgrp16 → kgrp32 in load_a/load_b — adapts the per-lane WMMA operand addressing to the new contiguous 32-byte layout: kgrp=0
   gets K-elements 0–15, kgrp=1 gets K-elements 16–31 (sequential split, verified correct by the WMMA instruction's input layout on
   gfx1250).
4. DS_A = DS_B = 2 → 1 — updates the read-count accounting used by s_wait_dscnt thresholds.

The TDM prefetch issue() is also moved to before the WMMA K-step loop so the sched_barrier pair doesn't interrupt the expert
scheduler's mid-loop read/WMMA interleaving.

Result

~3% bandwidth improvement on the target shape (M=64, N=65536, K=16384, tiles=64×256×128, warps=1×4, nb=3). Verified correct on
multiple shapes.

Test plan

• python3 tests/kernels/test_gemm_bf16_gfx1250.py -mnk 64,65536,16384 -tiles 64,256,128 -warps 1,4 -nb 3 -cluster 1,1 -bench —
  correctness PASS, BW improvement visible
• python3 -m pytest tests/kernels/test_gemm_bf16_gfx1250.py -v — all parametrized cases pass

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This PR aims to improve gfx1250 BF16/FP16 GEMM performance by changing the LDS operand load path to use wider 256-bit LDS reads (ds_read_b256), adjusting LDS row padding/alignment, and updating the expected DS-read accounting per WMMA K-step.

Changes:

• Increase LDS row padding (LDS_PAD) to 32 bytes to support 32-byte alignment for intended 256-bit LDS reads.
• Switch operand fragment loads from 128-bit LDS reads to a 256-bit LDS read path and update the per-K-step DS-read count.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.