[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,10 +148,6 @@ 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")
@@ -179,12 +176,29 @@ def compile_preshuffle_gemm(
     a_load_bytes = 16
     bytes_per_thread_a = (tile_m * tile_k * elem_bytes) // total_threads
     num_a_loads = bytes_per_thread_a // a_load_bytes
+    # The A tile is copied by total_threads threads in a_load_bytes chunks, and both
+    # divisions above truncate. Any remainder leaves the tail of the tile unfetched and
+    # the kernel then computes on stale LDS, so check the exact condition on the tile
+    # size itself rather than on either truncated intermediate.
+    a_tile_bytes = tile_m * tile_k * elem_bytes
+    a_copy_granularity = total_threads * a_load_bytes
+    if a_tile_bytes % a_copy_granularity != 0:
+        raise ValueError(
+            f"tile_m * tile_k * elem_bytes must be a multiple of {a_copy_granularity} "
+            f"(total_threads * a_load_bytes); got tile_m={tile_m}, tile_k={tile_k}, "
+            f"elem_bytes={elem_bytes} -> {a_tile_bytes} bytes, leaving "
+            f"{a_tile_bytes % a_copy_granularity} bytes of the A tile unloaded"
+        )
     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:
+        dsrd_preload, dvmem_preload = 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
@@ -260,15 +274,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)
 
@@ -312,7 +334,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),
             )
@@ -459,8 +490,12 @@ def hot_loop_scheduler():
 
         # ── Pipeline stage (double-buffered B via split fragments) ─
         def mma_kloop(a_stage, cur_frag_B):
-            fx.copy(uni_copy, pA_s2r_stages[a_stage], frag_A_retile)
+            # Issue the A-fragment LDS reads per k-step. Semantically this is the same as
+            # one whole-tile copy, but the single copy leaves the scheduler no room to
+            # interleave the reads with the MFMAs that consume them, which costs 3.3x on
+            # bf16 tile_k=256 (see the commit message for the measurement).
             for ki in range_constexpr(k_iters):
+                fx.copy(uni_copy, pA_s2r_stages[a_stage][None, None, ki], frag_A_retile[None, None, ki])
                 k_coord = ki if (use_mfma_scale_128 or use_mfma_k32) else (None, ki)
                 fx.gemm(tiled_mma, frag_C, frag_A[None, None, k_coord], cur_frag_B[None, None, k_coord], frag_C)
 

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)
@@ -980,3 +978,51 @@ def _args(c, a, b, sa, sb, bs):
         f"✓ Fused epilogue {epilogue} correctness verified: "
         f"max_abs_diff={max_diff:.4f}, max_rel={rel:.4f}, ref_max={ref.abs().max().item():.2f}"
     )
+
+
+@pytest.mark.parametrize("tile_m", [48, 80, 112, 144])
+def test_preshuffle_rejects_partial_a_tile(tile_m):
+    """Reject tile_m values that leave a partial 16B A load per thread.
+
+    Each of the 256 threads copies tile_m * tile_k * elem_bytes / 256 bytes of the A
+    tile in 16B chunks. When that is not a multiple of 16 the chunk count truncates and
+    the tail of the A tile is never fetched, which used to yield wrong results with no
+    diagnostic. For bf16 with tile_k=64 the product must be a multiple of 4096.
+    """
+    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, match=r"must be a multiple of"):
+        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("tile_m", [64, 96, 128, 160])
+def test_preshuffle_accepts_whole_a_tile(tile_m):
+    """tile_m values that divide the A tile evenly across threads must still compile."""
+    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,
+    )