diff --git a/aiter/aot/flydsl/gemm.py b/aiter/aot/flydsl/gemm.py index c2db1829fbd..f429780cc1c 100755 --- a/aiter/aot/flydsl/gemm.py +++ b/aiter/aot/flydsl/gemm.py @@ -12,6 +12,7 @@ Supported kernel families: - ``flydsl_hgemm_*`` gfx950 A16W16 GEMM kernels + - ``flydsl_hgemm_*_gfx1250`` gfx1250 A16W16 GEMM kernels - ``flydsl_bpreshuflle_*`` a8w8 preshuffle GEMM kernels - ``flydsl_bpreshuffle_8w_*`` gfx950 8-wave a8w8 ptpc GEMM kernels - ``flydsl_bpreshuffle_wmma_*`` gfx1250 a8w8 ptpc GEMM kernels @@ -71,8 +72,12 @@ gemm_a16w16_gfx950, make_gemm_a16w16_param_and_validate, ) +from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( + compile_gemm_a16w16 as compile_gemm_a16w16_gfx1250, +) from aiter.ops.flydsl.kernels.kernels_common import run_cached from aiter.ops.flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm +from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg from aiter.ops.flydsl.mxfp8_128_bpreshuffle_gemm_gfx1250 import ( BLOCK_K as SCALE_BLOCK_SIZE, ) @@ -209,7 +214,11 @@ def parse_csv(csv_path: str): params = get_flydsl_hgemm_kernel_params(kernel_name) if params is not None: params = dict(params) - params["kind"] = "hgemm" + params["kind"] = ( + "a16w16_gfx1250" + if params["target_gfx"] == "gfx1250" + else "hgemm" + ) else: params = None @@ -376,6 +385,104 @@ def _compile_hgemm_to_cache( ) +def _compile_a16w16_gfx1250_to_cache( + *, + m: int, + n: int, + k: int, + dtype: str, + out_dtype: str, + block_m: int, + block_n: int, + block_k: int, + stages: int, + split_k: int, + m_waves: int, + n_waves: int, + k_waves: int, + group_m: int, + use_half_tile_interleaved: bool, + has_bias: bool, + target_gfx: str, + **kwargs, +): + del kwargs, group_m + + import torch + + if target_gfx != "gfx1250": + raise ValueError( + f"The FlyDSL gfx1250 A16W16 kernel only supports gfx1250, got {target_gfx}" + ) + if k_waves != 1: + raise ValueError( + f"The FlyDSL gfx1250 A16W16 kernel only supports k_waves=1, got {k_waves}" + ) + + torch_dtype = _torch_dtype_for_kernel(dtype) + torch_out_dtype = _torch_dtype_for_kernel(out_dtype) + if torch_out_dtype not in (torch_dtype, torch.float32): + raise ValueError( + f"Unsupported output dtype {out_dtype!r} for input dtype {dtype!r}" + ) + in_dtype_str = "fp16" if torch_dtype == torch.float16 else "bf16" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get( + torch_out_dtype, "f32" + ) + fx_dtype = { + torch.bfloat16: fx.BFloat16, + torch.float16: fx.Float16, + torch.float32: fx.Float32, + torch.int32: fx.Int32, + } + launch = compile_gemm_a16w16_gfx1250( + N=n, + K=k, + tile_m=block_m, + tile_n=block_n, + tile_k=block_k, + m_warp=m_waves, + n_warp=n_waves, + in_dtype=in_dtype_str, + out_dtype=out_dtype_str, + num_buffers=stages, + add_bias=has_bias, + split_k=split_k, + main_loop_unroll=use_half_tile_interleaved, + ) + + with compile_only_env(): + dev = torch.device("cpu") + representative_extent = 8 + a = torch.empty((1, representative_extent), device=dev, dtype=torch_dtype) + b = torch.empty((1, representative_extent), device=dev, dtype=torch_dtype) + out = torch.empty((1, representative_extent), device=dev, dtype=torch_out_dtype) + bias = torch.empty( + (representative_extent if has_bias else 0,), + device=dev, + dtype=torch_dtype if has_bias else torch_out_dtype, + ) + workspace = torch.empty( + (representative_extent,), device=dev, dtype=torch.float32 + ) + semaphore = torch.zeros((representative_extent,), device=dev, dtype=torch.int32) + stream = fx.Stream(0) + dispatch_args = ( + ptr_arg(out, fx_dtype[out.dtype]), + ptr_arg(a, fx_dtype[a.dtype]), + ptr_arg(b, fx_dtype[b.dtype]), + ptr_arg(bias, fx_dtype[bias.dtype]), + ptr_arg(workspace, fx_dtype[workspace.dtype]), + ptr_arg(semaphore, fx_dtype[semaphore.dtype]), + m, + n, + k, + k, + stream, + ) + _compile_executable_to_cache(launch, *dispatch_args) + + def _compile_preshuffle_to_cache( *, m: int, @@ -679,7 +786,9 @@ def compile_one_config( t0 = time.time() try: - tensor_context = nullcontext() if kind == "hgemm" else FakeTensorMode() + tensor_context = ( + nullcontext() if kind in ("hgemm", "a16w16_gfx1250") else FakeTensorMode() + ) with ( override_env("FLYDSL_GPU_ARCH", aot_arch), tensor_context, @@ -688,6 +797,10 @@ def compile_one_config( hgemm_kwargs = dict(kwargs) hgemm_kwargs["target_gfx"] = aot_arch _compile_hgemm_to_cache(m=m, n=n, k=k, **hgemm_kwargs) + elif kind == "a16w16_gfx1250": + a16w16_kwargs = dict(kwargs) + a16w16_kwargs["target_gfx"] = aot_arch + _compile_a16w16_gfx1250_to_cache(m=m, n=n, k=k, **a16w16_kwargs) elif kind == "preshuffle": _compile_preshuffle_to_cache(m=m, n=n, k=k, **kwargs) elif kind == "8wave": diff --git a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv index 0bf96546bdb..744eaf4cd5d 100644 --- a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv +++ b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv @@ -138,7 +138,6 @@ gfx1250,256,8,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0 gfx1250,256,16,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -154,7 +153,6 @@ gfx1250,256,8,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0, gfx1250,256,16,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -170,7 +168,6 @@ gfx1250,256,8,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0, gfx1250,256,16,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -259,3 +256,6 @@ gfx1250,96,2048,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,trit gfx1250,96,4096,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,8192,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,16384,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 +gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,2,6.82,flydsl_hgemm_abf16_wbf16_bf16_t64x64x256x4_ks2_w4x2x1_bias1_ktail0_gm0_pht_gfx1250,0.0,276.75,4474.37 +gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,4,5.82,flydsl_hgemm_abf16_wbf16_bf16_t64x64x128x6_ks4_w4x2x1_bias1_ktail0_gm0_pht_gfx1250,0.0,259.44,4207.20 +gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,6,3.0,flydsl_hgemm_abf16_wbf16_bf16_t16x16x256x3_ks6_w1x1x1_bias1_ktail0_gm0_pht_gfx1250,0.0,15.73,374.10 diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py new file mode 100644 index 00000000000..8167e63e260 --- /dev/null +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper.""" + +from __future__ import annotations + +import functools + +import torch + +_compile_gemm_a16w16 = None +_run_compiled = None +_ptr_arg = None +_fx = None + + +def _lazy_import(): + global _compile_gemm_a16w16, _run_compiled, _ptr_arg, _fx + if _compile_gemm_a16w16 is not None: + return + import flydsl.expr as fx + + from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( + compile_gemm_a16w16, + ) + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled as run_compiled + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + _compile_gemm_a16w16 = compile_gemm_a16w16 + _run_compiled = run_compiled + _ptr_arg = ptr_arg + _fx = fx + + +_FX_DTYPE = {} + + +def _p(t): + if not _FX_DTYPE: + _FX_DTYPE.update( + { + torch.bfloat16: _fx.BFloat16, + torch.float16: _fx.Float16, + torch.float32: _fx.Float32, + torch.int32: _fx.Int32, + } + ) + return _ptr_arg(t, _FX_DTYPE[t.dtype]) + + +_CFG_KEYS = ( + "N", + "K", + "tile_m", + "tile_n", + "tile_k", + "m_warp", + "n_warp", + "in_dtype", + "out_dtype", + "num_buffers", + "waves_per_eu", + "activation", + "add_bias", + "physical_mk", + "physical_kn", + "kernarg_preload", + "split_k", + "sched_strategy", + "main_loop_unroll", + "variant", +) + + +@functools.lru_cache(maxsize=1024) +def _cached_launcher(*cfg): + return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) + + +_SPLIT_K_MAX_TILES = 4096 + + +@functools.cache +def _split_k_counters(device, stream): + return torch.zeros(_SPLIT_K_MAX_TILES, dtype=torch.int32, device=device) + + +def gemm_a16w16( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + dtype: torch.dtype = torch.float16, + y: torch.Tensor | None = None, + activation: str | None = None, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + num_buffers: int = 2, + waves_per_eu: int | None = None, + kernarg_preload: bool = False, + split_k: int = 1, + sched_strategy: str | None = None, + main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", +): + """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" + _lazy_import() + _half = (torch.float16, torch.bfloat16) + assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" + assert x.shape[1] == w.shape[1], "Incompatible K dimensions" + + M, K, N = x.shape[0], x.shape[1], w.shape[0] + assert ( + x.stride(1) == 1 or x.stride(0) == 1 or 1 in (M, K) + ), f"gemm_a16w16: x needs a unit-stride dim for TDM, got strides {tuple(x.stride())}" + assert ( + w.stride(1) == 1 or w.stride(0) == 1 or 1 in (N, K) + ), f"gemm_a16w16: w needs a unit-stride dim for TDM, got strides {tuple(w.stride())}" + physical_mk = x.stride(1) == 1 or K == 1 + physical_kn = w.stride(1) != 1 and N > 1 + + in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") + + if y is not None: + assert ( + y.shape[0] == M and y.shape[1] == N + ), f"y must be ({M}, {N}), got {tuple(y.shape)}" + assert y.stride(1) == 1 or 1 in ( + M, + N, + ), f"gemm_a16w16: y needs unit column stride, got strides {tuple(y.stride())}" + y_buf = y + else: + y_buf = torch.empty((M, N), device=x.device, dtype=dtype) + ldy = y_buf.stride(0) if M > 1 else N + assert ldy >= N, f"gemm_a16w16: y row stride {ldy} < N {N}" + + if bias is None: + bias = torch.empty(0, device=x.device, dtype=dtype) + + stream = torch.cuda.current_stream(device=x.device) + sem = _split_k_counters(x.device, stream) + if split_k > 1: + tiles = ((M + tile_m - 1) // tile_m) * ((N + tile_n - 1) // tile_n) + if tiles > _SPLIT_K_MAX_TILES: + raise RuntimeError( + f"[FlyDSL gfx1250] split_k needs {tiles} tile counters, " + f"more than {_SPLIT_K_MAX_TILES}" + ) + m_pad = ((M + tile_m - 1) // tile_m) * tile_m + n_pad = ((N + tile_n - 1) // tile_n) * tile_n + ws = torch.empty((split_k, m_pad, n_pad), device=x.device, dtype=torch.float32) + else: + ws = y_buf + + launch_fn = _cached_launcher( + N, + K, + tile_m, + tile_n, + tile_k, + m_warp, + n_warp, + in_dtype_str, + out_dtype_str, + num_buffers, + waves_per_eu, + activation, + bias.numel() > 0, + physical_mk, + physical_kn, + kernarg_preload, + split_k, + sched_strategy, + main_loop_unroll, + variant, + ) + + if physical_mk: + lda = x.stride(0) if M > 1 else K + else: + lda = x.stride(1) if K > 1 else M + if physical_kn: + ldb = w.stride(1) if K > 1 else N + else: + ldb = w.stride(0) if N > 1 else K + _run_compiled( + launch_fn, + _p(y_buf), + _p(x), + _p(w), + _p(bias), + _p(ws), + _p(sem), + M, + ldy, + lda, + ldb, + _fx.Stream(stream.cuda_stream), + ) + return y_buf + + +__all__ = ["gemm_a16w16"] diff --git a/aiter/ops/flydsl/gemm_kernels.py b/aiter/ops/flydsl/gemm_kernels.py index 9f10353d75a..97cc076443c 100644 --- a/aiter/ops/flydsl/gemm_kernels.py +++ b/aiter/ops/flydsl/gemm_kernels.py @@ -15,6 +15,7 @@ from aiter import logger from aiter.jit.utils.chip_info import get_gfx +from .gemm_a16w16_gfx1250 import gemm_a16w16 as gemm_a16w16_gfx1250 from .kernels.gemm_a16w16_gfx950 import ( SPLIT_K_SEMAPHORE_MAX_LEN, gemm_a16w16, @@ -138,10 +139,8 @@ def flydsl_hgemm( out_dtype: torch.dtype | None = None, stream: torch.cuda.Stream | None = None, ) -> torch.Tensor: - """Run the gfx950 A16W16 kernel with AITER's ``B[N, K]`` convention.""" + """Run the A16W16 kernel for this arch with AITER's ``B[N, K]`` convention.""" - if get_gfx() != "gfx950": - raise RuntimeError("The FlyDSL A16W16 kernel currently supports gfx950 only") if policy not in ("ft", "ht", "hti"): raise ValueError(f"Unsupported FlyDSL HGEMM policy: {policy!r}") launch_stream = ( @@ -150,6 +149,31 @@ def flydsl_hgemm( if launch_stream.device != a.device: raise ValueError(f"`stream` must be on {a.device}, got {launch_stream.device}") + gfx = get_gfx() + if gfx == "gfx1250": + if k_waves != 1: + raise ValueError("The gfx1250 FlyDSL A16W16 kernel supports k_waves=1 only") + with torch.cuda.stream(launch_stream): + return gemm_a16w16_gfx1250( + a, + b, + bias=bias, + dtype=out_dtype or a.dtype, + y=out, + tile_m=block_m, + tile_n=block_n, + tile_k=block_k, + m_warp=m_waves, + n_warp=n_waves, + num_buffers=stages, + split_k=split_k, + main_loop_unroll=policy in ("ht", "hti"), + ) + if gfx != "gfx950": + raise RuntimeError( + "The FlyDSL A16W16 kernel currently supports gfx950 and gfx1250 only" + ) + if not a.is_contiguous(): a = a.contiguous() if not b.is_contiguous(): diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py new file mode 100644 index 00000000000..cfb9ab26e89 --- /dev/null +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -0,0 +1,810 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.protocol import dsl_size_of +from flydsl.expr import const_expr, gpu, idx2crd, range_constexpr, rocdl, tdm_ops +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import check_smem_capacity + +from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops + +WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 +WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 +_SCHED_ALLOW_SALU = 1 << 2 +CPOL_DEV = 0x10 +KERNARG_PRELOAD_COUNT = 8 + + +def _f32(x): + return fx.Float32(x) + + +def apply_activation_scalar(val, activation: str): + if activation == "relu": + zero = _f32(0.0) + return (val > zero).select(val, zero) + one, half = _f32(1.0), _f32(0.5) + if activation in ("silu", "silu_exp2"): + return val / (one + fx.exp(_f32(0.0) - val)) + if activation == "gelu": + return half * val * (one + fx.erf(val * (1.0 / math.sqrt(2.0)))) + if activation == "gelu_tanh": + two = _f32(2.0) + inner = math.sqrt(2.0 / math.pi) * (val + 0.044715 * val * val * val) + tanh_val = one - two / (one + fx.exp(two * inner)) + return half * val * (one + tanh_val) + return val + + +def compile_gemm_a16w16( + *, + M: int = 0, + N: int = 0, + K: int, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + in_dtype: str = "fp16", + out_dtype: str | None = None, + num_buffers: int = 2, + waves_per_eu: int | None = None, + activation: str | None = None, + add_bias: bool = False, + physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) + physical_kn: bool = False, # False=N-major (row-major W), True=K-major (W^T) + kernarg_preload: bool = False, + split_k: int = 1, + sched_strategy: str | None = None, + main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", +): + def _req(cond, msg): + if not cond: + raise ValueError(msg) + + if out_dtype is None: + out_dtype = "f16" if in_dtype == "fp16" else "bf16" + is_f16 = in_dtype == "fp16" + + _req(2 <= num_buffers <= 8, f"num_buffers must be in [2, 8], got {num_buffers}") + _req(variant in ("bandwidth_bound", "compute_bound"), f"bad variant {variant!r}") + _req(in_dtype in ("fp16", "bf16"), f"bad in_dtype {in_dtype!r}") + _req(out_dtype in ("f32", "f16", "bf16"), f"bad out_dtype {out_dtype!r}") + _req( + sched_strategy in (None, "max-ilp", "max-memory-clause"), + f"bad sched_strategy {sched_strategy!r}", + ) + _req(split_k >= 1, f"split_k must be >= 1, got {split_k}") + + elem_bytes = 2 + elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 + num_warps, block_threads = m_warp * n_warp, m_warp * n_warp * WAVE_SIZE + warp_tile_m, warp_tile_n = tile_m // m_warp, tile_n // n_warp + split_k_chunk = K // split_k + num_k_tiles = (split_k_chunk + tile_k - 1) // tile_k + k_rem = split_k_chunk % tile_k + + for _nm, _v, _d in ( + ("K", K, split_k), + ("tile_k", tile_k, WMMA_K), + ("tile_m", tile_m, WMMA_M), + ("tile_n", tile_n, WMMA_N), + ("warp_tile_m", warp_tile_m, WMMA_M), + ("warp_tile_n", warp_tile_n, WMMA_N), + ): + _req(_v % _d == 0, f"{_nm} ({_v}) must be a multiple of {_d}") + + _req( + tile_k & (tile_k - 1) == 0, f"tile_k must be a power of 2 for TDM, got {tile_k}" + ) + + _req(N > 0, "N must be > 0 at compile time") + if physical_kn: + _req(tile_n & (tile_n - 1) == 0, f"tile_n must be a power of 2, got {tile_n}") + if not physical_mk: + _req(tile_m & (tile_m - 1) == 0, f"tile_m must be a power of 2, got {tile_m}") + _req( + num_k_tiles >= num_buffers - 1, + f"num_k_tiles {num_k_tiles} < {num_buffers - 1} (K={K}, tile_k={tile_k})", + ) + + gpu_arch = str(get_hip_arch()) + assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" + + k_wmma_steps = tile_k // WMMA_K + _fx_elem = fx.Float16 if is_f16 else fx.BFloat16 + wmma_m_rep, wmma_n_rep = warp_tile_m // WMMA_M, warp_tile_n // WMMA_N + N_ACCS = n_accs = wmma_m_rep * wmma_n_rep + N_A_FRAGS, N_B_FRAGS = wmma_m_rep * k_wmma_steps, wmma_n_rep * k_wmma_steps + + if physical_mk: + a_tile_shape, a_pad_interval = (tile_m, tile_k), tile_k + a_imm_bytes, lds_a_stride = tile_k * elem_bytes, tile_k + LDS_PAD_A + lds_a_elems = tile_m * lds_a_stride + LDS_PAD_A + else: + a_tile_shape, a_pad_interval = (tile_k, tile_m), tile_m + a_imm_bytes, lds_a_stride = None, tile_m + LDS_PAD_A + lds_a_elems = tile_k * lds_a_stride + LDS_PAD_A + + if physical_kn: + b_tile_shape, b_pad_interval = (tile_k, tile_n), tile_n + b_imm_bytes, lds_b_stride = None, tile_n + LDS_PAD_B + lds_b_elems = tile_k * lds_b_stride + LDS_PAD_B + else: + b_tile_shape, b_pad_interval = (tile_n, tile_k), tile_k + b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B + lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B + + if split_k > 1: + + @fx.struct + class SharedStorage: + a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] + b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] + flag: fx.Array[fx.Int32, 1, 16] + + else: + + @fx.struct + class SharedStorage: + a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] + b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] + + check_smem_capacity(dsl_size_of(SharedStorage), gpu_arch) + + _TDMS_PER_TILE = 2 # A + B + n_edge = N % tile_n + + @flyc.kernel + def kernel_gemm_a16w16( + arg_y: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_bias: fx.Pointer, + arg_ws: fx.Pointer, + arg_sem: fx.Pointer, + i32_m: fx.Int32, + i32_ldy: fx.Int32, + i32_lda: fx.Int32, + i32_ldb: fx.Int32, + ): + rocdl.disable_xdl_arb_stall() + elem_ty = _fx_elem.ir_type + + tx, bx, by = gpu.thread_id("x"), gpu.block_id("x"), gpu.block_id("y") + blk_m, blk_n = fx.Uint64(bx) * tile_m, fx.Uint64(by) * tile_n + if const_expr(split_k > 1): + split_k_base = fx.Uint64(gpu.block_id("z")) * split_k_chunk + else: + split_k_base = fx.Uint64(0) + + layout_thr = fx.make_layout( + (m_warp, n_warp, 2, 16), (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1) + ) + thr_coord = idx2crd(tx, layout_thr) + wave_m_idx, wave_n_idx, lane_kgrp, lane16 = ( + fx.Uint64(fx.get(thr_coord, 0)), + fx.Uint64(fx.get(thr_coord, 1)), + fx.Uint64(fx.get(thr_coord, 2)), + fx.Uint64(fx.get(thr_coord, 3)), + ) + + warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n + m_idx, ld_y = fx.Uint64(i32_m), fx.Uint64(i32_ldy) + gY = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))), + max_size=False, + num_records_bytes=m_idx * ld_y * elem_bytes_d, + ) + if const_expr(split_k > 1): + tiles_n = (N + tile_n - 1) // tile_n + n_pad = tiles_n * tile_n + m_pad = ((m_idx + (tile_m - 1)) // tile_m) * tile_m + plane = fx.Int64(m_pad) * n_pad + + def _ws_view(plane_idx): + return fx.rocdl.make_buffer_tensor( + fx.Tensor( + fx.make_view( + fx.add_offset(arg_ws, fx.Int64(plane_idx) * plane), + fx.make_layout((1, 1), (1, 1)), + ) + ), + max_size=False, + num_records_bytes=plane * 4, + ) + + gWs_mine = _ws_view(gpu.block_id("z")) + gWs = [_ws_view(s) for s in range_constexpr(split_k)] + tile_lin = fx.Uint64(bx) * tiles_n + fx.Uint64(by) + warp_lin = wave_m_idx * n_warp + wave_n_idx + ws_lane_base = (tile_lin * num_warps + warp_lin) * ( + n_accs * WAVE_SIZE * 8 + ) + (lane_kgrp * 16 + lane16) * 8 + sem_addr = ( + fx.Int64(fx.ptrtoint(arg_sem)) + + (fx.Int64(bx) * tiles_n + fx.Int64(by)) * 4 + ) + + lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() + big_a_mem, big_b_mem = lds.a.ptr, lds.b.ptr + big_a_base_idx = fx.Uint64(fx.ptrtoint(big_a_mem)) + big_b_base_idx = fx.Uint64(fx.ptrtoint(big_b_mem)) + blk_m64, blk_n64 = fx.Int64(blk_m), fx.Int64(blk_n) + split_k_base64 = fx.Int64(split_k_base) + + def _mk_side( + ptr, base_off, shape, outer_stride, pad_interval, pad_amount, extents + ): + gt = fx.Tensor( + fx.make_view( + fx.add_offset(ptr, base_off), + fx.make_layout(shape, (outer_stride, 1)), + ) + ) + return gt, fx.rocdl.make_tdm_atom( + gt, + extents, + strides=[outer_stride, None], + num_warps=num_warps, + pad_interval=pad_interval, + pad_amount=pad_amount, + early_timeout=True, + ) + + m_oob = fx.Int32(m_idx - blk_m) + n_oob = fx.Int32(fx.Uint64(N) - blk_n) if const_expr(N > 0) else None + a_extents = [m_oob, None] if const_expr(physical_mk) else [None, m_oob] + b_extents = [None, n_oob] if const_expr(physical_kn) else [n_oob, None] + + lda64 = fx.Int64(fx.Uint64(i32_lda)) + ldb64 = fx.Int64(fx.Uint64(i32_ldb)) + if const_expr(physical_mk): + a_base_off = blk_m64 * lda64 + split_k_base64 + else: + a_base_off = split_k_base64 * lda64 + blk_m64 + a_stride_rt = fx.Int32(i32_lda) + gA, atom_a = _mk_side( + arg_x, + a_base_off, + a_tile_shape, + a_stride_rt, + a_pad_interval, + LDS_PAD_A, + a_extents, + ) + if const_expr(physical_kn): + b_base_off = split_k_base64 * ldb64 + blk_n64 + else: + b_base_off = blk_n64 * ldb64 + split_k_base64 + b_stride_rt = fx.Int32(i32_ldb) + gB, atom_b = _mk_side( + arg_w, + b_base_off, + b_tile_shape, + b_stride_rt, + b_pad_interval, + LDS_PAD_B, + b_extents, + ) + + def _imm64(k_tile, mul): + if const_expr(not isinstance(mul, int)): + return fx.Int64(k_tile) * mul + return fx.Int64(k_tile * mul) + + def _slot_off(buf_idx, slot_elems): + return fx.Uint64(buf_idx * slot_elems) + + def _lane_bases(warp_base, reps, lds_stride, transpose): + if const_expr(not transpose): + row_off = (warp_base + lane16) * lds_stride + k_off = lane_kgrp * 8 + return [ + row_off + rep * WMMA_M * lds_stride + k_off + for rep in range_constexpr(reps) + ] + k_off = (lane_kgrp * 8 + lane16 % 8) * lds_stride + grp_off = (lane16 // 8) * 8 + return [ + k_off + warp_base + rep * WMMA_M + grp_off + for rep in range_constexpr(reps) + ] + + a_imm_rt = ( + a_imm_bytes + if const_expr(a_imm_bytes is not None) + else lda64 * (tile_k * elem_bytes) + ) + b_imm_rt = ( + b_imm_bytes + if const_expr(b_imm_bytes is not None) + else ldb64 * (tile_k * elem_bytes) + ) + + # One spec per operand side; A and B differ only in these fields. + MEM, BASE, BASES, STRIDE, TR, ELEMS, REPS, ATOM, GT, SHAPE, IMM, KDIM = range( + 12 + ) + A_SIDE = ( + big_a_mem, + big_a_base_idx, + _lane_bases(warp_m_base, wmma_m_rep, lds_a_stride, not physical_mk), + lds_a_stride, + not physical_mk, + lds_a_elems, + wmma_m_rep, + atom_a, + gA, + a_tile_shape, + a_imm_rt, + 1 if physical_mk else 0, + ) + B_SIDE = ( + big_b_mem, + big_b_base_idx, + _lane_bases(warp_n_base, wmma_n_rep, lds_b_stride, physical_kn), + lds_b_stride, + physical_kn, + lds_b_elems, + wmma_n_rep, + atom_b, + gB, + b_tile_shape, + b_imm_rt, + 0 if physical_kn else 1, + ) + + def issue_tdm_loads(buf_idx, k_tile): + if const_expr(k_rem != 0): + k_left = fx.Int32(split_k_chunk) - fx.Int32(k_tile) * tile_k + for sd in (A_SIDE, B_SIDE): + atom = sd[ATOM] + if const_expr(k_rem != 0): + atom = fx.atom_set_value(atom, f"extent_{sd[KDIM]}", k_left) + fx.copy( + atom, + sd[GT], + fx.Tensor( + fx.make_view( + fx.add_offset( + sd[MEM], + _slot_off(buf_idx, sd[ELEMS]), + ), + fx.make_layout(sd[SHAPE], (sd[STRIDE], 1)), + ) + ), + imm_offset=_imm64(k_tile, sd[IMM]), + ) + + lds_tr_ptr_ty = fx.PointerType.get( + elem_ty=_fx_elem.ir_type, + address_space=fx.AddressSpace.Shared, + alignment=elem_bytes, + ) + + def _frag(sd, slot_off, ks, rep): + vec8_ty = T.vec(8, elem_ty) + halves = [] + for k_half in range_constexpr(2): + if const_expr(sd[TR]): + off = ( + sd[BASES][rep] + + slot_off + + (ks * WMMA_K + k_half * 16) * sd[STRIDE] + ) + addr = fx.Int32(sd[BASE] + off * elem_bytes) + ptr = fx.to_llvm_ptr(fx.inttoptr(lds_tr_ptr_ty, addr)) + halves.append(fx.Vector(rocdl.ds_load_tr16_b128(vec8_ty, ptr))) + else: + off = sd[BASES][rep] + slot_off + (ks * WMMA_K + k_half * 16) + halves.append( + fx.ptr_load(fx.add_offset(sd[MEM], off), result_type=vec8_ty) + ) + return halves[0].shuffle(halves[1], list(range(16))) + + def load_frags(sd, buf_idx): + off = _slot_off(buf_idx, sd[ELEMS]) + return [ + _frag(sd, off, ks, rep) + for ks in range_constexpr(k_wmma_steps) + for rep in range_constexpr(sd[REPS]) + ] + + wmma_atom = fx.make_mma_atom( + fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, _fx_elem, fx.Float32) + ) + + def _rmem_vec(v, n, dtype): + t = fx.make_rmem_tensor(n, dtype) + t.store(v if isinstance(v, fx.Vector) else fx.Vector(v)) + return t + + def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): + a_ts = [_rmem_vec(f, 16, _fx_elem) for f in a_frags] + b_ts = [_rmem_vec(f, 16, _fx_elem) for f in b_frags] + acc_ts = [_rmem_vec(a, 8, fx.Float32) for a in accs_in] + rotate = rotate_buf is not None + next_a = [None] * N_A_FRAGS if rotate else None + next_b = [None] * N_B_FRAGS if rotate else None + if const_expr(rotate): + a_off = _slot_off(rotate_buf, lds_a_elems) + b_off = _slot_off(rotate_buf, lds_b_elems) + for ks in range_constexpr(k_wmma_steps): + for wm in range_constexpr(wmma_m_rep): + a_f = a_ts[ks * wmma_m_rep + wm] + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + fx.gemm( + wmma_atom, + acc_ts[idx], + b_ts[ks * wmma_n_rep + wn], + a_f, + acc_ts[idx], + ) + if const_expr(rotate): + rocdl.sched_barrier(_SCHED_ALLOW_SALU) + next_a[ks * wmma_m_rep + wm] = _frag(A_SIDE, a_off, ks, wm) + if const_expr(rotate): + rocdl.sched_barrier(_SCHED_ALLOW_SALU) + for wn in range_constexpr(wmma_n_rep): + next_b[ks * wmma_n_rep + wn] = _frag(B_SIDE, b_off, ks, wn) + return [t.load() for t in acc_ts], next_a, next_b + + _half_out = out_dtype in ("f16", "bf16") + _out_num = ( + fx.Float16 if out_dtype == "f16" else fx.BFloat16 if _half_out else None + ) + + if const_expr(add_bias): + bias_lay = fx.make_layout(4, 1) + bias_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(arg_bias, fx.make_layout(N, 1)), + max_size=False, + num_records_bytes=N * elem_bytes, + ) + bias_tiles = fx.logical_divide(bias_buf, bias_lay) + bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), _fx_elem) + if const_expr(n_edge != 0): + bias_elems = fx.logical_divide(bias_buf, fx.make_layout(1, 1)) + bias1_atom = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), _fx_elem) + + def load_bias4(tile_idx): + if const_expr(n_edge == 0): + r = fx.make_rmem_tensor(bias_lay, _fx_elem) + fx.copy(bias_atom, fx.slice(bias_tiles, (None, tile_idx)), r) + return r.load() + elems = [] + for e in range_constexpr(4): + r1 = fx.make_rmem_tensor(1, _fx_elem) + fx.copy( + bias1_atom, + fx.slice(bias_elems, (None, tile_idx * 4 + e)), + r1, + ) + elems.append(fx.Vector(r1.load())[0]) + return fx.Vector.from_elements(elems, _fx_elem) + + def epilogue_stores(final_accs): + _acc_ty = _out_num if _half_out else fx.Float32 + st_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _acc_ty) + if const_expr(n_edge != 0): + st1_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy16b() + if _half_out + else fx.rocdl.BufferCopy32b() + ), + _acc_ty, + ) + if const_expr(split_k > 1): + ws_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy128b(CPOL_DEV), fx.Float32 + ) + if const_expr(add_bias): + bias_vecs = [] + for wn in range_constexpr(wmma_n_rep): + col_tile = (blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8) // 4 + elems = [] + for half in range_constexpr(2): + bv = load_bias4(col_tile + half) + for i in range_constexpr(4): + elems.append(bv[i].to(fx.Float32)) + bias_vecs.append( + fx.Vector.from_elements(elems, fx.Float32).ir_value() + ) + + def _vec4(acc, half): + return fx.Vector.from_elements( + [fx.Vector(acc)[half * 4 + vi] for vi in range_constexpr(4)], + fx.Float32, + ) + + def _ws_off(acc_idx): + return ws_lane_base + acc_idx * (WAVE_SIZE * 8) + + def store_partial(acc, acc_idx): + for half in range_constexpr(2): + fx.copy( + ws_atom, + _rmem_vec(_vec4(acc, half), 4, fx.Float32), + gWs_mine[None, _ws_off(acc_idx) + half * 4], + ) + + def issue_partial_loads(acc_idx): + """Issue this lane's loads of all split_k planes (no waits here).""" + frags = [] + for half in range_constexpr(2): + for sidx in range_constexpr(split_k): + r = fx.make_rmem_tensor(4, fx.Float32) + fx.copy( + ws_atom, gWs[sidx][None, _ws_off(acc_idx) + half * 4], r + ) + frags.append(r) + return frags + + def sum_partials(frags): + """Consume issue_partial_loads' fragments into one 8-wide fp32 vector.""" + elems = [] + for half in range_constexpr(2): + total = None + for sidx in range_constexpr(split_k): + v = fx.Vector(frags[half * split_k + sidx].load()) + total = v if sidx == 0 else total + v + for vi in range_constexpr(4): + elems.append(total[vi]) + return fx.Vector.from_elements(elems, fx.Float32) + + def store_final(acc, row, col_base, wn): + if const_expr(add_bias): + acc = acc + bias_vecs[wn] + if const_expr(activation is not None): + acc = fx.Vector.from_elements( + [ + apply_activation_scalar(fx.Vector(acc)[i], activation) + for i in range_constexpr(8) + ], + fx.Float32, + ).ir_value() + if const_expr(_half_out): + h_vec = fx.Vector(acc).to(_out_num) + c_off = row * ld_y + col_base + if const_expr(n_edge == 0): + fx.copy(st_atom, _rmem_vec(h_vec, 8, _out_num), gY[None, c_off]) + else: + n_left = fx.Int32(N) - fx.Int32(col_base) + if n_left >= fx.Int32(8): + fx.copy( + st_atom, _rmem_vec(h_vec, 8, _out_num), gY[None, c_off] + ) + if n_left < fx.Int32(8): + for e in range_constexpr(8): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [h_vec[e]], _out_num + ) + fx.copy( + st1_atom, + _rmem_vec(one_vec, 1, _out_num), + gY[None, c_off + e], + ) + else: + for half in range_constexpr(2): + vec4 = _vec4(acc, half) + col = col_base + half * 4 + if const_expr(n_edge == 0): + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * ld_y + col], + ) + else: + n_left = fx.Int32(N) - fx.Int32(col) + if n_left >= fx.Int32(4): + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * ld_y + col], + ) + if n_left < fx.Int32(4): + for e in range_constexpr(4): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [vec4[e]], fx.Float32 + ) + fx.copy( + st1_atom, + _rmem_vec(one_vec, 1, fx.Float32), + gY[None, row * ld_y + col + e], + ) + + def _coords(wm, wn): + row = blk_m + warp_m_base + wm * WMMA_M + lane16 + col_base = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 + return row, col_base + + if const_expr(split_k == 1): + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + row, col_base = _coords(wm, wn) + store_final(final_accs[wm * wmma_n_rep + wn], row, col_base, wn) + else: + for acc_idx in range_constexpr(n_accs): + store_partial(final_accs[acc_idx], acc_idx) + comm_ops.waitcnt_all() + gpu.barrier() + if fx.Int32(tx) == fx.Int32(0): + arrival = fx.Int32(comm_ops.atomic_add_agent(sem_addr, fx.Int32(1))) + fx.ptr_store( + (arrival == fx.Int32(split_k - 1)).select( + fx.Int32(1), fx.Int32(0) + ), + lds.flag.ptr, + ) + gpu.barrier() + if fx.ptr_load(lds.flag.ptr) != fx.Int32(0): + pending = [issue_partial_loads(i) for i in range_constexpr(n_accs)] + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + row, col_base = _coords(wm, wn) + total = sum_partials(pending[wm * wmma_n_rep + wn]) + store_final(total, row, col_base, wn) + if fx.Int32(tx) == fx.Int32(0): + comm_ops.atomic_add_agent(sem_addr, fx.Int32(-split_k)) + + def _pack_state(accs_, a_, b_): + return list(accs_) + list(a_) + list(b_) + + def _unpack_state(state): + out, i = [], 0 + for n in (N_ACCS, N_A_FRAGS, N_B_FRAGS): + out.append(list(state[i : i + n])) + i += n + return out + + def _run_tile(accs_, ca, cb, cidx, lidx): + issue_tdm_loads(lidx % num_buffers, lidx) + tdm_ops.tensor_wait( + 0 if num_buffers == 2 else (num_buffers - 2) * _TDMS_PER_TILE + ) + gpu.barrier() + next_buf_i32 = (cidx + 1) % num_buffers + if const_expr(variant == "compute_bound"): + return wmma_tile(accs_, ca, cb, rotate_buf=next_buf_i32) + na = load_frags(A_SIDE, next_buf_i32) + nb = load_frags(B_SIDE, next_buf_i32) + accs_, _, _ = wmma_tile(accs_, ca, cb) + return accs_, na, nb + + acc_zero = fx.Vector.filled(8, fx.Float32(0.0), fx.Float32) + accs = [acc_zero] * n_accs + + # Prologue: + for i in range_constexpr(num_buffers - 1): + issue_tdm_loads(i, i) + tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) + gpu.barrier() + cur_a, cur_b = load_frags(A_SIDE, 0), load_frags(B_SIDE, 0) + + main_loop_iters = num_k_tiles - (num_buffers - 1) + load_idx_init, compute_idx_init = fx.Uint32(num_buffers - 1), fx.Uint32(0) + + TILES_PER_TRIP = 2 if variant == "bandwidth_bound" and main_loop_unroll else 1 + num_trips = main_loop_iters // TILES_PER_TRIP + load_idx_s, compute_idx_s = load_idx_init, compute_idx_init + + if const_expr(num_trips > 0): + init_state = _pack_state(accs, cur_a, cur_b) + [ + load_idx_init, + compute_idx_init, + ] + results = init_state + for trip, state in range(0, num_trips, 1, init=init_state): + cidx, lidx = fx.Uint32(state[-1]), fx.Uint32(state[-2]) + p_accs, ca, cb = _unpack_state(state[:-2]) + for _sub in range_constexpr(TILES_PER_TRIP): + p_accs, ca, cb = _run_tile(p_accs, ca, cb, cidx, lidx) + cidx, lidx = cidx + 1, lidx + 1 + results = yield _pack_state(p_accs, ca, cb) + [lidx, cidx] + + accs, cur_a, cur_b = _unpack_state(results[:-2]) + # loop results come back untyped; re-wrap so % stays unsigned (remui) + load_idx_s = fx.Uint32(results[-2]) + compute_idx_s = fx.Uint32(results[-1]) + else: + accs = list(accs) + + if const_expr(main_loop_iters % TILES_PER_TRIP == 1): + accs, cur_a, cur_b = _run_tile( + accs, cur_a, cur_b, compute_idx_s, load_idx_s + ) + + # Drain + drain_count_d = ( + num_buffers - 1 + if main_loop_iters > 0 + else min(num_buffers - 1, num_k_tiles) + ) + drain_base = max(0, main_loop_iters) + + accs = list(accs) + drain_a, drain_b = cur_a, cur_b + for j in range_constexpr(drain_count_d): + tile_idx = drain_base + j + if const_expr(j < drain_count_d - 1): + next_tile = tile_idx + 1 + tdm_ops.tensor_wait( + max(0, num_k_tiles - 1 - next_tile) * _TDMS_PER_TILE + ) + gpu.barrier() + next_a = load_frags(A_SIDE, next_tile % num_buffers) + accs, _, _ = wmma_tile(accs, drain_a, drain_b) + next_b = load_frags(B_SIDE, next_tile % num_buffers) + drain_a, drain_b = next_a, next_b + else: + accs, _, _ = wmma_tile(accs, drain_a, drain_b) + + if const_expr(num_buffers > 2): + rocdl.sched_barrier(0) + epilogue_stores(accs) + + @flyc.jit + def launch_gemm_a16w16( + arg_y: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_bias: fx.Pointer, + arg_ws: fx.Pointer, + arg_sem: fx.Pointer, + i32_m: fx.Int32, + i32_ldy: fx.Int32, + i32_lda: fx.Int32, + i32_ldb: fx.Int32, + stream: fx.Stream, + ): + gx = (fx.Uint64(i32_m) + (tile_m - 1)) // tile_m + gy = (N + tile_n - 1) // tile_n + + wpe = int(waves_per_eu) if waves_per_eu is not None else 0 + launcher = kernel_gemm_a16w16( + arg_y, + arg_x, + arg_w, + arg_bias, + arg_ws, + arg_sem, + i32_m, + i32_ldy, + i32_lda, + i32_ldb, + value_attrs={ + "rocdl.flat_work_group_size": f"{block_threads},{block_threads}", + "rocdl.waves_per_eu": wpe if wpe >= 1 else None, + }, + ) + + launcher.launch( + grid=(gx, gy, split_k), block=(block_threads, 1, 1), stream=stream + ) + + _llvm_opts = { + "amdgpu-expert-scheduling-mode": True, + "unroll-threshold": 0, + "amdgpu-kernarg-preload": kernarg_preload, + "amdgpu-kernarg-preload-count": KERNARG_PRELOAD_COUNT, + } + if sched_strategy is not None: + _llvm_opts["amdgpu-sched-strategy"] = sched_strategy + launch_gemm_a16w16.compile_hints["llvm_options"] = _llvm_opts + + return launch_gemm_a16w16 + + +__all__ = ["compile_gemm_a16w16"] diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py new file mode 100644 index 00000000000..1f76199758f --- /dev/null +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import argparse +import itertools +from functools import partial + +import pandas as pd +import torch +import torch.nn.functional as F + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.gemm_kernels import flydsl_hgemm +from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 as triton_gemm_a16w16 +from aiter.test_common import benchmark, checkAllclose, run_perftest + +torch.set_default_device("cuda") + +SUPPORTED_GFX = ["gfx950", "gfx1250"] +LAYOUTS = ["TN", "TT", "NN", "NT"] +ACTIVATIONS = { + "gelu": F.gelu, + "gelu_tanh": partial(F.gelu, approximate="tanh"), + "silu": F.silu, + "silu_exp2": F.silu, + "relu": F.relu, +} +SENTINEL = 4096.0 + +DEFAULT_MNK = { + "gfx1250": [ + (64, 64, 64), + (256, 256, 256), + (32, 256, 128), + (256, 32, 128), + (128, 128, 1024), + (1024, 128, 128), + (100, 190, 256), + (129, 257, 512), + (128, 192, 48), + (65, 190, 1000), + (64, 5120, 2880), + (64, 2880, 4096), + (64, 128, 2880), + ], + "gfx950": [ + (8, 4096, 4096), + (64, 4096, 4096), + (256, 4096, 4096), + (1024, 1024, 1024), + (32, 384, 7168), + (8, 7168, 2048), + ], +} +DEFAULT_CONFIGS = { + "gfx1250": [ + (64, 5120, 2880, 64, 64, 256, 4, 2, 4, 2, 1), + (64, 2880, 4096, 64, 64, 128, 4, 2, 6, 4, 1), + (64, 128, 2880, 16, 16, 256, 1, 1, 3, 6, 1), + (128, 128, 512, 64, 64, 128, 2, 2, 2, 2, 0), + (128, 128, 512, 64, 64, 128, 2, 2, 2, 4, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 3, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 6, 0), + (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4, 0), + (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8, 0), + (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4, 0), + (64, 64, 576, 64, 64, 64, 2, 2, 2, 2, 0), + (64, 100, 512, 64, 64, 64, 2, 2, 2, 2, 0), + (100, 100, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (129, 257, 512, 128, 128, 32, 2, 4, 3, 1, 0), + (65, 190, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (250, 120, 512, 64, 64, 128, 2, 2, 2, 1, 0), + (33, 65, 256, 32, 32, 128, 2, 2, 3, 1, 0), + (64, 64, 1024, 32, 32, 128, 2, 2, 3, 1, 0), + (128, 128, 1024, 128, 128, 32, 2, 4, 3, 1, 1), + ], + "gfx950": [ + (64, 4096, 4096, 32, 32, 128, 2, 2, 6, 1, 0), + (128, 4096, 4096, 32, 64, 128, 1, 4, 4, 1, 0), + (256, 4096, 4096, 64, 64, 128, 4, 2, 4, 1, 0), + (512, 4096, 4096, 64, 128, 64, 2, 4, 5, 1, 0), + (1024, 4096, 4096, 128, 128, 64, 2, 4, 4, 1, 0), + (1024, 1024, 1024, 64, 64, 64, 1, 4, 4, 1, 0), + (2048, 2048, 2048, 128, 128, 64, 2, 4, 4, 1, 0), + (4096, 4096, 4096, 256, 256, 64, 2, 4, 2, 1, 1), + (8, 7168, 2048, 16, 16, 64, 1, 1, 8, 1, 0), + (32, 14336, 4096, 32, 64, 128, 2, 2, 5, 1, 0), + ], +} + + +def generate_inputs(m, n, k, dtype, layout="TN", bias=False): + if layout[0] == "T": + x = torch.randn(m, k, dtype=dtype) + else: + x = torch.randn(k, m, dtype=dtype).T + if layout[1] == "T": + w = torch.randn(k, n, dtype=dtype).T + else: + w = torch.randn(n, k, dtype=dtype) + b = torch.randn(n, dtype=dtype) if bias else None + return x, w, b + + +def run_torch(x, w, bias=None, activation=None, dtype=dtypes.bf16): + out = F.linear( + x.to(dtypes.fp32), + w.to(dtypes.fp32), + None if bias is None else bias.to(dtypes.fp32), + ) + if activation is not None: + out = ACTIVATIONS[activation](out) + return out.to(dtype) + + +def run_candidates(candidates, ref, m, n, k, in_bytes, out_bytes, msg): + flops = 2 * m * n * k + nbytes = (m * k + n * k) * in_bytes + m * n * out_bytes + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + err = checkAllclose( + ref.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=1e-2, + atol=1e-2, + msg=f"{name}: {msg}", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +@benchmark() +def test_gemm_a16w16(m, n, k, dtype, layout): + x, w, _ = generate_inputs(m, n, k, dtype, layout) + ref = run_torch(x, w, dtype=dtype) + candidates = { + "flydsl": lambda: flydsl_hgemm(x, w, out_dtype=dtype), + "triton": lambda: triton_gemm_a16w16(x, w, dtype=dtype), + } + return run_candidates( + candidates, ref, m, n, k, x.element_size(), ref.element_size(), "gemm a16w16" + ) + + +@benchmark() +def test_gemm_a16w16_fused(m, n, k, dtype, activation, bias): + from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 + + x, w, b = generate_inputs(m, n, k, dtype, bias=bias) + ref = run_torch(x, w, b, activation, dtype) + candidates = { + "flydsl": lambda: gemm_a16w16(x, w, bias=b, dtype=dtype, activation=activation), + "triton": lambda: triton_gemm_a16w16( + x, w, bias=b, dtype=dtype, activation=activation + ), + } + return run_candidates( + candidates, + ref, + m, + n, + k, + x.element_size(), + ref.element_size(), + f"gemm a16w16 + {activation}", + ) + + +@benchmark() +def test_gemm_a16w16_config( + m, + n, + k, + dtype, + otype, + block_m, + block_n, + block_k, + m_waves, + n_waves, + stages, + split_k, + unroll, +): + x, w, _ = generate_inputs(m, n, k, dtype) + ref = run_torch(x, w, dtype=otype) + parent = torch.full((m + block_m, n + block_n), SENTINEL, dtype=otype) + y = parent[:m, :n] + cfg = { + "out": y, + "out_dtype": otype, + "block_m": block_m, + "block_n": block_n, + "block_k": block_k, + "stages": stages, + "split_k": split_k, + "m_waves": m_waves, + "n_waves": n_waves, + "k_waves": 1, + "group_m": 0, + "policy": "ht" if unroll else "ft", + } + assert flydsl_hgemm(x, w, **cfg) is y + candidates = { + "flydsl": lambda: flydsl_hgemm(x, w, **cfg), + "triton": lambda: triton_gemm_a16w16(x, w, dtype=otype), + } + ret = run_candidates( + candidates, + ref, + m, + n, + k, + x.element_size(), + ref.element_size(), + "gemm a16w16 config", + ) + assert torch.all(parent[m:] == SENTINEL) and torch.all( + parent[:, n:] == SENTINEL + ), "flydsl wrote outside the [m, n] output view" + return ret + + +def summarize(name, rows): + aiter.logger.info( + "%s summary (markdown):\n%s", name, pd.DataFrame(rows).to_markdown(index=False) + ) + + +def main(): + gfx = get_gfx() + if gfx not in SUPPORTED_GFX: + aiter.logger.warning("flydsl a16w16 gemm unsupported on %s; skipping", gfx) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes["bf16"], dtypes.d_dtypes["fp16"]], + nargs="*", + default="bf16,", + metavar="{bf16,fp16}", + help="""Input dtype of x and w. + e.g.: -d bf16""", + ) + parser.add_argument( + "-o", + "--otype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes[t] for t in ("bf16", "fp16", "fp32")], + nargs="*", + default="bf16,fp32,", + metavar="{bf16,fp16,fp32}", + help="""Output dtype for the kernel-config sweep (-c): fp32 keeps split-K + lossless, bf16/fp16 take the f32-accumulate-then-cast path. + e.g.: -o fp32""", + ) + parser.add_argument( + "-s", + "--mnk", + type=dtypes.str2tuple, + nargs="*", + default=DEFAULT_MNK[gfx], + help="""Shape of mnk. + e.g.: -s 64,5120,2880""", + ) + parser.add_argument( + "-l", + "--layout", + type=str, + choices=LAYOUTS, + nargs="*", + default=LAYOUTS, + help="""(x, w) memory layout. x: T = row-major [m,k], N = transposed view; + w: N = row-major [n,k], T = transposed view. TN is the nn.Linear call. + e.g.: -l TN NT""", + ) + parser.add_argument( + "-a", + "--activation", + type=str, + choices=list(ACTIVATIONS), + nargs="*", + default=list(ACTIVATIONS), + help="""Fused epilogue activation (gfx1250 kernel only). + e.g.: -a gelu silu""", + ) + parser.add_argument( + "--bias", + type=dtypes.str2bool, + nargs="*", + default=[False, True], + help="""Fuse a bias add into the epilogue. + e.g.: --bias 1""", + ) + parser.add_argument( + "-c", + "--config", + type=dtypes.str2tuple, + nargs="*", + default=DEFAULT_CONFIGS[gfx], + help="""Explicit flydsl_hgemm config as + m,n,k,block_m,block_n,block_k,m_waves,n_waves,stages,split_k,unroll + (unroll=1 -> policy ht; the fields of a tuned flydsl_hgemm_* csv name). + e.g.: -c 64,5120,2880,64,64,256,4,2,4,2,1""", + ) + args = parser.parse_args() + + for dtype in args.dtype: + summarize( + "gemm_a16w16", + [ + test_gemm_a16w16(m, n, k, dtype, layout) + for layout, (m, n, k) in itertools.product(args.layout, args.mnk) + ], + ) + if gfx == "gfx1250": + summarize( + "gemm_a16w16 fused epilogue", + [ + test_gemm_a16w16_fused(m, n, k, dtype, act, bias) + for act, bias, (m, n, k) in itertools.product( + args.activation, args.bias, args.mnk + ) + ], + ) + else: + aiter.logger.warning( + "%s: skipping fused-epilogue table (activation fusion is gfx1250-only)", + gfx, + ) + summarize( + "gemm_a16w16 kernel config", + [ + test_gemm_a16w16_config(m, n, k, dtype, otype, *tile) + for otype, (m, n, k, *tile) in itertools.product( + args.otype, args.config + ) + ], + ) + + +if __name__ == "__main__": + main()