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

kernels/gemm/preshuffle_gemm.py

@@ -120,6 +120,7 @@ def compile_preshuffle_gemm(
     use_async_copy: bool = False,
     xcd_swizzle: int = 0,
     lds_stage: int = 2,
+    preload: Optional[tuple[int, int]] = None,
 ):
     """Compile preshuffle GEMM (fp8/int8/fp16/bf16).
     Signature: fn(C, A, B, scale_a, scale_b, bias, M, N, stream). bias is the fused
@@ -147,13 +148,14 @@ def compile_preshuffle_gemm(
     is_8bit = is_fp8 or is_int8
     elem_bytes = 1 if is_8bit else 2
 
-    # The async gmem->LDS DMA (buffer_load_lds 128b) only lowers for 8-bit inputs.
-    if use_async_copy and not is_8bit:
-        raise ValueError("use_async_copy is only supported for 8-bit inputs (fp8/int8)")
-
     gpu_arch = get_rocm_arch()
     is_gfx942 = str(gpu_arch).startswith("gfx942")
     is_gfx950 = str(gpu_arch).startswith("gfx950")
+    if use_async_copy and not is_gfx950:
+        # buffer_load_lds only lowers on gfx950; without this the failure surfaces much
+        # later as an unactionable legalization error.
+        raise ValueError(f"use_async_copy requires gfx950, got {gpu_arch}")
+
     use_mfma_scale_128 = is_fp8 and is_gfx950 and (tile_k % 128 == 0)
     use_mfma_k32 = is_f16_or_bf16 and is_gfx950
 
@@ -195,9 +197,15 @@ def compile_preshuffle_gemm(
     num_b_loads = (tile_n * tile_k * elem_bytes) // total_threads // 16
     num_ds_load = (tile_m * tile_k * elem_bytes) // 64 // 16  # A LDS reads per wave
     num_gmem_loads = num_a_loads + num_b_loads
-    if is_8bit and is_gfx950:
+    if preload is not None:
+        if len(preload) != 2 or any(not isinstance(v, int) or v < 0 for v in preload):
+            raise ValueError(f"preload must be a pair of non-negative ints, got {preload!r}")
+        dsrd_preload, dvmem_preload = (int(v) for v in preload)
+    elif is_8bit and is_gfx950:
         dsrd_preload, dvmem_preload = _get_preload(tile_m, tile_n, tile_k)
     else:
+        # _TILE_PRELOAD_TABLE only covers 8-bit tiles, so other dtypes emit no
+        # sched_vmem/sched_dsrd hints unless the caller supplies them via preload=.
         dsrd_preload, dvmem_preload = (0, 0)
 
     a_lds_elems = tile_m * tile_k
@@ -273,15 +281,23 @@ def kernel_gemm(
         thr_g2r_B = fx.make_tiled_copy_B(buf_copy, tiled_mma).get_slice(tid)
 
         lds = fx.SharedAllocator().allocate(SharedStorage).peek()
-        if const_expr(is_8bit):
+        # dma_a_to_lds writes A with k_swz derived from k_blocks16, so whenever the DMA
+        # is used the LDS view must be swizzled the same way or the reader unswizzles at
+        # a different width. The sync path writes through this same view and is
+        # self-consistent under any swizzle, and measurably prefers the fixed one
+        # (bf16 tile_k=256: 1011 TF/s fixed vs 417 derived), so it keeps Swizzle<3,3,3>.
+        if const_expr(is_8bit or use_async_copy):
             k_blocks16 = (tile_k * elem_bytes) // 16
             if k_blocks16 <= 0 or (k_blocks16 & (k_blocks16 - 1)) != 0:
                 raise ValueError(
-                    f"Unsupported tile_k for 8-bit LDS swizzle: tile_k={tile_k}, elem_bytes={elem_bytes} (k_blocks16={k_blocks16}); "
-                    "expected tile_k*elem_bytes to be a positive multiple of 16 with (tile_k*elem_bytes/16) a power of two."
+                    f"Unsupported tile_k for LDS swizzle: tile_k={tile_k}, elem_bytes={elem_bytes} "
+                    f"(k_blocks16={k_blocks16}); expected tile_k*elem_bytes to be a positive multiple "
+                    "of 16 with (tile_k*elem_bytes/16) a power of two."
                 )
             swz_bits = k_blocks16.bit_length() - 1  # log2
-            swz = fx.SwizzleType.get(swz_bits, 4, swz_bits)
+            # base = log2(elements per 16B): 4 for 8-bit, 3 for f16/bf16
+            swz_base = (16 // elem_bytes).bit_length() - 1
+            swz = fx.SwizzleType.get(swz_bits, swz_base, swz_bits)
         else:
             swz = fx.SwizzleType.get(3, 3, 3)
 
@@ -325,7 +341,16 @@ def _make_sA(arr):
             # Bound to the real M extent (like the sync gA) so ragged-M blocks DMA-read
             # OOB rows as 0 instead of faulting past the allocation.
             gA_flat = fx.rocdl.make_buffer_tensor(
-                fx.Tensor(fx.make_view(fx.get_iter(arg_a), fx.make_layout(65536 * K, 1))),
+                # Byte-typed on both sides: the DMA is a raw 128b byte move and the
+                # index math below is already expressed in bytes. With an element-typed
+                # global view the copy only legalizes -- and only indexes correctly --
+                # when elem_bytes == 1, which is what restricted this path to 8-bit.
+                fx.Tensor(
+                    fx.make_view(
+                        fx.recast_iter(Int8, fx.get_iter(arg_a)),
+                        fx.make_layout(65536 * K * elem_bytes, 1),
+                    )
+                ),
                 max_size=False,
                 num_records_bytes=fx.Int64(i32_m) * fx.Int64(K) * fx.Int64(elem_bytes),
             )

tests/kernels/test_preshuffle_gemm.py

@@ -159,8 +159,6 @@ def test_mfma_a8_flyc_preshuffle(
         # operator's operand!"), while CDNA4 (gfx950) handles it. Restrict async
         # copy to gfx950 until the gfx942 codegen path is supported.
         pytest.skip(f"async copy (buffer_load_lds) is only supported on gfx950, not {get_rocm_arch()}")
-    if use_async_copy and in_dtype not in ("fp8", "int8"):
-        pytest.skip("async copy (buffer_load_lds) only supports 8-bit inputs (fp8/int8)")
     print("=" * 80)
     print(f"[flyc] MFMA {in_dtype.upper()} GEMM Test (Tile: {tile_m}x{tile_n}x{tile_k})")
     print("=" * 80)
@@ -1003,3 +1001,39 @@ def test_preshuffle_accepts_whole_a_tile(tile_m):
     if get_rocm_arch() not in ("gfx942", "gfx950"):
         pytest.skip(f"v2 preshuffle GEMM requires gfx942/gfx950, got {get_rocm_arch()}")
     compile_preshuffle_gemm(N=1024, K=2048, tile_m=tile_m, tile_n=256, tile_k=64, in_dtype="bf16", out_dtype="bf16")
+
+
+@pytest.mark.parametrize("in_dtype", ["fp16", "bf16"])
+@pytest.mark.parametrize("tile_k", [64, 128])
+def test_preshuffle_async_copy_2byte_dtypes(in_dtype, tile_k):
+    """The gmem->LDS DMA must match the sync path for 2-byte input dtypes.
+
+    tile_k is parametrized because the DMA's LDS swizzle width derives from
+    ``tile_k * elem_bytes``, so tile_k=128 exercises a different swizzle than the
+    tile_k=64 case that the fixed Swizzle<3,3,3> happened to match.
+    """
+    if get_rocm_arch() != "gfx950":
+        pytest.skip(f"async copy (buffer_load_lds) requires gfx950, got {get_rocm_arch()}")
+    test_mfma_a8_flyc_preshuffle(
+        in_dtype,
+        M=512,
+        N=1024,
+        K=2048,
+        tile_m=128,
+        tile_n=256,
+        tile_k=tile_k,
+        use_async_copy=True,
+        test_graph=False,
+        run_aiter_bench=False,
+    )
+
+
+@pytest.mark.parametrize("bad", [(1,), (1, 2, 3), (-1, 0), (0, -1), 4])
+def test_preshuffle_preload_validation(bad):
+    """preload is a public knob, so reject malformed values instead of emitting bad hints."""
+    if get_rocm_arch() not in ("gfx942", "gfx950"):
+        pytest.skip(f"v2 preshuffle GEMM requires gfx942/gfx950, got {get_rocm_arch()}")
+    with pytest.raises((ValueError, TypeError)):
+        compile_preshuffle_gemm(
+            N=1024, K=2048, tile_m=128, tile_n=256, tile_k=64, in_dtype="bf16", out_dtype="bf16", preload=bad
+        )