[Kernel][Perf] Enable the gmem->LDS async copy for f16/bf16 preshuffle GEMM

[JohnQinAMD]

Summary

use_async_copy=True was rejected for anything but fp8/int8. That was never a hardware limit — buffer_load_dwordx4 … lds moves raw bytes — it came from the global-side view being element-typed while the LDS side is byte-typed, so the copy only legalized, and only indexed correctly, when elem_bytes == 1:

failed to legalize operation 'fly.copy_atom_call' ...
(!fly.copy_atom<buffer_copy_lds<128>, 128>,
 !fly.memref<bf16, buffer_desc, 1:0>, !fly.memref<i8, shared, 1:1>)

Enabling it drops ds_write from 32 to 0 per k-tile and raises MfmaUtil from 73.3% to 78.8%. MemUnitStalled is 0.04% on both sides, so this is an issue-rate win, not a bandwidth one.

Stacked on #1007. Without the mma_kloop fix in that PR the DMA is worth
+0.9% rather than +8.2%.

Changes

1. Byte-type the global view so the copy is i8 → i8 and the
   byte-denominated index math in dma_a_to_lds is correct for
   elem_bytes != 1.

2. Derive the LDS swizzle from tile_k when the DMA is in use.
   dma_a_to_lds computes k_swz from k_blocks16 = tile_k * elem_bytes / 16,
   while the non-8-bit branch hardcoded Swizzle<3,3,3>; those agree only at
   tile_k * elem_bytes == 128, so tile_k=128 would write a 16-wide swizzle
   that the reader unswizzles as 8-wide.

   The sync path deliberately keeps the fixed swizzle. It writes through the
   same view, so it is self-consistent under any swizzle, and it measurably
   prefers the fixed one — bf16 128×256×256, sync: 1011 TF/s with
   Swizzle<3,3,3> vs 417 with the derived value. Only the DMA path needs
   the two sides to agree, so only the DMA path changes.

3. Add preload=. _TILE_PRELOAD_TABLE is tuned for 8-bit tiles and is
   consulted only when is_8bit and is_gfx950, so f16/bf16 emit no
   sched_vmem/sched_dsrd hints at all. The argument lets a caller supply
   them; the default preserves existing behaviour.

Performance

