[rocm-libraries] ROCm/rocm-libraries#8566 (commit 55ad753)

[GFX1250][CK_TILE] Coalesce MX scale16 scale load
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Problem

On gfx1250 the MX `scale16` scale load was strided across lanes while
`scale32` was coalesced. `scale16` used an identity host pre-shuffle
plus a K-fastest `[packs_m, MThreadPerXdl, packs_k]` scale descriptor,
so each lane's per-K-iteration offset was `lane * num_scale_k` (a
separate cache line per lane, every K iteration).

## Fix

Put `scale16` on the same **M/N-fastest** `[packs_m, packs_k,
MThreadPerXdl]` descriptor as `scale32`, and fold both host pre-shuffles
into **one formula** in `preShuffleScaleBuffer_gfx1250`. Each lane's
offset becomes `lane * 1` (unit-dword) instead of `lane * num_scale_k`,
so consecutive lanes hit consecutive addresses. This is a bijective
transpose of both the data and the descriptor, so the scale values are
unchanged. `scale32` is byte-for-byte unchanged.

### The layout is parameterized by `WarpTile::M`, not `ScaleBlockSize`

This is the key point from the review of #8202. The single pre-shuffle
formula keys off `MThreadPerXdl` = the WMMA `WarpTile` M (A scales) / N
(B scales) — the number of lanes holding a distinct scale row per warp —
**not** off `ScaleBlockSize`:

- a 32×32 WMMA has 32 distinct scale lanes → `MThreadPerXdl = 32`;
- a 16×16 WMMA has 16 (the other wavefront lanes replicate the same
scale rows) → `MThreadPerXdl = 16`.

Both `scale16` and `scale32` are instantiated with **both** 32×32 and
16×16 WMMA tiles, so the lane count — not the block size — selects the
layout. Deriving it from `ScaleBlockSize` mislays the `scale16` configs
that use a 32×32 `WarpTile`. `ScaleBlockSize` only changes how many
`int32` K-packs a lane holds: `scale16` halves the K span per scale, so
a lane holds twice as many packs and the **device** reads two adjacent
`int32` as one `int64` — a device-side read concern. Host packing stays
`int32` (`PackSize = 4`) for both block sizes.

The callers pass `M_Warp_Tile` (A) / `N_Warp_Tile` (B) into the
pre-shuffle from the pipeline and grouped-gemm MX tests.

## Validation

**Correctness (measured).** `test_ck_tile_mx_gemm_pipeline_tdm_wmma` on
gfx1250, full suite = 148 subtests over 37 configs: **147 / 148 pass**
against the fp32 reference — every `scale16` config (`WarpTile` 16 and
32, TDM V1 and V2) and every `scale32` config. The lone failure,
`/14.SmallM` (Col-major-A FP4×FP8 `scale32`), is **pre-existing and
unrelated**: it reproduces identically on `develop` with this PR
reverted.

**Coalescing (by construction).** The per-lane scale offset going from
`lane * num_scale_k` to `lane * 1` follows directly from the unified
`[packs_m, packs_k, MThreadPerXdl]` descriptor (reviewer-verifiable from
the diff); the host pre-shuffle was checked to reproduce this device
descriptor layout index-for-index for both block sizes.

**Performance (not measured).** Runtime throughput / bandwidth was not
benchmarked here; this PR is validated for correctness and for the
descriptor-level coalescing property, not for end-to-end kernel
performance.