Environment: MI355X (gfx950, 256 CU), ROCm 7.2.3 / PyTorch 2.11.0 / FlyDSL 0.2.4, base commit ac227c3 (main as of #1006) + #1007.

bf16, M=54560, K=8192, N=8192. Before this change bf16 has no DMA path at all, so the baseline is the best sync configuration. Both sides are best-of a tile_m × preload sweep measured in one process against the same torch reference (1615 TF/s), median of 20 timed iterations after 5 warmups:

compile_preshuffle_gemm(
    N=N, K=K, tile_m=tm, tile_n=256, tile_k=64, in_dtype="bf16", out_dtype="bf16",
    xcd_swizzle=4, preload=(p, p), use_async_copy=<False|True>,
)                       # tm in {128, 160, 192}, p in {2, 4, 8}

	Before (sync only)	After (DMA available)
best of sweep	1385 TF/s (128×256×64, preload 4)	1499 TF/s (160×256×64, preload 2)	+8.2%

Testing

pytest tests/kernels/test_preshuffle_gemm.py -c tests/pytest.ini -k async_copy_2byte -v
#   4 passed

pytest tests/kernels/test_preshuffle_gemm.py -c tests/pytest.ini -q
#   154 passed, 10 skipped

ruff check kernels/gemm/preshuffle_gemm.py tests/kernels/test_preshuffle_gemm.py
#   All checks passed!

Adds test_preshuffle_async_copy_2byte_dtypes (fp16/bf16 × tile_k 64/128, where 128 exercises a different swizzle than the 64 the fixed value happened to match). Removing the non-8-bit skip also means the existing sync_copy/async_copy parametrization covers fp16/bf16 for the first time — that is the jump from 126 to 154 passing against #1007.

The 10 skips are the pre-existing MXFP4 fp8-A cases. ruff format --check reports one pre-existing violation in test_preshuffle_gemm.py that this PR neither adds to nor fixes.

Not tested on gfx942 — buffer_load_lds lowering is gfx950-only and the new test skips accordingly.

Breaking Changes

None. use_async_copy and preload both default to previous behaviour, and the swizzle change is scoped to the DMA path, so no previously reachable configuration changes.

[Copilot]

Pull request overview

Enables the gmem→LDS async-copy (buffer_load_lds) path for fp16/bf16 in the preshuffle GEMM kernel on gfx950 by making the DMA source byte-typed, aligning LDS swizzle between writer/reader in the DMA path, and exposing scheduling prefetch controls to callers.

Changes:

• Add preload= to allow explicit sched_vmem/sched_dsrd hint tuning beyond the existing 8-bit tile table.
• Fix A-tile DMA legalization/indexing for 2-byte dtypes by byte-typing the global view and deriving the LDS swizzle when DMA is enabled.
• Expand tests to cover 2-byte async copy and A-tile copy-granularity validation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
kernels/gemm/preshuffle_gemm.py	Enables fp16/bf16 async gmem→LDS DMA by using byte-typed global views, DMA-consistent swizzle, new preload control, and improved A-tile validity checks.
tests/kernels/test_preshuffle_gemm.py	Removes the non-8-bit async-copy skip and adds coverage for fp16/bf16 async-copy correctness and A-tile granularity validation.

Suppressed comments (1)

kernels/gemm/preshuffle_gemm.py:155

• use_async_copy is still gfx950-only (buffer_load_lds lowering). Without an explicit arch check here, callers can enable use_async_copy=True on gfx942/other arches and hit a much less actionable legalization/codegen failure later in compilation. Add a clear early ValueError when use_async_copy is requested on a non-gfx950 arch.

    gpu_arch = get_rocm_arch()
    is_gfx942 = str(gpu_arch).startswith("gfx942")
    is_gfx950 = str(gpu_arch).startswith("gfx950")
    use_mfma_scale_128 = is_fp8 and is_gfx950 and (tile_k % 128 == 0)
    use_mfma_k32 = is_f16_or_bf16 and is_gfx950

---

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

[JohnQinAMD]

Thanks — both addressed in 15c0fa7.

• Arch guard: use_async_copy now raises early off gfx950. Worth noting this was already reachable for 8-bit before this PR (there was no arch check, only is_8bit), but enabling f16/bf16 makes it much easier to hit, so it belongs here.
• preload validation: rejects anything that is not a pair of non-negative ints, with test_preshuffle_preload_validation covering wrong arity, negatives and non-tuples.

Local suite after the change: 159 passed, 10 skipped (the 10 are the pre-existing MXFP4 fp8-A cases).

One thing I can't cover from my side: the test job shows as SKIPPED on both PRs, so no GPU test has run in CI. All my measurements and test runs are on ROCm 7.2.3 / PyTorch 2.11 rather than the 7.14 / 2.12 the CI image uses — I tried reproducing on rocm/pytorch:rocm7.14_... with the PyPI flydsl wheel but the codegen needs the custom LLVM that CI builds (lld invocation failed). If someone can trigger the GPU job on linux-flydsl-mi355-8, that would validate this on the supported toolchain.

[JohnQinAMD]

Correction to the performance figures, after re-measuring on freshly built trees.

The +8.2% in the description isolates only the use_async_copy toggle — both sides of that A/B already had preload=, which this PR also introduces. Measured against the actual base (#1007), the PR is worth considerably more:

tree (bf16, M=54560 K=N=8192, best of a tile_m × preload sweep)	TF/s
upstream ac227c3	923
#1007	934
#1007 + this PR	1509 (+61% over #1007)

Decomposed: exposing preload= is worth ~+57% on its own, because bf16 previously received no sched_vmem/sched_dsrd hints at all (_TILE_PRELOAD_TABLE is gated on is_8bit); the DMA then adds +8.2% on top of that.

So the larger half of this PR's value is the scheduling-hint argument rather than the async copy in the title. Happy to retitle or split preload= out if you'd prefer it reviewed separately — it is independent of the other three changes.

Reference points measured in the same process: torch 1618, aiter ASM 1603.