diff --git a/setup.py b/setup.py index 2f4ae06e0..ebe581cdc 100644 --- a/setup.py +++ b/setup.py @@ -156,6 +156,14 @@ def setup_requirements() -> Tuple[List[str], List[str]]: ] test_reqs: List[str] = ["pytest>=8.2.1"] + # Optional FlyDSL dependency for ROCm PyTorch builds. + if ( + rocm_build() + and "pytorch" in frameworks + and bool(int(os.getenv("NVTE_USE_FLYDSL", "0"))) + ): + install_reqs.extend(["flydsl"]) + # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: diff --git a/tests/pytorch/flydsl_kernels/test_gemm.py b/tests/pytorch/flydsl_kernels/test_gemm.py new file mode 100644 index 000000000..48510880f --- /dev/null +++ b/tests/pytorch/flydsl_kernels/test_gemm.py @@ -0,0 +1,551 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""User-facing FlyDSL GEMM tests -- ``general_gemm()`` under ``NVTE_USE_FLYDSL=1``. + +Exercises the same public entry point used by TE ``Linear`` / +``LayerNormLinear``. Coverage mirrors the Triton user-facing GEMM tests for the +currently supported FlyDSL surface: + +- fp32 / fp16 / bf16 regular tensors +- same-format and mixed-format tensor-wise FP8 +- same-format and mixed-format MXFP8 +- TN / NN / NT layouts +- batched multidimensional FP8 flattening + +Fused BIAS and BGRADB epilogues are intentionally not included yet because the +FlyDSL GEMM path does not currently support them. + +Each test compares the FlyDSL path against two independent references: + +1. ``torch.matmul`` on dequantized inputs, independent of hipBLASLt behavior. +2. The native C++ ``tex.generic_gemm`` backend through the same + ``general_gemm`` public surface. + +FlyDSL kernels currently require tile-aligned launch dimensions, so the test +shapes are aligned to the 256x256x128 kernel contract rather than reusing the +odd-sized Triton edge-mask cases. +""" + +import os + +import pytest +import torch + +from transformer_engine.pytorch import Float8Tensor +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import ( + MXFP8Quantizer, + MXFP8Tensor, +) +import transformer_engine_torch as tex + + +# --- Feature detection -------------------------------------------------------- + +major, minor = torch.cuda.get_device_capability() + +# The current FlyDSL MXFP8 implementation uses the gfx950 fp8-scaled MFMA. +has_mxfp8_support = major == 9 and minor >= 5 + +requires_mxfp8_support = pytest.mark.skipif( + not has_mxfp8_support, + reason="FlyDSL MXFP8 requires gfx950+ fp8-scaled MFMA support", +) + + +# --- Test parameters ---------------------------------------------------------- + +# The current FlyDSL kernels have no M/N edge masks and specialize K in K128 +# tiles. Keep all dimensions aligned to exercise the supported production path. +FLYDSL_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), + (1024, 512, 1024), +] + +MXFP8_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), +] + +LAYOUTS = ["TN", "NN", "NT"] + +FP8_FORMAT_COMBOS = [ + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E5M2), +] + +FP8_FORMAT_IDS = [ + "e4m3_e4m3", + "e4m3_e5m2", + "e5m2_e4m3", + "e5m2_e5m2", +] + +REGULAR_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + +# --- Fixtures ----------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def cleanup_env(): + """Save and restore FlyDSL-related environment variables between tests.""" + old_flydsl = os.environ.get("NVTE_USE_FLYDSL") + old_mxfp8 = os.environ.get("NVTE_ROCM_ENABLE_MXFP8") + + yield + + if old_flydsl is None: + os.environ.pop("NVTE_USE_FLYDSL", None) + else: + os.environ["NVTE_USE_FLYDSL"] = old_flydsl + + if old_mxfp8 is None: + os.environ.pop("NVTE_ROCM_ENABLE_MXFP8", None) + else: + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = old_mxfp8 + + +# --- Helpers ------------------------------------------------------------------ + +def get_shapes(layout, M, K, N): + """Return the A/B storage shapes used by TE's public GEMM tests.""" + if layout == "TN": + return (M, K), (N, K) + if layout == "NN": + return (M, K), (K, M) + if layout == "NT": + return (M, K), (M, K) + raise ValueError(f"Unsupported layout: {layout}") + + +def compute_pytorch_reference(A_ref, B_ref, layout): + """Compute the equivalent public-layout GEMM with ``torch.matmul``.""" + if layout == "TN": + return torch.matmul(B_ref, A_ref.T) + if layout == "NN": + return torch.matmul(B_ref, A_ref) + if layout == "NT": + return torch.matmul(B_ref.T, A_ref) + raise ValueError(f"Unsupported layout: {layout}") + + +def create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b): + """Create independently typed Float8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + return A_fp8, B_fp8, A_fp8.dequantize(), B_fp8.dequantize() + + +def _make_mxfp8_quantizer(fp8_dtype): + """Create one independently typed MXFP8 quantizer with both orientations.""" + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype) + quantizer.set_usage(rowwise=True, columnwise=True) + return quantizer + + +def create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, +): + """Create independently typed MXFP8Tensor inputs and references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device="cuda") * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device="cuda") * 0.5 + + A_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_a)(A_f32) + B_mxfp8 = _make_mxfp8_quantizer(fp8_dtype_b)(B_f32) + + return ( + A_mxfp8, + B_mxfp8, + A_mxfp8.dequantize(), + B_mxfp8.dequantize(), + ) + + +def call_gemm(A, B, layout, out_dtype, use_flydsl=True): + """Call ``general_gemm`` through either FlyDSL or the native C++ path.""" + os.environ["NVTE_USE_FLYDSL"] = "1" if use_flydsl else "0" + + output, bias_grad, gelu_input, extra_output = general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=None, + quantization_params=None, + gelu=False, + grad=False, + accumulate=False, + ) + + assert bias_grad is None + assert gelu_input is None + assert extra_output is None + return output + + +def assert_gemm_close(actual, expected, *, atol, rtol): + """Compare through FP32 so output narrowing does not hide diagnostics.""" + torch.testing.assert_close( + actual.float(), + expected.float(), + atol=atol, + rtol=rtol, + equal_nan=False, + ) + + +# ============================================================================== +# Approach 1: FlyDSL vs PyTorch torch.matmul reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_pytorch_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against an FP32 PyTorch reference.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + output = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + expected = compute_pytorch_reference(A.float(), B.float(), layout) + + assert_gemm_close(output, expected, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format tensor-wise FP8 FlyDSL GEMMs.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, A_deq, B_deq = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 FlyDSL GEMMs.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, A_deq, B_deq = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + output = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + expected = compute_pytorch_reference( + A_deq.float(), + B_deq.float(), + layout, + ) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "dtype", + REGULAR_DTYPES, + ids=["fp32", "fp16", "bf16"], +) +def test_flydsl_vs_cpp_regular(M, K, N, layout, dtype): + """Test regular FlyDSL GEMM against the native C++ backend.""" + torch.manual_seed(42) + + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device="cuda") * 0.5 + B = torch.randn(B_shape, dtype=dtype, device="cuda") * 0.5 + + flydsl_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=True, + ) + cpp_out = call_gemm( + A, + B, + layout, + out_dtype=dtype, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=1e-3, rtol=1e-2) + + +@pytest.mark.parametrize("M, K, N", FLYDSL_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_fp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format FP8 against native C++.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, _, _ = create_fp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_fp8, + B_fp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +@requires_mxfp8_support +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_cpp_mxfp8(M, K, N, layout, fp8_format): + """Test same-format and mixed-format MXFP8 against native C++.""" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + A_mxfp8, B_mxfp8, _, _ = create_mxfp8_tensors( + M, + K, + N, + layout, + fp8_dtype_a, + fp8_dtype_b, + ) + + flydsl_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=True, + ) + cpp_out = call_gemm( + A_mxfp8, + B_mxfp8, + layout, + out_dtype=torch.float32, + use_flydsl=False, + ) + + assert_gemm_close(flydsl_out, cpp_out, atol=5e-3, rtol=1e-2) + + +# ============================================================================== +# Batched multidimensional FP8 coverage +# ============================================================================== + +@pytest.mark.parametrize( + "batch_size, M, K, N", + [ + (2, 256, 512, 256), + (4, 256, 512, 256), + ], +) +@pytest.mark.parametrize( + "fp8_format", + FP8_FORMAT_COMBOS, + ids=FP8_FORMAT_IDS, +) +def test_flydsl_vs_pytorch_fp8_multidim( + batch_size, + M, + K, + N, + fp8_format, +): + """Exercise flatten-leading-dim semantics for multidimensional FP8.""" + torch.manual_seed(42) + + fp8_dtype_a, fp8_dtype_b = fp8_format + + # TN layout: the wrapper flattens all leading dimensions into rows. + A_f32 = ( + torch.randn( + batch_size, + M, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + B_f32 = ( + torch.randn( + batch_size, + N, + K, + dtype=torch.float32, + device="cuda", + ) + * 0.5 + ) + + A_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full((1,), 1.0, dtype=torch.float32, device="cuda"), + amax=torch.empty((1,), dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + output = call_gemm( + A_fp8, + B_fp8, + layout="TN", + out_dtype=torch.float32, + use_flydsl=True, + ) + + A_flat = A_fp8.dequantize().reshape(-1, K) + B_flat = B_fp8.dequantize().reshape(-1, K) + expected = torch.matmul(B_flat, A_flat.T) + + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +if __name__ == "__main__": + # Quick smoke tests using one case from each supported input family. + os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_ROCM_ENABLE_MXFP8"] = "1" + + test_flydsl_vs_pytorch_regular( + 256, + 512, + 256, + "TN", + torch.float16, + ) + test_flydsl_vs_pytorch_fp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + ) + + if has_mxfp8_support: + test_flydsl_vs_pytorch_mxfp8( + 256, + 512, + 256, + "TN", + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), + ) + + print("All FlyDSL GEMM smoke tests passed!") diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5c9686f15..c9d9a5563 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1344,6 +1344,148 @@ def test_linear_accuracy(dtype, bs, model, return_bias, bias): assert_allclose(te_output, torch_output, tolerance, rtol[dtype]) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["small", "126m"]) +@pytest.mark.parametrize( + "fp8_recipe", + [ + None, + recipe.Float8CurrentScaling(), + recipe.DelayedScaling(), + recipe.MXFP8BlockScaling(), + ], +) +def test_linear_accuracy_flydsl( + dtype, + bs, + model, + fp8_recipe, +): + """Compare FlyDSL and native TE Linear forward, dgrad, and wgrad.""" + + if not IS_HIP_EXTENSION: + pytest.skip("FlyDSL GEMM is only supported on HIP.") + + fp8 = fp8_recipe is not None + config = model_configs[model] + + if isinstance(fp8_recipe, recipe.MXFP8BlockScaling): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif fp8 and not fp8_available: + pytest.skip(reason_for_no_fp8) + + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + # Validate the GEMM backend, not quantized parameter storage. + # FlyDSL GEMM does not currently support bias. + with quantized_model_init(enabled=False, recipe=fp8_recipe): + linear_ref = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + linear_flydsl = Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + device="cuda", + ).eval() + + with torch.no_grad(): + linear_flydsl.weight.copy_(linear_ref.weight) + + input_shape = ( + config.max_seqlen_q, + bs, + config.hidden_size, + ) + + inp_ref = torch.randn( + input_shape, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + inp_flydsl = inp_ref.detach().clone().requires_grad_(True) + + try: + # Native TE backend. + os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_ref = linear_ref(inp_ref) + + out_ref.sum().backward() + torch.cuda.synchronize() + + # FlyDSL backend. + os.environ["NVTE_USE_FLYDSL"] = "1" + os.environ["NVTE_FLYDSL_GEMM_WARN_FALLBACK"] = "1" + + reset_rng_states() + FP8GlobalStateManager.reset() + + with autocast(enabled=fp8, recipe=fp8_recipe): + out_flydsl = linear_flydsl(inp_flydsl) + + out_flydsl.sum().backward() + torch.cuda.synchronize() + + finally: + os.environ.pop("NVTE_USE_FLYDSL", None) + os.environ.pop("NVTE_FLYDSL_GEMM_WARN_FALLBACK", None) + FP8GlobalStateManager.reset() + + atol, rtol = get_tolerances(dtype) + + if fp8: + atol = max(atol, 1e-2) + rtol = max(rtol, 1e-2) + + torch.testing.assert_close( + out_flydsl, + out_ref, + atol=atol, + rtol=rtol, + ) + + torch.testing.assert_close( + inp_flydsl.grad, + inp_ref.grad, + atol=atol, + rtol=rtol, + ) + + # Wgrad is an NT GEMM with a reduction over the flattened + # sequence/batch dimension. FlyDSL and the native TE backend may use + # different FP32 accumulation orders, so allow the small expected + # non-associative rounding difference. + wgrad_atol = atol + wgrad_rtol = rtol + + if dtype == torch.float32 and not fp8: + wgrad_atol = max(wgrad_atol, 1e-4) + wgrad_rtol = max(wgrad_rtol, 1e-4) + + torch.testing.assert_close( + linear_flydsl.weight.grad, + linear_ref.weight.grad, + atol=wgrad_atol, + rtol=wgrad_rtol, + ) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["small"]) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 3e787f820..52ca77379 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -10,6 +10,7 @@ import ctypes import os import functools +import warnings import torch from torch.utils.cpp_extension import IS_HIP_EXTENSION import transformer_engine_torch as tex @@ -460,7 +461,46 @@ def general_gemm( "beta": beta, } - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + use_gemm_flydsl = ( + IS_HIP_EXTENSION + and bool(int(os.environ.get("NVTE_USE_FLYDSL", "0"))) + ) + + if use_gemm_flydsl: + # Lazy import keeps FlyDSL off the normal Transformer Engine import path. + from ..flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, + ) + + try: + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( + *args, + **kwargs, + ) + except FlyDSLUnsupportedError as exc: + warn_fallback = os.environ.get( + "NVTE_FLYDSL_GEMM_WARN_FALLBACK", + "0", + ).lower() not in ("", "0", "false", "no", "off") + + if warn_fallback: + warnings.warn( + "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; " + f"falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=2, + ) + + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + *args, + **kwargs, + ) + else: + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + *args, + **kwargs, + ) if IS_HIP_EXTENSION and use_bf16_tn_output_workaround: out = cast_if_needed(out, torch.float32) diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py new file mode 100644 index 000000000..92fa250e8 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -0,0 +1,3 @@ +from . import gemm + +__all__ = ["gemm"] \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py new file mode 100644 index 000000000..5acdce6a2 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" + +from .exceptions import FlyDSLUnsupportedError +from .gemm_wrappers import te_generic_gemm_flydsl + +__all__ = [ + "FlyDSLUnsupportedError", + "te_generic_gemm_flydsl", +] \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py new file mode 100644 index 000000000..4201d571d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/bf16_gemm.py @@ -0,0 +1,1132 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL BF16 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K64 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as BF16 tensors shaped [M, K] and [N, K], and writes FP16, +BF16, or FP32 C shaped [M, N]. The public ``bf16_matmul`` entry point accepts +Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 2 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "bf16_pp_smem_a0" +LDS_SYM_A1 = "bf16_pp_smem_a1" +LDS_SYM_B0 = "bf16_pp_smem_b0" +LDS_SYM_B1 = "bf16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized 4-wave kernel for compile-time ``K`` and output dtype. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 2 + VEC_BYTES = 16 + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL BF16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 BF16 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_bf16_byte_buffer_tensor(A) + gB = make_bf16_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x64 half-page (16 KiB). Each half has its own LDS base. + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K * ELEM_BYTES)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def _bf16_k32_frag(full_frag, k32): + # A/B 16x64 BF16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight BF16 values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, + ) + + def _pinned_bf16_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + ( + f"v_mfma_f32_16x16x32_bf16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 BF16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_bf16_mfma_once( + acc_idx, + _bf16_k32_frag(a_frag, k32), + _bf16_k32_frag(b_frag, k32), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _bf16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _bf16_k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _bf16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_bf16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _bf16_k32_frag(b_frags[ni], k32), + ) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + +def bf16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN BF16 GEMM adapter. + + Public/backend contract: + a: [M, K] BF16 + b: [K, N] BF16 + c: [M, N] FP16, BF16, or FP32 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL BF16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: + raise TypeError( + "FlyDSL BF16 GEMM expects both operands to have torch.bfloat16 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + raise TypeError( + "FlyDSL BF16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL BF16 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized BF16 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16 + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL BF16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch( + int(K_runtime), + C.dtype, + bool(use_xcd_remap), + ) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py new file mode 100644 index 000000000..1ae38569a --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/exceptions.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +class FlyDSLUnsupportedError(RuntimeError): + """The GEMM request is valid but unsupported by the available FlyDSL kernels.""" \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py new file mode 100644 index 000000000..709f76484 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm.py @@ -0,0 +1,1139 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP16 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K64 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP16 tensors shaped [M, K] and [N, K], and writes FP16, +BF16, or FP32 C shaped [M, N]. The public ``fp16_matmul`` entry point accepts +Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor as make_fp16_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 64 + +# Public metadata consumed by wrappers. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 2 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp16_pp_smem_a0" +LDS_SYM_A1 = "fp16_pp_smem_a1" +LDS_SYM_B0 = "fp16_pp_smem_b0" +LDS_SYM_B1 = "fp16_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 64 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def make_fp16_inputs(M, N, K, device="cuda"): + """Generate FP16 A[M,K] and B[N,K] inputs.""" + A = (torch.randn(M, K, device=device) * 0.5).to(torch.float16) + B = (torch.randn(N, K, device=device) * 0.5).to(torch.float16) + return A, B + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K64 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 2 + VEC_BYTES = 16 + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {output_dtype}" + ) + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K64 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 FP16 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_fp16_byte_buffer_tensor(A) + gB = make_fp16_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K32 slices x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(8) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K32 slices for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(16) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x64 half-page (16 KiB). Each half has its own LDS base. + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K * ELEM_BYTES)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def _fp16_k32_frag(full_frag, k32): + # A/B 16x64 FP16 wave fragments are i32x8. Each K32 MFMA + # consumes one contiguous i32x4 slice (eight FP16 values/lane). + lo = k32 * 4 + v = Vec(full_frag) + return Vec.from_elements( + [v[lo], v[lo + 1], v[lo + 2], v[lo + 3]], + fx.Int32, + ) + + def _pinned_fp16_mfma_once(acc_idx, a_k32, b_k32): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k32), arith._to_raw(b_k32)], + ( + f"v_mfma_f32_16x16x32_f16 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x64 FP16 product into pinned AGPRs.""" + for k32 in range_constexpr(2): + _pinned_fp16_mfma_once( + acc_idx, + _fp16_k32_frag(a_frag, k32), + _fp16_k32_frag(b_frag, k32), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K64 update is two in-place K32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32, a0, a1, a2, a3, b0, b1): + """Issue one K32 slice for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for mi in range_constexpr(4): + a_k32 = _fp16_k32_frag(a_frags[mi], k32) + for nj in range_constexpr(2): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k32, + _fp16_k32_frag(b_frags[nj], k32), + ) + + def mfma_4n_4mi_k32(subtile_id, k32, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue one K32 slice for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for mi in range_constexpr(4): + a_k32 = _fp16_k32_frag(a_frags[mi], k32) + for ni in range_constexpr(4): + _pinned_fp16_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k32, + _fp16_k32_frag(b_frags[ni], k32), + ) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if const_expr(output_dtype != torch.float32): + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Compute is K32-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:32]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[32:64]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:32]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[32:64]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[32:64]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:32]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[32:64]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K64 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + +def fp16_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN FP16 GEMM adapter. + + Public/backend contract: + a: [M, K] FP16 + b: [K, N] FP16 + c: [M, N] FP16, BF16, or FP32 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP16 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.float16 or b.dtype != torch.float16: + raise TypeError( + "FlyDSL FP16 GEMM expects both operands to have torch.float16 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + raise TypeError( + "FlyDSL FP16 GEMM output dtype must be torch.float16, " + f"torch.bfloat16, or torch.float32, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got {a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP16 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized FP16 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.float16 and B.dtype == torch.float16 + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP16 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch( + int(K_runtime), + C.dtype, + bool(use_xcd_remap), + ) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py new file mode 100644 index 000000000..5aaeab3b1 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp16_gemm_utils.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors +"""Minimal byte-staging helpers for the first-pass BF16 four-wave GEMM.""" + +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr + +# ceildiv is the canonical cdiv from the shared layer +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + +def divmod(a, b): + return (a // b, a % b) + + +def swizzle_128(row, col_in_bytes): + """HK 128-byte row XOR swizzle; ``col_in_bytes`` is a byte coordinate.""" + offset = row * 128 + col_in_bytes + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def make_bf16_byte_buffer_tensor(arg_u8): + """Create a byte-addressed buffer tensor from a contiguous BF16 uint8 view.""" + return fx.rocdl.make_buffer_tensor(arg_u8, max_size=False) + + +def compute_global_swizzle(lane_id, wave_id, row_stride_bytes, n_rounds, preshuffled=False): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + raise AssertionError("BF16 first-pass port does not support preshuffled operands") + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col_bytes = (lane_id % 8) * 16 + r, c = swizzle_128(row, col_bytes) + offsets.append(r * row_stride_bytes + c) + return offsets + + +class G2SLoader: + """Issue raw 16-byte buffer-to-LDS copies. + + Both the global source and LDS destination must be byte-addressed. Fly's copy lowering does not legalize an i8 buffer source paired with a bf16 LDS + destination even when the transfer width is the same 128 bits. + """ + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + lds_ptr = fx.inttoptr(self.LdsPtr_t, base_i32 + fx.Int32(step_off)) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, byte_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + fx.copy(self.g2lds_atom, src, self._lds_dst_at(lds_dst, step), soffset=fx.Int32(byte_offset)) + + def load_one(self, lds_dst, byte_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + fx.copy(self.g2lds_atom, src, self._lds_dst_at(lds_dst, step), soffset=fx.Int32(byte_offset)) + + +def pack_i32x4_i32x8(lo, hi): + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + """Raw 16-byte LDS reader used to assemble an i32x8 BF16 K64 fragment.""" + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16bytes(self, lds_src, offset): + ptr_off = fx.add_offset(lds_src.ptr, fx.make_int_tuple(offset)) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + return fx.make_view(i8_iter, fx.make_layout(16, 1)).load() + + def load_one(self, lds_src, lds_offset): + return self._vec_load_16bytes(lds_src, lds_offset).bitcast(fx.Int32) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py new file mode 100644 index 000000000..6cbd61102 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp32_gemm.py @@ -0,0 +1,1091 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL FP32 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K32 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes A and B as FP32 tensors shaped [M, K] and [N, K], and writes FP32 C +shaped [M, N]. The public ``fp32_matmul`` entry point accepts Transformer +Engine's TN contract and performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp16_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_bf16_byte_buffer_tensor as make_fp32_byte_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 32 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 4 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp32_pp_smem_a0" +LDS_SYM_A1 = "fp32_pp_smem_a1" +LDS_SYM_B0 = "fp32_pp_smem_b0" +LDS_SYM_B1 = "fp32_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 32 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def make_fp32_inputs(M, N, K, device="cuda"): + """Generate FP32 A[M,K] and B[N,K] inputs.""" + A = (torch.randn(M, K, device=device) * 0.5).to(torch.float32) + B = (torch.randn(N, K, device=device) * 0.5).to(torch.float32) + return A, B + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel(K: int, use_xcd_remap: bool = True): + """Build the specialized 4-wave kernel for compile-time ``K``. + + ``K`` must contain at least four K32 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 4 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K32 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LDS_BYTES_HALF = LDS_ELEMS_HALF * ELEM_BYTES + LOAD_PASSES_HALF = LDS_BYTES_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x64 FP32 page is two independent 128x64 half-pages. + # Store LDS as bytes so BufferCopyLDS128b sees i8 on both source and + # destination. Each half-page remains exactly 16 KiB. + a0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + a1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b0_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_0: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + b1_1: fx.Array[fx.Uint8, LDS_BYTES_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + # A/B arrive as contiguous uint8 byte views. Keeping staging byte-addressed + # preserves the original 16-byte G2L instruction cadence and vmcnt values. + gA = make_fp32_byte_buffer_tensor(A) + gB = make_fp32_byte_buffer_tensor(B) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K * ELEM_BYTES, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, fx.Uint8.ir_type, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(4) # C is FP32. + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + # Eight refill VMEM operations overlap four independent 8-MFMA + # groups (two K16 halves x two N-halves). + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_mfma(32) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Eight refill VMEM operations and eight distributed A-bottom LDS + # reads overlap four independent 8-MFMA K32 groups. + for _ in range_constexpr(4): + rocdl.sched_vmem(2) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(32) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Eight two-read prefetch groups overlap four complete-quadrant + # 16-MFMA groups (two K16 halves for each of Q2 and Q3). + for _ in range_constexpr(4): + rocdl.sched_dsrd(4) + rocdl.sched_mfma(64) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x64 half-page (16 KiB). Each half has its own LDS base. + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K * ELEM_BYTES) + k_base * fx.Index(ELEM_BYTES) + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one 16-byte half of the wave operand tile. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K * ELEM_BYTES)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def _fp32_k4_operand(full_frag, k32_half, k4): + # A/B 16x32 FP32 wave fragments are i32x8: one FP32 value per + # VGPR and eight K4 MFMA steps per logical K32 tile. Keep the + # existing two-half schedule by grouping four K4 steps per half. + return Vec(full_frag)[k32_half * 4 + k4] + + def _pinned_fp32_mfma_once(acc_idx, a_k4, b_k4): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [arith._to_raw(a_k4), arith._to_raw(b_k4)], + ( + f"v_mfma_f32_16x16x4_f32 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}]" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Accumulate one logical 16x16x32 FP32 product into pinned AGPRs.""" + for k32_half in range_constexpr(2): + for k4 in range_constexpr(4): + _pinned_fp32_mfma_once( + acc_idx, + _fp32_k4_operand(a_frag, k32_half, k4), + _fp32_k4_operand(b_frag, k32_half, k4), + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + # The final logical K32 update is eight in-place K4 FP32 MFMAs. + assert dst_slot == old_acc_idx + pinned_mfma(old_acc_idx, a_frag, b_frag) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def mfma_2n_4mi_k32(subtile_id, n_base, k32_half, a0, a1, a2, a3, b0, b1): + """Issue four K4 steps (one K16 half) for a 4x2 accumulator slab.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for nj in range_constexpr(2): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, n_base + nj), + a_k4, + _fp32_k4_operand(b_frags[nj], k32_half, k4), + ) + + def mfma_4n_4mi_k32(subtile_id, k32_half, a0, a1, a2, a3, b0, b1, b2, b3): + """Issue four K4 steps (one K16 half) for a complete 4x4 quadrant.""" + a_frags = (a0, a1, a2, a3) + b_frags = (b0, b1, b2, b3) + for k4 in range_constexpr(4): + for mi in range_constexpr(4): + a_k4 = _fp32_k4_operand(a_frags[mi], k32_half, k4) + for ni in range_constexpr(4): + _pinned_fp32_mfma_once( + _acc_idx(subtile_id, mi, ni), + a_k4, + _fp32_k4_operand(b_frags[ni], k32_half, k4), + ) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + buffer_ops.buffer_store(Vec(acc)[ii], c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K * ELEM_BYTES) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Compute is K16-half-major across all 16 + # independent accumulators, eliminating the two-deep same-AGPR + # dependency chains produced by pinned_mfma(). + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n_4mi_k32(0, 0, 0, a00, a01, a02, a03, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n_4mi_k32(0, 2, 0, a00, a01, a02, a03, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + # K32 slice 0 already covers K[0:16]. + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + # Keep this refill/LDS-read slot compute-free. + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n_4mi_k32(0, 0, 1, a00, a01, a02, a03, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n_4mi_k32(0, 2, 1, a00, a01, a02, a03, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + # K32 slice 1 already covers K[16:32]. + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + # Keep this refill/LDS-read slot compute-free. + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n_4mi_k32(1, 0, 0, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n_4mi_k32(1, 2, 0, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + # K32 slice 0 already covers K[0:16]. + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + # Keep this refill slot compute-free. + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n_4mi_k32(1, 0, 1, a00, a01, a02, a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n_4mi_k32(1, 2, 1, a00, a01, a02, a03, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + # K32 slice 1 already covers K[16:32]. + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + # Keep this refill slot compute-free. + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n_4mi_k32(2, 0, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + # K32 slice 0 already covers K[0:16]. + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n_4mi_k32(2, 1, a10, a11, a12, a13, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + # K32 slice 1 already covers K[16:32]. + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n_4mi_k32(3, 0, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + # K32 slice 0 already covers K[0:16]. + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n_4mi_k32(3, 1, a10, a11, a12, a13, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + # K32 slice 1 already covers K[16:32]. + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K32 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch(K: int, use_xcd_remap: bool = True): + return _compile_kernel(K, use_xcd_remap=use_xcd_remap) + + + +def fp32_matmul( + a: torch.Tensor, + b: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN FP32 GEMM adapter. + + Public/backend contract: + a: [M, K] FP32 + b: [K, N] FP32 + c: [M, N] FP32 output + + The optimized core streams both operands with K contiguous and therefore + privately consumes B as [N, K]. In the normal TE TN path, ``b`` is a + transpose view of contiguous rowwise weight storage, so ``b.T`` is already + contiguous and does not require a physical transpose. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP32 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + if a.dtype != torch.float32 or b.dtype != torch.float32: + raise TypeError( + "FlyDSL FP32 GEMM expects both operands to have torch.float32 dtype, " + f"got {a.dtype} and {b.dtype}" + ) + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype != torch.float32: + raise TypeError( + f"The current FlyDSL FP32 kernel stores torch.float32 output, got {c.dtype}" + ) + if a.device != b.device or a.device != c.device: + raise ValueError( + f"A, B, and C must be on the same device, got " + f"{a.device}, {b.device}, and {c.device}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP32 GEMM requires contiguous output storage") + + b_hk = b.transpose(0, 1).contiguous() + doGemm(a, b_hk, c, stream=stream) + + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch the private K-specialized FP32 core. + + A and B are shaped [M, K] and [N, K]; C is shaped [M, N]. M and N + remain runtime values, while K selects the cached compile-time specialization. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + assert A.dtype == torch.float32 and B.dtype == torch.float32 + assert C.dtype == torch.float32 + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP32 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert C.shape == (M_runtime, N_runtime) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.contiguous().view(torch.uint8).view(-1) + B_arg = B.contiguous().view(torch.uint8).view(-1) + C_arg = C.view(-1) + launch = _cached_launch(int(K_runtime), bool(use_xcd_remap)) + launch(A_arg, B_arg, C_arg, M_runtime, N_runtime, stream=stream) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py new file mode 100644 index 000000000..b29fa67d8 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm.py @@ -0,0 +1,1183 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [M, K] and +[N, K], one FP32 inverse +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts Transformer Engine's TN contract and +performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 1 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, b_f8_ir_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x128 half-page (16 KiB). Each half has its own LDS base. + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + A_scale_inv, + B_scale_inv, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing TN tensor-wise FP8 adapter. + + Public/backend contract: + a: [M, K] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [K, N] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + The optimized private core streams both operands as row-major [Rows, K], + so B is adapted from TE's logical [K, N] representation to [N, K]. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 TN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + # In the normal TE TN path, b is a transpose view of contiguous rowwise + # weight storage, so b.T is already contiguous and this does not require a + # physical transpose/copy. + b_hk = b.transpose(0, 1).contiguous() + doGemm( + a, + b_hk, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py new file mode 100644 index 000000000..122d1ad95 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nn.py @@ -0,0 +1,1208 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave NN GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [K, M] and +[N, K], one FP32 inverse +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts an NN contract and +performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_linear_128x128, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 1 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, b_f8_ir_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_linear_128x128(lane, wave_id, c_m, LOAD_PASSES_HALF) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # One logical A K64 half is two ds_read_b64_tr_b8 instructions. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # NN A is contiguous [K, M]. Stage a physical row-major [K128, M128] + # half-page without the TN XOR swizzle; S2R performs the transpose. + m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + global_base = k_base * fx.Index(c_m) + m_base + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # Physical LDS A is [K128, M128]. The CDNA4 transpose-read lane + # mapping addresses 8 K rows x 16 M columns per K64 operand half. + # + # The wave-level base follows the documented transpose-load layout: + # k_row = lane_div_32 * 4 + (lane_mod_16 // 4) -> 0..7 + # m_col = m_tile + n-group/lane-in-quad offset -> 0..15 + # Two addresses separated by 32 K rows produce the complementary + # halves required for the complete K64 i32x4 operand. + local_m_tile = ( + (reg_subtile_m_idx0 + fx.Index(sm * 2)) * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + k_row = (fx.Index(lane) // fx.Index(32)) * fx.Index(4) + (lane_mod_16 // fx.Index(4)) + m_col = local_m_tile + ((fx.Index(lane) // fx.Index(16)) % fx.Index(2)) * fx.Index(8) + (lane_mod_16 % fx.Index(4)) * fx.Index(2) + k_half_base = fx.Index(half * 64) + first = (k_half_base + k_row) * fx.Index(BLOCK_M // 2) + m_col + second = (k_half_base + k_row + fx.Index(32)) * fx.Index(BLOCK_M // 2) + m_col + return s2r.load_one_transpose(lds_a[sm], fx.Int32(first), fx.Int32(second)) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + A_scale_inv, + B_scale_inv, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch_nn( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing NN tensor-wise FP8 adapter. + + Public/backend contract: + a: [K, M] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [N, K] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + The NN core consumes TE's existing physical payloads directly: + A is contiguous columnwise storage [K, M] and B is contiguous rowwise + storage [N, K]. No transpose or materialization is performed. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 NN expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + k, m = a.shape + n, kb = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + if not a.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous A [K, M] storage; " + "refusing to materialize a replacement" + ) + if not b.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous B [N, K] storage; " + "refusing to materialize a replacement" + ) + + doGemm( + a, + b, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + K_runtime, M_runtime = A.shape + N_runtime, Kb_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch_nn( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py new file mode 100644 index 000000000..975fd898d --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_nt.py @@ -0,0 +1,1225 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL tensor-wise FP8 4-wave NT GEMM kernel for Transformer Engine. + +The kernel specializes on K at compile time because the K128 loop is fully +hand-unrolled. M/N are runtime launch dimensions. The private optimized core +consumes independently typed FP8 E4M3 or E5M2 A/B tensors shaped [K, M] and +[K, N], one FP32 inverse +scale per operand, and writes float16, bfloat16, or float32 C shaped [M, N]. The public +``fp8_matmul`` entry point accepts an NT contract and +performs the required private adaptation. + +This module imports ``flydsl`` at import time and must therefore be imported +lazily only after FlyDSL availability has been confirmed. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from .exceptions import FlyDSLUnsupportedError + +# Transformer Engine-local FlyDSL utilities. +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_linear_128x128, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K + +NUM_THREADS = 256 +WARP_SIZE = 64 +NUM_WAVES = NUM_THREADS // WARP_SIZE + +SUBTILE_M = 64 +SUBTILE_N = 64 + +MFMA_M = 16 +MFMA_N = 16 + +SUBTILES_PER_WAVE = 4 +MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M +MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N +ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + +ELEM_BYTES = 1 +VEC_BYTES = 16 + +LDS_ELEMS_A = BLOCK_M * BLOCK_K +LDS_ELEMS_B = BLOCK_N * BLOCK_K +LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES +LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + +LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) +LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 +LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 +PASSES_PER_A_MI = LOAD_PASSES_A_SUBTILE // MFMA_M_PER_SUBTILE + +LDS_SYM_A0 = "fp8_pp_smem_a0" +LDS_SYM_A1 = "fp8_pp_smem_a1" +LDS_SYM_B0 = "fp8_pp_smem_b0" +LDS_SYM_B1 = "fp8_pp_smem_b1" +LDS_ALIAS_DOMAIN = '#llvm.alias_scope_domain' +SCOPE_IDS = ("a0", "a1", "b0", "b1") + +assert BLOCK_K == 128 +# DO NOT CHANGE THE FOLLOWING LINE. +assert NUM_THREADS == 256 +assert LOAD_PASSES_A * NUM_THREADS * VEC_BYTES == LDS_BYTES_A +assert LOAD_PASSES_B * NUM_THREADS * VEC_BYTES == LDS_BYTES_B +assert LOAD_PASSES_A % 2 == 0 +assert LOAD_PASSES_B % 2 == 0 + + +def swizzle_xor16(row, col_in_bytes): + """XOR swizzle for the LDS K-byte coordinate.""" + chunk = col_in_bytes // fx.Index(VEC_BYTES) + byte_in_chunk = col_in_bytes % fx.Index(VEC_BYTES) + row_bits = (row % fx.Index(16)) // fx.Index(2) + swz_chunk = chunk ^ row_bits + return swz_chunk * fx.Index(VEC_BYTES) + byte_in_chunk + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL FP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, b_f8_ir_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + a_scale_rsrc = buffer_ops.create_buffer_resource(A_scale_inv, max_size=True) + b_scale_rsrc = buffer_ops.create_buffer_resource(B_scale_inv, max_size=True) + output_scale = ( + buffer_ops.buffer_load(a_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + * buffer_ops.buffer_load(b_scale_rsrc, fx.Index(0), vec_width=1, dtype=T.f32) + ) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + if const_expr(use_xcd_remap): + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + else: + pid_m, pid_n = divmod(fx.block_idx.x, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_linear_128x128(lane, wave_id, c_m, LOAD_PASSES_HALF) + gl_off_b = compute_global_linear_128x128(lane, wave_id, c_n, LOAD_PASSES_HALF) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def read_pinned_accumulator(acc_idx): + acc_pin = PIN_ACC_BASE + acc_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + def hot_loop_scheduler_q_refill_2n(): + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # One logical transposed K64 half is two ds_read_b64_tr_b8 instructions. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(2) + rocdl.sched_mfma(2) + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + rocdl.sched_barrier(0) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # NT A is contiguous [K, M]. Stage a physical row-major [K128, M128] + # half-page without the TN XOR swizzle; S2R performs the transpose. + m_base = bx_m_idx + fx.Index(subtile * (BLOCK_M // 2)) + global_base = k_base * fx.Index(c_m) + m_base + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + # NT B is contiguous [K, N]. Stage a physical row-major [K128, N128] + # half-page without XOR swizzling; S2R performs the transpose. + n_base = by_n_idx + fx.Index(subtile * (BLOCK_N // 2)) + global_base = k_base * fx.Index(c_n) + n_base + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag_half(lds_b, local_row, half, k_half): + # B is physically [K, N]. Each LDS half-page is [K128, N128]. + # One K64 MFMA half requires two ds_read_b64_tr_b8 instructions, + # exactly like the transposed A path. + half_col = local_row - fx.Index(half * (BLOCK_N // 2)) + k_base = fx.Index(k_half * 64) + first = k_base * fx.Index(BLOCK_N // 2) + half_col + second = first + fx.Index(32 * (BLOCK_N // 2)) + return s2r.load_one_transpose( + lds_b[half], + fx.Int32(first), + fx.Int32(second), + ) + + def load_b_frag(lds_b, local_row, half): + x0 = load_b_frag_half(lds_b, local_row, half, 0) + x1 = load_b_frag_half(lds_b, local_row, half, 1) + return pack_frag_halves(x0, x1) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def pinned_mfma(acc_idx, a_frag, b_frag): + """Issue ordinary FP8 MFMA into the fixed physical accumulator bank.""" + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}}," + f"~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}" + ), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag): + """Final-page ordinary FP8 MFMA with independently named AGPR source/destination.""" + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + arith._to_raw(a_frag), + arith._to_raw(b_frag), + ], + ( + f"v_mfma_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + ( + f"v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}}," + f"~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}" + ), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, b0, b1, b2, b3): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + pinned_mfma(acc_base + 2, a_frag, b2) + pinned_mfma(acc_base + 3, a_frag, b3) + + def mfma_2n(acc_base, a_frag, b0, b1): + pinned_mfma(acc_base + 0, a_frag, b0) + pinned_mfma(acc_base + 1, a_frag, b1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] * output_scale + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, sn, ni): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + return load_b_frag(lds_b, b_row_addr, sn) + + def load_b_subtile_regs(lds_b, sn): + return ( + load_b_subtile_ni_regs(lds_b, sn, 0), + load_b_subtile_ni_regs(lds_b, sn, 1), + load_b_subtile_ni_regs(lds_b, sn, 2), + load_b_subtile_ni_regs(lds_b, sn, 3), + ) + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # Physical LDS A is [K128, M128]. The CDNA4 transpose-read lane + # mapping addresses 8 K rows x 16 M columns per K64 operand half. + # + # The wave-level base follows the documented transpose-load layout: + # k_row = lane_div_32 * 4 + (lane_mod_16 // 4) -> 0..7 + # m_col = m_tile + n-group/lane-in-quad offset -> 0..15 + # Two addresses separated by 32 K rows produce the complementary + # halves required for the complete K64 i32x4 operand. + local_m_tile = ( + (reg_subtile_m_idx0 + fx.Index(sm * 2)) * fx.Index(SUBTILE_M) + + fx.Index(mi * MFMA_M) + - fx.Index(sm * (BLOCK_M // 2)) + ) + k_row = (fx.Index(lane) // fx.Index(32)) * fx.Index(4) + (lane_mod_16 // fx.Index(4)) + m_col = local_m_tile + ((fx.Index(lane) // fx.Index(16)) % fx.Index(2)) * fx.Index(8) + (lane_mod_16 % fx.Index(4)) * fx.Index(2) + k_half_base = fx.Index(half * 64) + first = (k_half_base + k_row) * fx.Index(BLOCK_M // 2) + m_col + second = (k_half_base + k_row + fx.Index(32)) * fx.Index(BLOCK_M // 2) + m_col + return s2r.load_one_transpose(lds_a[sm], fx.Int32(first), fx.Int32(second)) + + def load_a_subtile_mi_regs(lds_a, sm, mi): + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + return pack_frag_halves(x0, x1) + + def load_a_subtile_regs(lds_a, sm): + return ( + load_a_subtile_mi_regs(lds_a, sm, 0), + load_a_subtile_mi_regs(lds_a, sm, 1), + load_a_subtile_mi_regs(lds_a, sm, 2), + load_a_subtile_mi_regs(lds_a, sm, 3), + ) + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + ): + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, b00, b01) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, b02, b03) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, b00, b01) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, b02, b03) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, b00, b01) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, b02, b03) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, b00, b01) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, b02, b03) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, b12, b13) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, b10, b11) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, b12, b13) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 1, 0), a01, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 2, 0), a02, b00, b01, b02, b03) + mfma_4n(_acc_idx(0, 3, 0), a03, b00, b01, b02, b03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 1, 0), a01, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 2, 0), a02, b10, b11, b12, b13) + mfma_4n(_acc_idx(1, 3, 0), a03, b10, b11, b12, b13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00 = load_a_subtile_mi_regs(next_a, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, b00, b01, b02, b03) + + next_a01 = load_a_subtile_mi_regs(next_a, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, b00, b01, b02, b03) + + next_a02 = load_a_subtile_mi_regs(next_a, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, b00, b01, b02, b03) + + next_a03 = load_a_subtile_mi_regs(next_a, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, b00, b01, b02, b03) + + next_b00 = load_b_subtile_ni_regs(next_b, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, b10, b11, b12, b13) + + next_b01 = load_b_subtile_ni_regs(next_b, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, b10, b11, b12, b13) + + next_b02 = load_b_subtile_ni_regs(next_b, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, b10, b11, b12, b13) + + next_b03 = load_b_subtile_ni_regs(next_b, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, b10, b11, b12, b13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = (next_a00, next_a01, next_a02, next_a03) + next_b0_regs = (next_b00, next_b01, next_b02, next_b03) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03 = a0_regs + b00, b01, b02, b03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10 = load_b_subtile_ni_regs(cur_b, 1, 0) + b11 = load_b_subtile_ni_regs(cur_b, 1, 1) + b12 = load_b_subtile_ni_regs(cur_b, 1, 2) + b13 = load_b_subtile_ni_regs(cur_b, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = load_a_subtile_mi_regs(cur_a, 1, 0) + a11 = load_a_subtile_mi_regs(cur_a, 1, 1) + a12 = load_a_subtile_mi_regs(cur_a, 1, 2) + a13 = load_a_subtile_mi_regs(cur_a, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + a0_regs = load_a_subtile_regs(lds_a0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + b0_regs = load_b_subtile_regs(lds_b0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + else: + a0_regs, b0_regs = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + + # Common two-page tail. The penultimate tile uses Q2/Q3 carry-prefetch + # to prepare A-top/B-left for the final tile, but performs no K+2 refill. + if (NUM_K_TILES % 2) == 0: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs) + else: + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs) + + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + C: fx.Tensor, + A_scale_inv: fx.Tensor, + B_scale_inv: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + B, + C, + A_scale_inv, + B_scale_inv, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch_nt( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, + use_xcd_remap: bool = True, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + use_xcd_remap=use_xcd_remap, + ) + + + +def fp8_matmul( + a: torch.Tensor, + a_scale_inv: torch.Tensor, + b: torch.Tensor, + b_scale_inv: torch.Tensor, + c: torch.Tensor, + stream=None, +): + """TE-facing NT tensor-wise FP8 adapter. + + Public/backend contract: + a: [K, M] FP8 E4M3 or E5M2 activation payload + a_scale_inv: one-element FP32 inverse quantization scale + b: [K, N] FP8 E4M3 or E5M2 weight payload + b_scale_inv: one-element FP32 inverse quantization scale + c: [M, N] float16, bfloat16, or float32 output + + The NT core consumes TE's existing physical payloads directly: + A and B are contiguous columnwise payloads [K, M] and [K, N]. + No transpose or materialization is performed. + """ + if not isinstance(a, torch.Tensor) or not isinstance(b, torch.Tensor): + raise TypeError("FlyDSL FP8 GEMM expects plain torch.Tensor payloads") + + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL FP8 NT expects rank-2 operands, got A{tuple(a.shape)} " + f"and B{tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL FP8 GEMM expects E4M3 or E5M2 payloads, " + f"got A={a.dtype} and B={b.dtype}" + ) + + k, m = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Inner dimensions do not match: A{tuple(a.shape)} and B{tuple(b.shape)}" + ) + + for name, scale in ( + ("A_scale_inv", a_scale_inv), + ("B_scale_inv", b_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise TypeError( + f"{name} must contain exactly one FP32 value, got " + f"dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + if tuple(c.shape) != (m, n): + raise ValueError(f"C shape {tuple(c.shape)} != expected {(m, n)}") + if c.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL FP8 supports only float16, bfloat16, and float32 " + f"outputs, got {c.dtype}" + ) + if not c.is_contiguous(): + raise ValueError("FlyDSL FP8 requires contiguous output storage") + + tensors = (a, b, a_scale_inv, b_scale_inv, c) + if any(t.device != a.device for t in tensors[1:]): + raise ValueError( + "A, B, inverse scales, and C must be on the same device" + ) + + if not a.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous A [K, M] storage; " + "refusing to materialize a replacement" + ) + if not b.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous B [K, N] storage; " + "refusing to materialize a replacement" + ) + + doGemm( + a, + b, + c, + a_scale_inv, + b_scale_inv, + stream=stream, + ) + +def doGemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale_inv: torch.Tensor, + B_scale_inv: torch.Tensor, + stream=None, + use_xcd_remap: bool = True, +): + """Launch tensor-wise FP8 GEMM with TE-style inverse input scales.""" + K_runtime, M_runtime = A.shape + Kb_runtime, N_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A FP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B FP8 dtype: {B.dtype}" + assert C.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL FP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + assert A_scale_inv.dtype == torch.float32 and A_scale_inv.numel() == 1 + assert B_scale_inv.dtype == torch.float32 and B_scale_inv.numel() == 1 + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + if stream is None: + stream = torch.cuda.current_stream() + + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + C_arg = C.contiguous().view(-1) + A_scale_arg = A_scale_inv.contiguous().view(-1) + B_scale_arg = B_scale_inv.contiguous().view(-1) + + launch = _cached_launch_nt( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + bool(use_xcd_remap), + ) + launch( + A_arg, + B_arg, + C_arg, + A_scale_arg, + B_scale_arg, + M_runtime, + N_runtime, + stream=stream, + ) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py new file mode 100644 index 000000000..b8ed21535 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/fp8_gemm_utils.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as _llvm, vector +from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace +from flydsl.expr import arith, const_expr, range_constexpr, rocdl +from flydsl.expr.typing import Vector as Vec +from flydsl.expr.utils.arith import _to_raw as as_mlir_value + +# ceildiv is the canonical cdiv from the shared layer +def cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +ceildiv = cdiv + +def divmod(a, b): + """Integer divmod that works on DSL values (e.g. ``Int32``). + + The builtin ``divmod`` rejects DSL scalar types, so this uses the overloaded + ``//`` / ``%`` operators to emit the corresponding ops. + """ + return (a // b, a % b) + + +def preshuffle_b(b_t): + """Permute row-major ``B_T`` ``(N, K)`` for ``b_preshuffled=True``.""" + n, k = b_t.shape[-2:] + assert n % 16 == 0 and k % 64 == 0, f"need N%16==0 and K%64==0, got N={n} K={k}" + return b_t.reshape(n // 16, 16, k // 64, 4, 16).permute(0, 2, 3, 1, 4).contiguous() + + +def make_fp8_buffer_tensor(arg_i8, fp8_ir_t): + # max_size=False with no num_records_bytes: cosize(layout) becomes a + # runtime expression because TensorAdaptor defaults to layout-dynamic + # memref (post #554), so the descriptor adapts to the actual tensor + # extent and no longer bakes the first-call's shape into IR. + t_i8 = fx.rocdl.make_buffer_tensor(arg_i8, max_size=False) + iter_i8 = fx.get_iter(t_i8) + f8_buf_ptr_ty = fx.PointerType.get( + elem_ty=fp8_ir_t, + address_space=TargetAddressSpace.BufferDesc, + alignment=fx.PointerType(iter_i8.type).alignment, + ) + iter_f8 = fx.recast_iter(f8_buf_ptr_ty, iter_i8) + return fx.Tensor(fx.make_view(iter_f8, fx.get_layout(t_i8))) + + +def swizzle_128(row, col): + offset = row * 128 + col + swizzle = ((offset % (16 * 128)) >> 8) << 4 + swizzled_offset = offset ^ swizzle + return swizzled_offset // 128, swizzled_offset % 128 + + +def compute_global_swizzle(lane_id, wave_id, K, n_rounds, preshuffled): + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + if const_expr(preshuffled): + row = lane_id % 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id // 8) * 16 + offsets.append( + (row // 16) * (K * 16) + (row % 16) * 16 + (col // 64) * 1024 + ((col % 64) // 16) * 256 + (col % 16) + ) + else: + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + r, c = swizzle_128(row, col) + offsets.append(r * K + c) + return offsets + + +def compute_global_linear_128x128(lane_id, wave_id, leading_dim, n_rounds): + """Offsets for an unswizzled row-major 128x128 tile. + + This uses the same 16-byte/thread DMA decomposition as + ``compute_global_swizzle`` but does not XOR-permute the logical source + coordinates. It is used by the NN A path, whose LDS page is physically + [K128, M128] for the CDNA4 transpose-read instruction. + """ + offsets = [] + n_waves = fx.block_dim.x // 64 + for round in range_constexpr(n_rounds): + row = lane_id // 8 + wave_id * 8 + round * (n_waves * 8) + col = (lane_id % 8) * 16 + offsets.append(row * leading_dim + col) + return offsets + + +class G2SLoader: + def __init__(self, gl_src, gl_offsets, n_load_steps, lds_dtype, wave_id): + self.g2lds_atom = fx.make_copy_atom(fx.rocdl.BufferCopyLDS128b(), 128) + self.LdsPtr_t = fx.PointerType.get(lds_dtype, 2, 512) + self.gl_src = gl_src + self.gl_offsets = gl_offsets + self.n_load_steps = n_load_steps + self.wave_id = wave_id + self.n_waves = fx.block_dim.x // 64 + + def _lds_dst_at(self, lds_dst, step): + step_off = self.wave_id * 1024 + step * (self.n_waves * 1024) + base_i32 = fx.Int32(fx.ptrtoint(lds_dst.ptr)) + sum_i32 = base_i32 + fx.Int32(step_off) + lds_ptr = fx.inttoptr(self.LdsPtr_t, sum_i32) + return fx.make_view(lds_ptr, fx.make_layout(1, 1)) + + def load(self, lds_dst, k_offset): + for step in range_constexpr(self.n_load_steps): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + def load_one(self, lds_dst, k_offset, step): + src = fx.slice(self.gl_src, (None, fx.Int32(self.gl_offsets[step]))) + dst = self._lds_dst_at(lds_dst, step) + fx.copy(self.g2lds_atom, src, dst, soffset=fx.Int32(k_offset)) + + +def pack_i32x4_i32x8(lo, hi): + # Pack two i32x4 as one i32x8 + return lo.shuffle(hi, list(range(8))) + + +class S2RLoader: + def __init__(self, wave_idx, n_tiles): + self.lane_id = fx.thread_idx.x % 64 + self.wave_idx = wave_idx + self.n_tiles = n_tiles + + def _vec_load_16xf8(self, lds_src, offset): + off_tup = fx.make_int_tuple(offset) + ptr_off = fx.add_offset(lds_src.ptr, off_tup) + i8_iter = fx.recast_iter(fx.Uint8, ptr_off) + view = fx.make_view(i8_iter, fx.make_layout(16, 1)) + return view.load() + + def load(self, lds_src, preshuffled=False): + frag = [] + for i in range_constexpr(self.n_tiles): + halves = [] + row = self.wave_idx * (self.n_tiles * 16) + i * 16 + self.lane_id % 16 + for step in range_constexpr(2): + col = (self.lane_id // 16) * 16 + step * 64 + if const_expr(preshuffled): + offset = (row // 8) * 1024 + (row % 8) * 16 + (col // 16) * 128 + else: + row_swz, col_swz = swizzle_128(row, col) + offset = row_swz * 128 + col_swz + v = self._vec_load_16xf8(lds_src, offset) + halves.append(v.bitcast(fx.Int32)) + frag.append(pack_i32x4_i32x8(halves[0], halves[1])) + return frag + + def load_one(self, lds_src, lds_offset): + v = self._vec_load_16xf8(lds_src, lds_offset) + return v.bitcast(fx.Int32) + + def _ds_read_b64_tr_b8(self, lds_src, byte_offset): + """Issue one gfx950 ``ds_read_b64_tr_b8`` and return i32x2. + + The inline-asm output uses one even-aligned 64-bit VGPR tuple. The + compiler owns allocation of the ``=v`` tuple; the memory clobber keeps + the operation ordered with respect to LDS traffic. + """ + base_i32 = fx.Int32(fx.ptrtoint(lds_src.ptr)) + addr_i32 = base_i32 + fx.Int32(byte_offset) + raw_type = ir.VectorType.get([2], ir.IntegerType.get_signless(32)) + raw = _llvm.inline_asm( + raw_type, + [as_mlir_value(addr_i32)], + "ds_read_b64_tr_b8 $0, $1\n", + "=v,v,~{memory}", + has_side_effects=True, + ) + return Vec(vector.BitCastOp(raw_type, raw).result, (2,), fx.Int32) + + def load_one_transpose(self, lds_src, first_byte_offset, second_byte_offset): + """Load one K64 FP8 MFMA operand half from physical LDS [K, M]. + + CDNA4 requires two ``ds_read_b64_tr_b8`` instructions for the complete + K64 operand. Each instruction returns i32x2; concatenation preserves the + existing i32x4 half-fragment interface used by the GEMM hot loop. + """ + lo = self._ds_read_b64_tr_b8(lds_src, first_byte_offset) + hi = self._ds_read_b64_tr_b8(lds_src, second_byte_offset) + return lo.shuffle(hi, [0, 1, 2, 3]) + + +class StoreC: + def __init__(self, A_scale, B_scale, C, c_rows, c_cols, c_idx_fn, n_tiles_a, n_tiles_b): + self.c_rows = c_rows + self.c_cols = c_cols + self.lane_id = fx.thread_idx.x % 64 + self.c_idx_fn = c_idx_fn + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + # Exact byte counts from compile-time shape (BF16 C output, FP32 scales). + # ``num_records_bytes`` is required when ``max_size=False`` -- see + # ``make_buffer_tensor`` docstring for the silent-OOB rationale. + c_nbytes = c_rows * c_cols * 2 # BFloat16 = 2 bytes + sa_nbytes = c_rows * 4 # Float32 row-wise scale + sb_nbytes = c_cols * 4 # Float32 col-wise scale + gC = fx.rocdl.make_buffer_tensor(C, max_size=False, num_records_bytes=c_nbytes) + gSA = fx.rocdl.make_buffer_tensor(A_scale, max_size=False, num_records_bytes=sa_nbytes) + gSB = fx.rocdl.make_buffer_tensor(B_scale, max_size=False, num_records_bytes=sb_nbytes) + self.c_div = fx.logical_divide(gC, fx.make_layout(1, 1)) + self.sa_div = fx.logical_divide(gSA, fx.make_layout(1, 1)) + self.sb_div = fx.logical_divide(gSB, fx.make_layout(1, 1)) + + self.scale_atom_4 = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32) + self.scale_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + self.out_atom_1 = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), fx.BFloat16) + self.reg_f32_4 = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Float32) + self.reg_f32_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + self.reg_bf16_1 = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.BFloat16) + + def _load_scale_vec4(self, row): + fx.copy(self.scale_atom_4, fx.slice(self.sa_div, (None, fx.Int32(row))), self.reg_f32_4) + return Vec(fx.memref_load_vec(self.reg_f32_4)) + + def _load_scale_scalar(self, col): + fx.copy(self.scale_atom_1, fx.slice(self.sb_div, (None, fx.Int32(col))), self.reg_f32_1) + return Vec(fx.memref_load_vec(self.reg_f32_1))[0] + + def _store_bf16(self, value_bf16, c_index): + fx.memref_store_vec(Vec.filled(1, value_bf16, fx.BFloat16), self.reg_bf16_1) + fx.copy(self.out_atom_1, self.reg_bf16_1, fx.slice(self.c_div, (None, fx.Int32(c_index)))) + + def store(self, c_frag, base_row, base_col): + a_scales = [ + self._load_scale_vec4(base_row + i * 16 + (self.lane_id // 16) * 4) for i in range_constexpr(self.n_tiles_a) + ] + b_scales = [ + self._load_scale_scalar(base_col + i * 16 + self.lane_id % 16) for i in range_constexpr(self.n_tiles_b) + ] + for ti in range_constexpr(self.n_tiles_a): + row = base_row + ti * 16 + (self.lane_id // 16) * 4 + for tj in range_constexpr(self.n_tiles_b): + col = base_col + tj * 16 + self.lane_id % 16 + col_valid = col < self.c_cols + oob = fx.Int32(self.c_rows * self.c_cols) + vec_f32 = Vec(c_frag[self.c_idx_fn(ti, tj)]) + for i in range_constexpr(4): + scaled = (vec_f32[i] * (a_scales[ti][i] * b_scales[tj])).to(fx.BFloat16) + c_index = (row + i) * self.c_cols + col + self._store_bf16(scaled, arith.select(col_valid, c_index, oob)) + + +def wait_barrier(count): + _llvm.inline_asm( + res=None, + operands_=[], + asm_string=f"s_waitcnt vmcnt({count})\ns_barrier", + constraints="", + has_side_effects=True, + ) + + +class Mfma16x16x128: + def __init__(self, n_tiles_a, n_tiles_b): + self.atom = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN)) + self.zero_value = Vec.filled(4, 0.0, fx.Float32) + self.n_tiles_a = n_tiles_a + self.n_tiles_b = n_tiles_b + + def idx(self, i, j): + return i * self.n_tiles_b + j + + def _make_operand_frag(self, value): + frag = fx.make_rmem_tensor(8, fx.Int32) + frag.store(Vec(value)) + return frag + + def _make_accum_frag(self, value): + frag = fx.make_rmem_tensor(4, fx.Float32) + frag.store(Vec(value)) + return frag + + def _do_mma(self, a, b, c): + a_frag = self._make_operand_frag(a) + b_frag = self._make_operand_frag(b) + c_frag = self._make_accum_frag(c) + fx.gemm(self.atom, c_frag, a_frag, b_frag, c_frag) + return c_frag.load().ir_value() + + def call(self, a, b, c, *, set_prio=True): + assert len(a) == self.n_tiles_a + assert len(b) == self.n_tiles_b + assert len(c) == self.n_tiles_a * self.n_tiles_b + + a_frags = [self._make_operand_frag(a[idx]) for idx in range_constexpr(self.n_tiles_a)] + b_frags = [self._make_operand_frag(b[idx]) for idx in range_constexpr(self.n_tiles_b)] + c_frags = [self._make_accum_frag(c[idx]) for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + if const_expr(set_prio): + rocdl.s_setprio(1) + for i in range_constexpr(self.n_tiles_a): + for j in range_constexpr(self.n_tiles_b): + cf = c_frags[self.idx(i, j)] + fx.gemm(self.atom, cf, a_frags[i], b_frags[j], cf) + if const_expr(set_prio): + rocdl.s_setprio(0) + rocdl.s_barrier() + return [c_frags[idx].load().ir_value() for idx in range_constexpr(self.n_tiles_a * self.n_tiles_b)] + + def call_one(self, a, b, c, i, j): + assert i < self.n_tiles_a and j < self.n_tiles_b + + return self._do_mma(a[i], b[j], c[self.idx(i, j)]) diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py new file mode 100644 index 000000000..5beff5a75 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -0,0 +1,1119 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""TE entry points for the FlyDSL GEMM backend.""" + +import os + +import torch +import transformer_engine_torch as tex + +from transformer_engine.pytorch.utils import get_device_compute_capability + +from .bf16_gemm import bf16_matmul +from .fp16_gemm import fp16_matmul +from .fp32_gemm import fp32_matmul +from .fp8_gemm import fp8_matmul +from .fp8_gemm_nn import fp8_matmul as fp8_matmul_nn +from .fp8_gemm_nt import fp8_matmul as fp8_matmul_nt +from .mxfp8_gemm import mxfp8_matmul + + +def _product(shape): + """Return the product of dimensions in ``shape``.""" + result = 1 + for dim in shape: + result *= dim + return result + + +def _get_gemm_output_shape(A, transa, B, transb) -> torch.Size: + """Compute TE's logical GEMM output shape. + + This matches ``getGemmOutputShape`` in the C++/Triton backends: the + physical GEMM is flattened to ``[M, N]``, while the returned tensor keeps + B's leading dimensions when ``transb`` is false. + """ + A_shape = A if isinstance(A, torch.Size) else A.shape + B_shape = B if isinstance(B, torch.Size) else B.shape + + if len(A_shape) < 2 or len(B_shape) < 2: + raise ValueError( + "FlyDSL GEMM expects both logical operands to have rank >= 2, " + f"got A={tuple(A_shape)} and B={tuple(B_shape)}" + ) + + A0 = _product(A_shape[:-1]) + A1 = A_shape[-1] + B1 = B_shape[-1] + + output_shape = [B1] if transb else list(B_shape[:-1]) + output_shape.append(A0 if transa else A1) + return torch.Size(output_shape) + + +def reinterpret_as_fp8_tensor( + a: torch.Tensor, + dtype: tex.DType, +) -> torch.Tensor: + """View TE's uint8 payload as the native torch FP8 dtype for this GPU.""" + capability = get_device_compute_capability() + + # gfx950 uses OCP FP8. gfx942 and earlier ROCm architectures use FNUZ. + use_ocp_fp8 = capability == (9, 5) + + if dtype == tex.DType.kFloat8E4M3: + torch_dtype = ( + torch.float8_e4m3fn + if use_ocp_fp8 + else torch.float8_e4m3fnuz + ) + elif dtype == tex.DType.kFloat8E5M2: + torch_dtype = ( + torch.float8_e5m2 + if use_ocp_fp8 + else torch.float8_e5m2fnuz + ) + else: + raise TypeError(f"Unsupported TE FP8 dtype: {dtype}") + + return a.view(torch_dtype) + +def _validate_common_epilogue( + *, + quantizer, + bias, + gelu, + grad, + accumulate, + alpha, + beta, +): + """Validate features not yet implemented by the FlyDSL GEMM backend.""" + if quantizer is not None: + raise NotImplementedError( + "FlyDSL GEMM output quantization is not implemented" + ) + + if float(alpha) != 1.0 or float(beta) != 0.0: + raise NotImplementedError( + "FlyDSL GEMM currently supports only alpha=1 and beta=0" + ) + + if accumulate: + raise NotImplementedError( + "FlyDSL GEMM accumulation is not implemented" + ) + + if bias is not None and bias.numel() != 0: + raise NotImplementedError( + "FlyDSL GEMM bias is not implemented" + ) + + +def _classify_input(t): + """Classify a GEMM operand for the FlyDSL backend.""" + try: + from transformer_engine.pytorch.float8_tensor import Float8Tensor + from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( + Float8TensorStorage, + ) + if isinstance(t, (Float8Tensor, Float8TensorStorage)): + return "fp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + if isinstance(t, (MXFP8Tensor, MXFP8TensorStorage)): + return "mxfp8", t + except ImportError: + pass + + try: + from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensorStorage, + ) + if isinstance(t, QuantizedTensorStorage): + raise ValueError( + f"The FlyDSL GEMM backend does not support " + f"{type(t).__name__}. Only Float8Tensor / " + f"Float8TensorStorage and MXFP8Tensor / " + f"MXFP8TensorStorage are implemented." + ) + except ImportError: + pass + + return "regular", None + + +def _reinterpret_fp8_payload(data, fp8_dtype, name): + """Reinterpret TE's uint8 payload using its ``tex.DType`` metadata.""" + if data is None: + raise RuntimeError(f"{name} does not contain the required FP8 payload") + + if fp8_dtype not in ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ): + raise TypeError( + f"{name} has unsupported TE FP8 dtype metadata: {fp8_dtype}" + ) + + # TE stores Float8Tensor payloads as uint8. Use TE's shared conversion + # helper so ROCm's correct native torch FP8 type is selected from tex.DType. + if data.dtype == torch.uint8: + return reinterpret_as_fp8_tensor(data, fp8_dtype) + + # A materialized payload may already have been reinterpreted. Accept it + # only when its TE metadata is one of the recognized FP8 enum values. + if data.element_size() == 1 and data.dtype.is_floating_point: + return data + + raise TypeError( + f"{name} FP8 storage must be uint8 or an already reinterpreted " + f"one-byte floating-point tensor, got {data.dtype}" + ) + + +def _valid_fp8_transpose(t): + """Return whether a TE Float8 operand has usable columnwise storage.""" + return ( + hasattr(t, "_transpose") + and t._transpose is not None + and not getattr(t, "_transpose_invalid", False) + ) + + +def _mxfp8_debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _mxfp8_debug(message: str) -> None: + if _mxfp8_debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") + + +def _canonicalize_blas_pair( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Swap TE BLAS operands and apply their original transpose flags.""" + a_flydsl = B_data.transpose(0, 1) if transb else B_data + b_flydsl = A_data.transpose(0, 1) if transa else A_data + return a_flydsl, b_flydsl + + +def _flatten_rowwise(t: torch.Tensor, name: str) -> torch.Tensor: + """Flatten all leading dimensions while preserving the final dimension.""" + if t.ndim < 2: + raise ValueError( + f"FlyDSL GEMM expects {name} to have rank >= 2, got {tuple(t.shape)}" + ) + return t.reshape(-1, t.shape[-1]) + + +def _canonicalize_blas_operands( + A_data: torch.Tensor, + transa: bool, + B_data: torch.Tensor, + transb: bool, +): + """Convert TE's BLAS-shaped operands to FlyDSL row-major operands. + + TE's generic GEMM interface follows BLAS column-major interpretation. + FlyDSL kernels consume ordinary row-major matrices: + + a_flydsl: [M, K] + b_flydsl: [K, N] + + The standard conversion is to swap A/B and apply the original transpose + flags to the swapped operands: + + a_flydsl = op(B) + b_flydsl = op(A) + """ + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + A_flat = _flatten_rowwise(A_data, "A") + B_flat = _flatten_rowwise(B_data, "B") + + a_flydsl, b_flydsl = _canonicalize_blas_pair( + A_flat, + transa, + B_flat, + transb, + ) + + m, k = a_flydsl.shape + kb, n = b_flydsl.shape + if kb != k: + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + raise ValueError( + f"FlyDSL {layout} canonicalization produced incompatible operands: " + f"{tuple(a_flydsl.shape)} @ {tuple(b_flydsl.shape)}" + ) + + return a_flydsl, b_flydsl, m, n, k + + +def _validate_or_allocate_output( + D, + *, + shape, + dtype, + device, + backend_name, +): + if D is None: + return torch.empty(shape, dtype=dtype, device=device) + + if tuple(D.shape) != tuple(shape): + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {tuple(shape)}" + ) + if D.dtype != dtype: + raise TypeError( + f"FlyDSL {backend_name} requires {dtype} output, got {D.dtype}" + ) + if D.device != device: + raise ValueError( + f"D must be on {device}, got {D.device}" + ) + if not D.is_contiguous(): + raise ValueError( + f"FlyDSL {backend_name} requires contiguous output storage" + ) + return D + + +def _run_regular_gemm( + A, + transa, + B, + transb, + D, + *, + dtype, + matmul, + backend_name, + output_dtype=None, +): + """Run FP16/BF16/FP32 through shared TN/NN/NT shape handling.""" + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + f"FlyDSL {backend_name} GEMM expects plain torch.Tensor operands" + ) + if A.dtype != dtype or B.dtype != dtype: + raise TypeError( + f"FlyDSL {backend_name} GEMM requires {dtype} inputs, " + f"got A={A.dtype} and B={B.dtype}" + ) + if A.device != B.device: + raise ValueError( + f"A and B must be on the same device, got {A.device} and {B.device}" + ) + + output_shape = _get_gemm_output_shape(A, transa, B, transb) + + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( + A, transa, B, transb + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL {backend_name} logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) + + if output_dtype is None: + output_dtype = dtype + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=A.device, + backend_name=backend_name, + ) + + matmul( + a_flydsl, + b_flydsl, + D.view(m, n), + ) + return D + + +def _materialize_rowwise_from_columnwise( + transpose_data: torch.Tensor, + name: str, +) -> torch.Tensor: + """Reconstruct logical rowwise FP8 data from TE columnwise storage. + + This matches Triton's ``materialize_rowwise_from_columnwise`` exactly. + TE stores an n-D rowwise tensor ``[D0, ..., Dn-2, K]`` columnwise as + ``[K, D0, ..., Dn-2]``. Recover rowwise storage by rotating the leading + K dimension back to the tail. + """ + if transpose_data.ndim < 2: + raise ValueError( + f"{name} must have rank >= 2, got {tuple(transpose_data.shape)}" + ) + + if transpose_data.ndim == 2: + return transpose_data.transpose(0, 1).contiguous() + + perm = list(range(1, transpose_data.ndim)) + [0] + return transpose_data.permute(*perm).contiguous() + + +def _get_fp8_logical_rowwise_payload(t, name): + """Return logical rowwise FP8 data, matching the Triton wrapper. + + Prefer TE's rowwise ``_data``. If only valid columnwise ``_transpose`` + storage exists, materialize a rowwise copy once for canonicalization. + """ + fp8_dtype = getattr(t, "_fp8_dtype", None) + data = getattr(t, "_data", None) + + if data is not None: + return _reinterpret_fp8_payload( + data, + fp8_dtype, + f"{name}._data", + ) + + if not _valid_fp8_transpose(t): + raise RuntimeError( + f"{name} has neither valid rowwise (_data) nor " + f"columnwise (_transpose) FP8 storage" + ) + + transpose_data = _reinterpret_fp8_payload( + t._transpose, + fp8_dtype, + f"{name}._transpose", + ) + + return _materialize_rowwise_from_columnwise( + transpose_data, + f"{name}._transpose", + ) + + +def _select_mxfp8_data_and_scale( + t, + *, + will_transpose: bool, + name: str, +): + """Select the TE MXFP8 representation required by BLAS semantics.""" + if will_transpose: + data = getattr(t, "_columnwise_data", None) + scale = getattr(t, "_columnwise_scale_inv", None) + orientation = "columnwise" + else: + data = getattr(t, "_rowwise_data", None) + scale = getattr(t, "_rowwise_scale_inv", None) + orientation = "rowwise" + + _mxfp8_debug( + f"{name}: will_transpose={will_transpose}, " + f"selected={orientation}, data_present={data is not None}, " + f"scale_present={scale is not None}" + ) + + if data is None or scale is None: + raise RuntimeError( + f"{name} does not contain required {orientation} MXFP8 data and scales" + ) + + _mxfp8_debug( + f"{name} selected data shape={tuple(data.shape)}, " + f"dtype={data.dtype}, stride={tuple(data.stride())}; " + f"scale shape={tuple(scale.shape)}, dtype={scale.dtype}, " + f"stride={tuple(scale.stride())}" + ) + return data, scale + + +def _flatten_mxfp8_scale(t: torch.Tensor, name: str) -> torch.Tensor: + if t.ndim < 2: + raise ValueError( + f"FlyDSL MXFP8 expects {name} scale rank >= 2, " + f"got {tuple(t.shape)}" + ) + original_shape = tuple(t.shape) + if t.ndim > 2: + t = t.reshape(-1, t.shape[-1]) + _mxfp8_debug( + f"{name} scale flatten: {original_shape} -> {tuple(t.shape)}, " + f"contiguous={t.is_contiguous()}" + ) + return t + + +def _run_mxfp8( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Canonicalize independently typed E4M3/E5M2 MXFP8 operands.""" + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) + if ( + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes + ): + raise NotImplementedError( + "FlyDSL MXFP8 supports E4M3 and E5M2 independently for A/B; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + _mxfp8_debug( + f"entry: layout={layout}, A_type={type(A).__name__}, " + f"B_type={type(B).__name__}, D_provided={D is not None}" + ) + + # Match TE CanonicalizeGemmInput / Triton data_and_scale_for_transpose: + # A: transa=True -> rowwise, transa=False -> columnwise + # B: transb=True -> columnwise, transb=False -> rowwise + A_data, A_scale = _select_mxfp8_data_and_scale( + A, + will_transpose=not transa, + name="A", + ) + B_data, B_scale = _select_mxfp8_data_and_scale( + B, + will_transpose=transb, + name="B", + ) + + # MXFP8Tensor stores rowwise/columnwise payloads as raw uint8. Reinterpret + # those exact bytes using each operand's own FP8 metadata before applying + # BLAS canonicalization. No copy or numerical conversion is performed here. + if A_data.dtype == torch.uint8: + A_data = reinterpret_as_fp8_tensor(A_data, a_fp8_dtype) + if B_data.dtype == torch.uint8: + B_data = reinterpret_as_fp8_tensor(B_data, b_fp8_dtype) + + output_shape = _get_gemm_output_shape( + A_data.shape, transa, B_data.shape, transb + ) + + a_flydsl, b_flydsl, m, n, k = _canonicalize_blas_operands( + A_data, + transa, + B_data, + transb, + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL MXFP8 logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) + + A_scale = _flatten_mxfp8_scale(A_scale, "A") + B_scale = _flatten_mxfp8_scale(B_scale, "B") + a_scale, b_scale = _canonicalize_blas_pair( + A_scale, + transa, + B_scale, + transb, + ) + + _mxfp8_debug( + f"canonicalized layout={layout}: " + f"a={tuple(a_flydsl.shape)}, dtype={a_flydsl.dtype}, " + f"stride={tuple(a_flydsl.stride())}; " + f"b={tuple(b_flydsl.shape)}, dtype={b_flydsl.dtype}, " + f"stride={tuple(b_flydsl.stride())}" + ) + _mxfp8_debug( + f"canonicalized scales: " + f"a_scale={tuple(a_scale.shape)}, stride={tuple(a_scale.stride())}; " + f"b_scale={tuple(b_scale.shape)}, stride={tuple(b_scale.stride())}" + ) + _mxfp8_debug(f"derived GEMM dimensions: M={m}, N={n}, K={k}") + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + scale_group_size = 32 + if k % scale_group_size != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{scale_group_size}" + ) + + # Shared BLAS canonicalization yields: + # a_scale [M, K/32] + # b_scale [K/32, N] + expected_a_scale = (m, k // scale_group_size) + expected_b_scale = (k // scale_group_size, n) + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"A scale shape {tuple(a_scale.shape)} != expected " + f"{expected_a_scale}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"B scale shape {tuple(b_scale.shape)} != expected " + f"{expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="MXFP8", + ) + + mxfp8_matmul( + a_flydsl, + a_scale, + b_flydsl, + b_scale, + D.view(m, n), + ) + return D + + +def _run_fp8( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Run tensor-wise E4M3/E5M2 FP8 combinations for TN/NN/NT. + + NN and NT are dispatched directly from TE's existing physical + representations without transpose kernels or materialized payloads: + + - NN core: [K, M] x [N, K] + - NT core: [K, M] x [K, N] + + TN retains the shared canonicalized path through + ``fp8_gemm.fp8_matmul``. + """ + a_fp8_dtype = getattr(A, "_fp8_dtype", None) + b_fp8_dtype = getattr(B, "_fp8_dtype", None) + supported_fp8_dtypes = ( + tex.DType.kFloat8E4M3, + tex.DType.kFloat8E5M2, + ) + if ( + a_fp8_dtype not in supported_fp8_dtypes + or b_fp8_dtype not in supported_fp8_dtypes + ): + raise NotImplementedError( + "FlyDSL FP8 supports E4M3 and E5M2 independently for A/B; " + f"got A={a_fp8_dtype} and B={b_fp8_dtype}" + ) + + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + A_scale_inv = getattr(A, "_scale_inv", None) + B_scale_inv = getattr(B, "_scale_inv", None) + for name, scale in ( + ("A._scale_inv", A_scale_inv), + ("B._scale_inv", B_scale_inv), + ): + if not isinstance(scale, torch.Tensor): + raise RuntimeError(f"{name} is not populated") + if scale.dtype != torch.float32 or scale.numel() != 1: + raise ValueError( + f"{name} must contain exactly one FP32 tensor-wise inverse " + f"scale, got dtype={scale.dtype}, shape={tuple(scale.shape)}" + ) + + # TE exposes GEMM operands in BLAS/column-major convention. The + # row-major FlyDSL result is formed from the swapped operands: + # + # flydsl_a = op(B) + # flydsl_b = op(A) + # + # For NN, the dedicated kernel consumes: + # + # flydsl_a physical [K, M] = B columnwise storage + # flydsl_b physical [N, K] = A columnwise storage + # + # Here FlyDSL M is TE's n and FlyDSL N is TE's m, so the kernel writes + # the existing TE output allocation in its ordinary [M, N] view. Both + # payloads already exist; this path performs no transpose or materialization. + if not transa and not transb: + if not _valid_fp8_transpose(B): + raise RuntimeError( + "FlyDSL FP8 NN requires valid B columnwise (_transpose) storage" + ) + if not _valid_fp8_transpose(A): + raise RuntimeError( + "FlyDSL FP8 NN requires valid A columnwise (_transpose) storage" + ) + + a_flydsl = _reinterpret_fp8_payload( + B._transpose, + b_fp8_dtype, + "B._transpose", + ) + b_flydsl = _reinterpret_fp8_payload( + A._transpose, + a_fp8_dtype, + "A._transpose", + ) + + if a_flydsl.ndim != 2 or b_flydsl.ndim != 2: + raise ValueError( + "FlyDSL FP8 NN direct path expects rank-2 columnwise storage, " + f"got B._transpose={tuple(a_flydsl.shape)} and " + f"A._transpose={tuple(b_flydsl.shape)}" + ) + if not a_flydsl.is_contiguous() or not b_flydsl.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NN requires contiguous TE columnwise storage; " + "refusing to materialize replacement operands" + ) + + k, m = a_flydsl.shape + n, kb = b_flydsl.shape + if kb != k: + raise ValueError( + "FlyDSL FP8 NN storage mismatch after BLAS operand swap: " + f"B._transpose{tuple(a_flydsl.shape)} and " + f"A._transpose{tuple(b_flydsl.shape)}" + ) + + # Float8TensorStorage does not expose a public ``shape`` attribute. + # The direct NN operands already determine the flattened kernel output + # shape exactly. Preserve TE's preallocated logical output shape when + # one is provided; otherwise use the flattened [M, N] shape. + output_shape = ( + D.shape + if D is not None + else torch.Size((m, n)) + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 NN logical output shape {tuple(output_shape)} " + f"does not match kernel shape {(m, n)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="FP8 NN", + ) + + # Scales follow the swapped FlyDSL operands. + fp8_matmul_nn( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) + return D + + # For TE NT (transa=False, transb=True), the dedicated kernel consumes + # both swapped operands directly from TE columnwise storage: + # + # kernel A physical [K, M] = B._transpose + # kernel B physical [K, N] = A._transpose + # + # Both operands are therefore staged as physical K-major tiles and read + # from LDS with ``ds_read_b64_tr_b8``. No torch transpose, + # ``.contiguous()``, or temporary FP8 payload is introduced. + if not transa and transb: + if not _valid_fp8_transpose(B): + raise RuntimeError( + "FlyDSL FP8 NT requires valid B columnwise (_transpose) storage" + ) + if not _valid_fp8_transpose(A): + raise RuntimeError( + "FlyDSL FP8 NT requires valid A columnwise (_transpose) storage" + ) + + # TE columnwise payloads are contiguous transposes of the logical + # rowwise tensors. After the BLAS operand swap, their exposed shapes are: + # + # B._transpose: [M, K] + # A._transpose: [N, K] + # + # The NT kernel consumes the same physical bytes as: + # + # kernel A: [K, M] + # kernel B: [K, N] + # + # Reinterpret only the 2-D shape. ``view`` is zero-copy and preserves + # the exact columnwise allocation; no torch transpose or materialization + # is performed. + b_columnwise = _reinterpret_fp8_payload( + B._transpose, + b_fp8_dtype, + "B._transpose", + ) + a_columnwise = _reinterpret_fp8_payload( + A._transpose, + a_fp8_dtype, + "A._transpose", + ) + + if b_columnwise.ndim != 2 or a_columnwise.ndim != 2: + raise ValueError( + "FlyDSL FP8 NT direct path expects rank-2 columnwise storage, " + f"got B._transpose={tuple(b_columnwise.shape)} and " + f"A._transpose={tuple(a_columnwise.shape)}" + ) + if not b_columnwise.is_contiguous() or not a_columnwise.is_contiguous(): + raise ValueError( + "FlyDSL FP8 NT requires contiguous TE columnwise storage; " + "refusing to materialize replacement operands" + ) + + m, k = b_columnwise.shape + n, ka = a_columnwise.shape + if ka != k: + raise ValueError( + "FlyDSL FP8 NT columnwise K mismatch after BLAS operand swap: " + f"B._transpose{tuple(b_columnwise.shape)} and " + f"A._transpose{tuple(a_columnwise.shape)}" + ) + + a_flydsl = b_columnwise.view(k, m) + b_flydsl = a_columnwise.view(k, n) + + # Float8TensorStorage does not expose a public ``shape`` attribute. + # Preserve TE's preallocated logical output shape when available. + output_shape = ( + D.shape + if D is not None + else torch.Size((m, n)) + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 NT logical output shape {tuple(output_shape)} " + f"does not match kernel shape {(m, n)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="FP8 NT", + ) + + # Scales follow the BLAS-swapped kernel operands. + fp8_matmul_nt( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) + return D + + # Match Triton's regular-FP8 handling: establish logical rowwise + # payloads first, then apply the same shared BLAS-to-row-major + # canonicalization used for FP16/BF16/FP32. + A_data = _get_fp8_logical_rowwise_payload(A, "A") + B_data = _get_fp8_logical_rowwise_payload(B, "B") + + output_shape = _get_gemm_output_shape( + A_data.shape, transa, B_data.shape, transb + ) + + a_flydsl, b_flydsl, m, n, _ = _canonicalize_blas_operands( + A_data, transa, B_data, transb + ) + if _product(output_shape) != m * n: + raise RuntimeError( + f"FlyDSL FP8 logical output shape {tuple(output_shape)} " + f"does not match flattened GEMM shape {(m, n)}" + ) + + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"A and B must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name="FP8", + ) + + # Operand swap means B's tensor-wise scale belongs to a_flydsl and A's + # tensor-wise scale belongs to b_flydsl. + fp8_matmul( + a_flydsl, + B_scale_inv, + b_flydsl, + A_scale_inv, + D.view(m, n), + ) + return D + + +def te_generic_gemm_flydsl( + A, + transa, + B, + transb, + D, + quantizer, + output_dtype, + bias=None, + bias_type=None, + gelu=False, + gelu_in=None, + grad=False, + workspace=None, + workspaceSize=0, + accumulate=False, + use_split_accumulator=False, + comm_overlap=None, + comm_type=None, + extra_output=None, + bulk_overlap=False, + alpha=1.0, + beta=0.0, +): + """Run a supported FlyDSL GEMM through TE's generic GEMM interface. + + Supported layouts: + - TN: transa=True, transb=False + - NN: transa=False, transb=False + - NT: transa=False, transb=True + + TT is intentionally rejected. + + Supported dtypes: + - MXFP8 E4M3/E5M2 A/B combinations with FP16, BF16, or FP32 output + - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output + - BF16 input with FP16, BF16, or FP32 output + - FP16 input with FP16, BF16, or FP32 output + - FP32 input with FP32 output + """ + del bias_type + del gelu_in + del workspace + del workspaceSize + del use_split_accumulator + del comm_overlap + del comm_type + del extra_output + del bulk_overlap + + if transa and transb: + raise NotImplementedError( + "FlyDSL GEMM does not support transa=True, transb=True (TT)" + ) + + _validate_common_epilogue( + quantizer=quantizer, + bias=bias, + gelu=gelu, + grad=grad, + accumulate=accumulate, + alpha=alpha, + beta=beta, + ) + + a_kind, _ = _classify_input(A) + b_kind, _ = _classify_input(B) + + if a_kind == "mxfp8" or b_kind == "mxfp8": + # Validate both are MXFP8 + if a_kind != b_kind: + raise ValueError( + "Mixed MXFP8 and non-MXFP8 FlyDSL GEMM inputs are not supported" + ) + + # Sanity: both operands must have at least one pre-quantized copy. + if getattr(A, '_rowwise_data', None) is None and getattr(A, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + if getattr(B, '_rowwise_data', None) is None and getattr(B, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + + mxfp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in mxfp8_output_dtypes: + raise NotImplementedError( + "FlyDSL MXFP8 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + + D = _run_mxfp8( + A, + transa, + B, + transb, + D, + output_dtype=mxfp8_output_dtypes[output_dtype], + ) + return D, None, None, None + + if a_kind == "fp8" or b_kind == "fp8": + if a_kind != b_kind: + raise ValueError( + "Mixed regular FP8 and non-FP8 FlyDSL GEMM inputs are not supported" + ) + + fp8_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp8_output_dtypes: + raise NotImplementedError( + "FlyDSL tensor-wise FP8 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + + D = _run_fp8( + A, + transa, + B, + transb, + D, + output_dtype=fp8_output_dtypes[output_dtype], + ) + return D, None, None, None + + if a_kind != "regular" or b_kind != "regular": + raise TypeError( + "Unsupported FlyDSL GEMM operand types: " + f"{type(A).__name__} and {type(B).__name__}" + ) + if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor): + raise TypeError( + "FlyDSL regular GEMM expects plain torch.Tensor operands" + ) + + if A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16: + bf16_output_dtypes = { + None: torch.bfloat16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in bf16_output_dtypes: + raise NotImplementedError( + "FlyDSL BF16 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.bfloat16, + matmul=bf16_matmul, + backend_name="BF16", + output_dtype=bf16_output_dtypes[output_dtype], + ) + return D, None, None, None + + if A.dtype == torch.float16 and B.dtype == torch.float16: + fp16_output_dtypes = { + None: torch.float16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in fp16_output_dtypes: + raise NotImplementedError( + "FlyDSL FP16 supports FP16, BF16, or FP32 output, " + f"got {output_dtype}" + ) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.float16, + matmul=fp16_matmul, + backend_name="FP16", + output_dtype=fp16_output_dtypes[output_dtype], + ) + return D, None, None, None + + if A.dtype == torch.float32 and B.dtype == torch.float32: + if output_dtype not in (None, tex.DType.kFloat32): + raise NotImplementedError( + "FlyDSL FP32 currently supports only FP32 output, " + f"got {output_dtype}" + ) + D = _run_regular_gemm( + A, + transa, + B, + transb, + D, + dtype=torch.float32, + matmul=fp32_matmul, + backend_name="FP32", + ) + return D, None, None, None + + raise NotImplementedError( + "FlyDSL GEMM currently supports only MXFP8, tensor-wise E4M3 FP8, " + "BF16, FP16, or FP32 inputs; " + f"got A={A.dtype} and B={B.dtype}" + ) \ No newline at end of file diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py new file mode 100644 index 000000000..a54dc43d7 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp8_gemm.py @@ -0,0 +1,1406 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""FlyDSL MXFP8 GEMM implementation. + +This module contains both the HK-derived optimized 4-wave kernel and its +MXFP8-specific launch preparation. Transformer Engine BLAS canonicalization +is performed by ``gemm_wrappers.py`` before entering ``mxfp8_matmul``. + +Canonical launch inputs: + + a: [M, K] FP8 E4M3 or E5M2 payload + a_scale: [M, K/32] raw E8M0 bytes + b: [K, N] FP8 E4M3 or E5M2 payload + b_scale: [K/32, N] raw E8M0 bytes + D: [M, N] float16, bfloat16, or float32 output +""" + +import functools +import os + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import llvm +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Transformer Engine-local FlyDSL utilities. +from .exceptions import FlyDSLUnsupportedError +from .fp8_gemm_utils import ( + G2SLoader, + S2RLoader, + compute_global_swizzle, + make_fp8_buffer_tensor, + pack_i32x4_i32x8, + swizzle_128, +) + + +_BLOCK_M = 256 +_BLOCK_N = 256 +_BLOCK_K = 128 + +# Public metadata consumed by wrappers — keep. +BLOCK_M = _BLOCK_M +BLOCK_N = _BLOCK_N +BLOCK_K = _BLOCK_K +SCALE_GROUP_SIZE = 32 + + +def _debug_enabled() -> bool: + value = os.getenv("DEBUG_FLYDSL_MXFP8_GEMM", "") + return value.lower() not in ("", "0", "false", "no", "off") + + +def _debug(message: str) -> None: + if _debug_enabled(): + print(f"[DEBUG_FLYDSL_MXFP8_GEMM] {message}") + + +def pack_mx32_scales_iter(scales_u8: torch.Tensor) -> torch.Tensor: + """Pack raw [Rows, K/32] E8M0 scales as [K/128, Rows] uint32.""" + if scales_u8.dtype != torch.uint8: + raise TypeError( + f"MXFP8 scales must be torch.uint8 E8M0 bytes, got {scales_u8.dtype}" + ) + if scales_u8.ndim != 2: + raise ValueError( + f"MXFP8 scales must be rank 2, got shape {tuple(scales_u8.shape)}" + ) + + rows, qk = scales_u8.shape + if qk % 4 != 0: + raise ValueError( + f"Scale K dimension must be divisible by 4 K32 groups, got {qk}" + ) + + s32 = scales_u8.contiguous().view(rows, qk // 4, 4).to(torch.int32) + packed = ( + s32[:, :, 0] + | (s32[:, :, 1] << 8) + | (s32[:, :, 2] << 16) + | (s32[:, :, 3] << 24) + ) + return packed.transpose(0, 1).contiguous() + + +def pack_mx32_scales_for_hk(scales_u8: torch.Tensor) -> torch.Tensor: + """Convert raw rowwise E8M0 scales to [K/128, Rows] MFMA-ready words.""" + scale_iter = pack_mx32_scales_iter(scales_u8) + rows = scales_u8.shape[0] + + if rows % 64 != 0: + raise ValueError( + f"Rows={rows} must be a multiple of 64 for HK MFMA scale packing" + ) + + device = scales_u8.device + row = torch.arange(rows, device=device, dtype=torch.int64) + row_within_16 = row % 16 + k_subgroup = (row // 16) % 4 + tile = row // 64 + + packed = torch.zeros_like(scale_iter) + for group in range(4): + source_row = tile * 64 + group * 16 + row_within_16 + source_value = scale_iter[:, source_row] + byte_value = ( + source_value >> (k_subgroup * 8).view(1, rows) + ) & 0xFF + packed |= byte_value << (group * 8) + + return packed.contiguous() + + +def _encode_waitcnt(vmcnt=63, lgkmcnt=15): + """Encode the CDNA4/gfx950 ``S_WAITCNT`` SIMM16 operand. + + ``rocdl.s_waitcnt`` accepts the raw 16-bit immediate operand of the + 32-bit ``S_WAITCNT`` ISA instruction. On CDNA4, that SIMM16 field is: + + SIMM16[3:0] = vmcnt[3:0] + SIMM16[6:4] = expcnt[2:0] + SIMM16[11:8] = lgkmcnt[3:0] + SIMM16[15:14] = vmcnt[5:4] + + ``vmcnt`` is therefore one six-bit counter split across two noncontiguous + fields; bits [5:4] are placed in SIMM16[15:14], while bits [3:0] remain + in SIMM16[3:0]. + + A wait-counter field set to its maximum representable value is effectively + unconstrained: the instruction does not wait on that counter. This helper + always encodes ``expcnt=7`` and defaults to ``vmcnt=63`` and ``lgkmcnt=15``, + so callers specify only the counters on which they intend to wait. + + For example, ``_encode_waitcnt(lgkmcnt=0)`` returns ``0xC07F``, which the + assembler renders as ``s_waitcnt lgkmcnt(0)``. + See: https://llvm.org/docs/AMDGPU/gfx9_waitcnt.html + """ + if not 0 <= vmcnt <= 63: + raise ValueError(f"vmcnt must be in [0, 63], got {vmcnt}") + if not 0 <= lgkmcnt <= 15: + raise ValueError(f"lgkmcnt must be in [0, 15], got {lgkmcnt}") + + return ( + (7 << 4) # expcnt=7 -> SIMM16[6:4] (unconstrained) + | (vmcnt & 0x0F) # vmcnt[3:0] -> SIMM16[3:0] + | ((lgkmcnt & 0x0F) << 8) # lgkmcnt[3:0] -> SIMM16[11:8] + | ((vmcnt & 0x30) << 10) # vmcnt[5:4] -> SIMM16[15:14] + ) + + +# Keep the documented gfx950 encoding invariant executable and import-time cheap. +assert _encode_waitcnt(lgkmcnt=0) == 0xC07F + + +def _barrier(vmcnt=63, lgkmcnt=15): + if vmcnt != 63 or lgkmcnt != 15: + rocdl.s_waitcnt(_encode_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt)) + rocdl.s_barrier() + +def _min(a, b): + return arith.select(a < b, a, b) + + +def _divmod(a, b): + return a // b, a % b + + +def _xcd_swizzle(num_pid_m, num_pid_n): + NUM_XCDS = 8 + WGM = 4 + NUM_CUS = 32 * NUM_XCDS + SWIZZLE_THRESHOLD = 4 * NUM_CUS + + wgid = fx.block_idx.x + num_wg = num_pid_m * num_pid_n + + # Simple row-major path. + simple_m, simple_n = _divmod(wgid, num_pid_n) + + # XCD-remapped grouped-M path. + intra_xcd, xcd = _divmod(wgid, NUM_XCDS) + wgid_remap = xcd * (num_wg // NUM_XCDS) + intra_xcd + num_wgid_in_group = WGM * num_pid_n + group_id, intra_group = _divmod(wgid_remap, num_wgid_in_group) + first_pid_m = group_id * WGM + group_size_m = _min(num_pid_m - first_pid_m, WGM) + pid_n, intra_group_m = _divmod(intra_group, group_size_m) + pid_m = first_pid_m + intra_group_m + + use_simple = (num_wg < SWIZZLE_THRESHOLD) | (num_wg % NUM_XCDS != 0) + return ( + arith.select(use_simple, simple_m, pid_m), + arith.select(use_simple, simple_n, pid_n), + ) + + +def _compile_kernel( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, +): + """Build the specialized kernel for compile-time K, A/B FP8 types, and output dtype. + + ``K`` must contain at least four K128 tiles. Runtime M/N are expected to + be exact multiples of ``BLOCK_M``/``BLOCK_N``; the kernel has no edge masks. + """ + BLOCK_M, BLOCK_N, BLOCK_K = _BLOCK_M, _BLOCK_N, _BLOCK_K + + fp8_input_types = { + torch.float8_e4m3fn: (fx.Float8E4M3FN, 0), + torch.float8_e5m2: (fx.Float8E5M2, 1), + } + try: + a_fx_dtype, a_matrix_format = fp8_input_types[a_fp8_dtype] + b_fx_dtype, b_matrix_format = fp8_input_types[b_fp8_dtype] + except KeyError as exc: + raise TypeError( + "FlyDSL MXFP8 input dtype must be torch.float8_e4m3fn or " + f"torch.float8_e5m2, got A={a_fp8_dtype}, B={b_fp8_dtype}" + ) from exc + + if output_dtype == torch.float16: + output_element_bytes = 2 + output_fx_dtype = fx.Float16 + elif output_dtype == torch.bfloat16: + output_element_bytes = 2 + output_fx_dtype = fx.BFloat16 + elif output_dtype == torch.float32: + output_element_bytes = 4 + output_fx_dtype = fx.Float32 + else: + raise TypeError( + "FlyDSL MXFP8 supports only float16, bfloat16, and float32 " + f"outputs, got {output_dtype}" + ) + + NUM_THREADS = 256 + WARP_SIZE = 64 + + SUBTILE_M = 64 + SUBTILE_N = 64 + + MFMA_M = 16 + MFMA_N = 16 + + SUBTILES_PER_WAVE = 4 + MFMA_M_PER_SUBTILE = SUBTILE_M // MFMA_M + MFMA_N_PER_SUBTILE = SUBTILE_N // MFMA_N + ACCS_PER_WAVE = SUBTILES_PER_WAVE * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + + ELEM_BYTES = 1 + VEC_BYTES = 16 + + LDS_ELEMS_A = BLOCK_M * BLOCK_K + LDS_ELEMS_B = BLOCK_N * BLOCK_K + LDS_BYTES_A = LDS_ELEMS_A * ELEM_BYTES + LDS_BYTES_B = LDS_ELEMS_B * ELEM_BYTES + + LOAD_PASSES_A = LDS_BYTES_A // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_B = LDS_BYTES_B // (NUM_THREADS * VEC_BYTES) + LOAD_PASSES_A_SUBTILE = LOAD_PASSES_A // 2 + LOAD_PASSES_B_SUBTILE = LOAD_PASSES_B // 2 + LOAD_PASSES_SCALES = 16 + + assert K % BLOCK_K == 0, f"K must be a multiple of {BLOCK_K}, got {K}" + NUM_K_TILES = K // BLOCK_K + assert NUM_K_TILES >= 4, f"K={K} gives {NUM_K_TILES} K128 tiles; the two-page pipeline needs at least 4" + + LDS_ELEMS_HALF = (BLOCK_M // 2) * BLOCK_K + LOAD_PASSES_HALF = LDS_ELEMS_HALF // (NUM_THREADS * VEC_BYTES) + assert LOAD_PASSES_HALF == LOAD_PASSES_A_SUBTILE == LOAD_PASSES_B_SUBTILE + + @fx.struct + class SharedStorage: + # Each logical 256x128 page is two independent 128x128 half-pages. + # The hot loop refills one 16-byte pass of one half-page at a time. + a0_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a0_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_0: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + a1_1: fx.Array[a_fx_dtype, LDS_ELEMS_HALF, 16] + b0_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b0_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_0: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + b1_1: fx.Array[b_fx_dtype, LDS_ELEMS_HALF, 16] + + @flyc.kernel(known_block_size=[NUM_THREADS, 1, 1]) + def kernel_gemm( + A: fx.Tensor, As: fx.Tensor, B: fx.Tensor, Bs: fx.Tensor, C: fx.Tensor, c_m: fx.Int32, c_n: fx.Int32 + ): + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_a0 = (lds.a0_0, lds.a0_1) + lds_a1 = (lds.a1_0, lds.a1_1) + lds_b0 = (lds.b0_0, lds.b0_1) + lds_b1 = (lds.b1_0, lds.b1_1) + + a_f8_ir_t = a_fx_dtype.ir_type + b_f8_ir_t = b_fx_dtype.ir_type + gA = make_fp8_buffer_tensor(A, a_f8_ir_t) + gB = make_fp8_buffer_tensor(B, b_f8_ir_t) + a_div = fx.logical_divide(gA, fx.make_layout(1, 1)) + b_div = fx.logical_divide(gB, fx.make_layout(1, 1)) + as_rsrc = buffer_ops.create_buffer_resource(As, max_size=True) + bs_rsrc = buffer_ops.create_buffer_resource(Bs, max_size=True) + tx = gpu.thread_id("x") + + num_blocks_m = c_m // BLOCK_M + num_blocks_n = c_n // BLOCK_N + + pid_m, pid_n = _xcd_swizzle(num_blocks_m, num_blocks_n) + + bx_m = pid_m * BLOCK_M + by_n = pid_n * BLOCK_N + + # The flattened/XCD-swizzled block coordinates are i32, while global + # address arithmetic below is expressed in MLIR index type. Convert + # once here and use these index-typed tile bases for every address. + bx_m_idx = fx.Index(bx_m) + by_n_idx = fx.Index(by_n) + + # Keep wave/lane arithmetic in i32. compute_global_swizzle() combines + # these values with i32 constants, so Index-typed coordinates would make + # arith.addi receive mixed operand types. + tx_i32 = fx.Int32(tx) + wave_id = tx_i32 // fx.Int32(WARP_SIZE) + lane = tx_i32 % fx.Int32(WARP_SIZE) + + # The utility mapping is identical to the previous manual staging: + # each step contributes one contiguous 16-byte vector per thread, while + # the global K coordinate is XOR-unswizzled for the physical LDS slot. + gl_off_a = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + gl_off_b = compute_global_swizzle(lane, wave_id, K, LOAD_PASSES_HALF, preshuffled=False) + a_g2s = G2SLoader(a_div, gl_off_a, LOAD_PASSES_HALF, a_f8_ir_t, wave_id) + b_g2s = G2SLoader(b_div, gl_off_b, LOAD_PASSES_HALF, b_f8_ir_t, wave_id) + s2r = S2RLoader(fx.Int32(0), 1) + + layout_lane16 = fx.make_layout((4, 16), (16, 1)) + coord_lane16 = fx.idx2crd(fx.Int32(lane), layout_lane16) + lane_div_16 = fx.get(coord_lane16, 0) + lane_mod_16 = fx.get(coord_lane16, 1) + + # C can exceed the signed-i32 element/byte offset range for large M*N. + # Bias the buffer descriptor base once per CTA using an index/i64 GEP, + # then store with only tile-local i32 offsets. This keeps the hot store + # instruction form unchanged while avoiding i32 wrap in buffer_store(). + c_n_idx_for_base = fx.Index(c_n) + c_tile_base_elems = bx_m_idx * c_n_idx_for_base + by_n_idx + c_tile_base_bytes = c_tile_base_elems * fx.Index(output_element_bytes) + c_rsrc = buffer_ops.create_buffer_resource( + C, + max_size=True, + base_byte_offset=c_tile_base_bytes, + ) + + PIN_ACC_BASE = 0 + + def _reg_list(prefix, start, end): + return ",".join(f"~{{{prefix}{r}}}" for r in range(start, end + 1)) + + def reserve_pinned_accumulators(): + # Reserve a fixed physical AGPR bank for all accumulators. In the + # SSA-lowered path, the compiler generated heavy AGPR <-> VGPR traffic, + # including v_accvgpr_mov/read sequences, s_nop stalls, and accumulator + # spills. Pinning each f32x4 accumulator to a stable AGPR range keeps the + # scaled MFMA accumulation in place and avoids those transfers and spills. + # + # ACCS_PER_WAVE = 64 accumulator objects and each object is f32x4, + # so the physical bank is exactly 64 * 4 = 256 AGPRs: a[0:255]. + clobbers = _reg_list("a", PIN_ACC_BASE, PIN_ACC_BASE + ACCS_PER_WAVE * 4 - 1) + llvm.InlineAsmOp( + None, + [], + "", + clobbers, + has_side_effects=True, + ) + + def zero_pinned_accumulators(): + for ai in range_constexpr(ACCS_PER_WAVE * 4): + llvm.InlineAsmOp( + None, + [], + f"v_accvgpr_write_b32 a[{PIN_ACC_BASE + ai}], 0", + f"~{{a{PIN_ACC_BASE + ai}}}", + has_side_effects=True, + ) + + def _inline_asm_i32(asm_string, constraints, operands=None): + op = llvm.InlineAsmOp( + T.i32, + operands or [], + asm_string, + constraints, + has_side_effects=True, + ) + return _one_i32_result(op) + + def _one_i32_result(op): + # Accept the result attribute names exposed by the supported MLIR Python bindings. + return getattr(op, "result", getattr(op, "res", op.results[0])) + + def _to_raw_inline_asm_operand(value): + # TODO: Replace arith._to_raw once FlyDSL exposes a supported public + # API for passing wrapped values to llvm.InlineAsmOp. _to_raw is + # deprecated, but remains heavily used internally by FlyDSL. + return arith._to_raw(value) + + def read_physical_accumulator_slot(slot_idx): + acc_pin = PIN_ACC_BASE + slot_idx * 4 + r0 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 0}]", "=v") + r1 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 1}]", "=v") + r2 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 2}]", "=v") + r3 = _inline_asm_i32(f"v_accvgpr_read_b32 $0, a[{acc_pin + 3}]", "=v") + return Vec.from_elements([r0, r1, r2, r3], fx.Int32).bitcast(fx.Float32) + + # As/Bs are MFMA-ready packed scale words: [K128, row] uint32. + # Each loaded dword already contains the four 16-row/16-col MFMA scale + # bytes for this lane's 64-row A/B half. The MFMA instruction selects + # the byte via op_sel/op_sel_hi, so there is intentionally no hot-loop + # byte extraction and no 0x01010101 broadcast here. + c_m_idx = fx.Index(c_m) + c_n_idx = fx.Index(c_n) + + def hot_loop_scheduler_q_refill_2n(): + # Steady-state Q1 schedule: eight chunks of one K+2 VMEM/LDS + # refill pass followed by two MFMAs. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q0_refill_a1_2n(): + # Steady-state Q0 schedule. Each chunk contains exactly: + # 1 K+2 VMEM/LDS refill pass + # 1 current-tile A-bottom K64 ds_read_b128 + # 2 current-tile Q0 MFMAs + # Repeated eight times, this distributes all eight A-bottom LDS reads + # across Q0 and maximizes their distance from reuse of that half-page. + for _ in range_constexpr(8): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(2) + + rocdl.sched_barrier(0) + + def hot_loop_scheduler_q_prefetch_4n(): + # Q2/Q3 carry-prefetch schedule used by both the steady loop and the + # penultimate tail tile. Each of eight chunks contains: + # 2 LDS reads for one complete next-tile A-top or B-left fragment + # 4 MFMAs using the current tile + for _ in range_constexpr(8): + rocdl.sched_dsrd(2) + rocdl.sched_mfma(4) + + rocdl.sched_barrier(0) + + def load_a_scale_row(k128, row): + packed = buffer_ops.buffer_load( + as_rsrc, + k128 * c_m_idx + bx_m_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + def load_b_scale_row(k128, row): + packed = buffer_ops.buffer_load( + bs_rsrc, + k128 * c_n_idx + by_n_idx + row, + vec_width=1, + dtype=T.i32, + ) + return packed + + def load_a_scale_subtile(k128, sm): + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(lane) + a_scale = load_a_scale_row(k128, a_row) + return (a_scale, a_scale, a_scale, a_scale) + + def load_b_scale_subtile(k128, sn): + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_row = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(lane) + b_scale = load_b_scale_row(k128, b_row) + return (b_scale, b_scale, b_scale, b_scale) + + def load_scale_tile(k128): + # Load all scale VGPRs needed by this wave for this K128 tile once. + # Return order: A-top, A-bottom, B-left, B-right. + return ( + load_a_scale_subtile(k128, 0), + load_a_scale_subtile(k128, 1), + load_b_scale_subtile(k128, 0), + load_b_scale_subtile(k128, 1), + ) + + def stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a): + # One pass writes 256 threads * 16 B = 4 KiB. Four passes fill one + # 128x128 half-page (16 KiB). Each half has its own LDS base. + global_base = (bx_m_idx + fx.Index(subtile * (BLOCK_M // 2))) * fx.Index(K) + k_base + a_g2s.load_one(lds_a[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b): + global_base = (by_n_idx + fx.Index(subtile * (BLOCK_N // 2))) * fx.Index(K) + k_base + b_g2s.load_one(lds_b[subtile], fx.Int32(global_base), pass_in_subtile) + + def stage_a_subtile(k_base, subtile, lds_a): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_a_subtile_pass(k_base, subtile, pass_in_subtile, lds_a) + + def stage_b_subtile(k_base, subtile, lds_b): + for pass_in_subtile in range_constexpr(LOAD_PASSES_HALF): + stage_b_subtile_pass(k_base, subtile, pass_in_subtile, lds_b) + + def load_frag_half_at_byte_base(lds_page, row_byte_base, half): + # Issue exactly one 16-byte LDS read for one K64 half of an MFMA operand. + # Keeping the halves separate allows steady-state Q0 to schedule one + # A-bottom ds_read_b128 in each refill/MFMA chunk. + k_col = reg_lds_k_col0 if half == 0 else reg_lds_k_col1 + return s2r.load_one(lds_page, fx.Int32(row_byte_base + k_col)) + + def pack_frag_halves(x0, x1): + return pack_i32x4_i32x8(x0, x1) + + def load_frag_at_byte_base(lds_page, row_byte_base): + # Default complete-fragment path used outside the dedicated Q0 schedule. + x0 = load_frag_half_at_byte_base(lds_page, row_byte_base, 0) + x1 = load_frag_half_at_byte_base(lds_page, row_byte_base, 1) + return pack_frag_halves(x0, x1) + + def load_b_frag(lds_b, local_row, half): + # B is [N, K]. Each 128-row half-page has a local row origin of 0. + half_row = local_row - fx.Index(half * (BLOCK_N // 2)) + return load_frag_at_byte_base(lds_b[half], half_row * fx.Index(BLOCK_K)) + + def _acc_idx(subtile_id, mi, ni): + return subtile_id * MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE + mi * MFMA_N_PER_SUBTILE + ni + + def pinned_mfma(acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Fixed physical accumulator bank, visible SSA A/B/scale operands. + # acc_idx maps directly to a[PIN_ACC_BASE + 4*acc_idx : +3]. + # The scale operands are MFMA-ready packed dwords. mi/ni choose + # which of the four bytes inside the A/B scale dword the MFMA uses. + acc_pin = PIN_ACC_BASE + acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$0, $1, " + f"a[{acc_pin}:{acc_pin + 3}], " + f"$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + (f"v,v,v,v,~{{a{acc_pin}}},~{{a{acc_pin + 1}}},~{{a{acc_pin + 2}}},~{{a{acc_pin + 3}}}"), + has_side_effects=True, + ) + + def pinned_final_mfma(dst_slot, old_acc_idx, a_frag, b_frag, a_scale, b_scale, mi, ni): + # Final-page form used by HK: destination and previous partial sum + # may be different AGPR ranges. Once old_acc_idx is consumed, its + # physical slot is dead and can be reused as a later destination. + dst_pin = PIN_ACC_BASE + dst_slot * 4 + old_pin = PIN_ACC_BASE + old_acc_idx * 4 + llvm.InlineAsmOp( + None, + [ + _to_raw_inline_asm_operand(a_frag), + _to_raw_inline_asm_operand(b_frag), + _to_raw_inline_asm_operand(a_scale), + _to_raw_inline_asm_operand(b_scale), + ], + ( + f"v_mfma_scale_f32_16x16x128_f8f6f4 " + f"a[{dst_pin}:{dst_pin + 3}], " + f"$0, $1, " + f"a[{old_pin}:{old_pin + 3}], " + f"$2, $3 " + f"op_sel:[{mi & 1},{ni & 1},0] " + f"op_sel_hi:[{mi >> 1},{ni >> 1},0] " + f"cbsz:{a_matrix_format} blgp:{b_matrix_format}" + ), + (f"v,v,v,v,~{{a{dst_pin}}},~{{a{dst_pin + 1}}},~{{a{dst_pin + 2}}},~{{a{dst_pin + 3}}}"), + has_side_effects=True, + ) + + def mfma_4n(acc_base, a_frag, a_scale, b0, b1, b2, b3, bs0, bs1, bs2, bs3): + """Emit four N-direction scaled MFMAs into fixed physical AGPR accumulators.""" + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, 1) + pinned_mfma(acc_base + 2, a_frag, b2, a_scale, bs2, mi, 2) + pinned_mfma(acc_base + 3, a_frag, b3, a_scale, bs3, mi, 3) + + def mfma_2n(acc_base, a_frag, a_scale, b0, b1, bs0, bs1, ni_base): + mi = (acc_base // MFMA_N_PER_SUBTILE) % MFMA_M_PER_SUBTILE + pinned_mfma(acc_base + 0, a_frag, b0, a_scale, bs0, mi, ni_base + 0) + pinned_mfma(acc_base + 1, a_frag, b1, a_scale, bs1, mi, ni_base + 1) + + def store_acc_vector_for_logical_idx(logical_acc_idx, acc): + subtile_id = logical_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = logical_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + row_base = subtile_m_idx * SUBTILE_M + fx.Index(mi * MFMA_M) + lane_div_16 * 4 + col = subtile_n_idx * SUBTILE_N + fx.Index(ni * MFMA_N) + lane_mod_16 + for ii in range_constexpr(4): + row = row_base + fx.Index(ii) + c_idx = row * fx.Index(c_n) + col + value = Vec(acc)[ii] + if output_dtype != torch.float32: + value = value.to(output_fx_dtype) + buffer_ops.buffer_store(value, c_rsrc, c_idx) + + + # Explicit register coordinates for HK-style four-quadrant mapping. + # BLOCK_M/BLOCK_N are 256x256. Four waves map to warp positions + # inside each 128x128 quadrant: + # cA: (warp_m, warp_n) + # cB: (warp_m, warp_n + 2) + # cC: (warp_m + 2, warp_n) + # cD: (warp_m + 2, warp_n + 2) + reg_k_col0 = lane_div_16 * 16 + reg_k_col1 = 64 + lane_div_16 * 16 + + # Every fragment row differs only by multiples of 16, so row % 16 is + # always lane_mod_16. Hoist the logical->physical XOR mapping once. + _, reg_lds_k_col0 = swizzle_128(lane_mod_16, reg_k_col0) + _, reg_lds_k_col1 = swizzle_128(lane_mod_16, reg_k_col1) + + reg_subtile_m_idx0 = wave_id // 2 + reg_subtile_n_idx0 = wave_id % 2 + + reserve_pinned_accumulators() + zero_pinned_accumulators() + + def load_b_subtile_ni_regs(lds_b, scale_tile, sn, ni): + # Fine-grained B register load for one 16-row N-direction MFMA slice. + # Return one packed B fragment and its matching scale operand. + subtile_n_idx = reg_subtile_n_idx0 + fx.Index(sn * 2) + b_scales = scale_tile[2] if sn == 0 else scale_tile[3] + + b_row_addr = subtile_n_idx * fx.Index(SUBTILE_N) + fx.Index(ni * MFMA_N) + lane_mod_16 + b_ni = load_b_frag(lds_b, b_row_addr, sn) + b_scale_ni = b_scales[ni] + return b_ni, b_scale_ni + + def load_b_subtile_regs(lds_b, scale_tile, sn): + b0, bs0 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 0) + b1, bs1 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 1) + b2, bs2 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 2) + b3, bs3 = load_b_subtile_ni_regs(lds_b, scale_tile, sn, 3) + return b0, b1, b2, b3, bs0, bs1, bs2, bs3 + + def load_a_subtile_mi_half(lds_a, sm, mi, half): + # One ds_read_b128 for one K64 half of one A MFMA slice. + subtile_m_idx = reg_subtile_m_idx0 + fx.Index(sm * 2) + a_row_addr = subtile_m_idx * fx.Index(SUBTILE_M) + fx.Index(mi * MFMA_M) + lane_mod_16 + half_row = a_row_addr - fx.Index(sm * (BLOCK_M // 2)) + row_byte_base = half_row * fx.Index(BLOCK_K) + return load_frag_half_at_byte_base(lds_a[sm], row_byte_base, half) + + def load_a_subtile_mi_regs(lds_a, scale_tile, sm, mi): + # Fine-grained A register load for one 16-row M-direction MFMA slice. + a_scales = scale_tile[0] if sm == 0 else scale_tile[1] + x0 = load_a_subtile_mi_half(lds_a, sm, mi, 0) + x1 = load_a_subtile_mi_half(lds_a, sm, mi, 1) + a_mi = pack_frag_halves(x0, x1) + a_scale_mi = a_scales[mi] + return a_mi, a_scale_mi + + def load_a_subtile_regs(lds_a, scale_tile, sm): + a0, as0 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 0) + a1, as1 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 1) + a2, as2 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 2) + a3, as3 = load_a_subtile_mi_regs(lds_a, scale_tile, sm, 3) + return a0, a1, a2, a3, as0, as1, as2, as3 + + def hk_one_k_with_refill( + k128, + cur_a, + cur_b, + next_a, + next_b, + refill_a, + refill_b, + a0_regs, + b0_regs, + cur_scales, + prev_refill_scales, + ): + # Scale invariant: + # cur_scales is HK MFMA-ready for K. + # prev_refill_scales is HK MFMA-ready for K+1. + # This iteration issues K+2 scale loads and returns them for the + # next steady iteration or final tail. + + # Wait only far enough for the current page; the next-page refill may remain in flight. + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + # Immediately issue MFMA-ready K+2 scale loads. + # They are returned for the next iteration without any in-kernel + # byte extraction or broadcast. + refill_scales = load_scale_tile(fx.Index(k128 + 2)) + next_scales_ready = prev_refill_scales + # A-top and B-left are both carried as complete 64-row register tiles, + # so their LDS half-pages can be refilled immediately. + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + # Refill the current ping-pong page with K+2, alternating A and B passes. + k_refill = fx.Index((k128 + 2) * BLOCK_K) + + # Q0: interleave the current tile's A-bottom LDS reads with K+2 + # refills and Q0 compute. Each complete A-bottom fragment is assembled + # from two independently scheduled K64 halves. + rocdl.sched_barrier(0) + a10_x0 = load_a_subtile_mi_half(cur_a, 1, 0, 0) + stage_a_subtile_pass(k_refill, 0, 0, refill_a) + mfma_2n(_acc_idx(0, 0, 0), a00, as00, b00, b01, bs00, bs01, 0) + + a10_x1 = load_a_subtile_mi_half(cur_a, 1, 0, 1) + stage_b_subtile_pass(k_refill, 0, 0, refill_b) + mfma_2n(_acc_idx(0, 0, 2), a00, as00, b02, b03, bs02, bs03, 2) + + a11_x0 = load_a_subtile_mi_half(cur_a, 1, 1, 0) + stage_a_subtile_pass(k_refill, 0, 1, refill_a) + mfma_2n(_acc_idx(0, 1, 0), a01, as01, b00, b01, bs00, bs01, 0) + + a11_x1 = load_a_subtile_mi_half(cur_a, 1, 1, 1) + stage_b_subtile_pass(k_refill, 0, 1, refill_b) + mfma_2n(_acc_idx(0, 1, 2), a01, as01, b02, b03, bs02, bs03, 2) + + a12_x0 = load_a_subtile_mi_half(cur_a, 1, 2, 0) + stage_a_subtile_pass(k_refill, 0, 2, refill_a) + mfma_2n(_acc_idx(0, 2, 0), a02, as02, b00, b01, bs00, bs01, 0) + + a12_x1 = load_a_subtile_mi_half(cur_a, 1, 2, 1) + stage_b_subtile_pass(k_refill, 0, 2, refill_b) + mfma_2n(_acc_idx(0, 2, 2), a02, as02, b02, b03, bs02, bs03, 2) + + a13_x0 = load_a_subtile_mi_half(cur_a, 1, 3, 0) + stage_a_subtile_pass(k_refill, 0, 3, refill_a) + mfma_2n(_acc_idx(0, 3, 0), a03, as03, b00, b01, bs00, bs01, 0) + + a13_x1 = load_a_subtile_mi_half(cur_a, 1, 3, 1) + stage_b_subtile_pass(k_refill, 0, 3, refill_b) + mfma_2n(_acc_idx(0, 3, 2), a03, as03, b02, b03, bs02, bs03, 2) + + hot_loop_scheduler_q0_refill_a1_2n() + + # Retire the eight distributed A-bottom LDS reads before K+2 refills + # overwrite the current page's A-bottom half-page. Keep this wait as + # late as possible to maximize read/compute overlap. + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10 = pack_frag_halves(a10_x0, a10_x1) + a11 = pack_frag_halves(a11_x0, a11_x1) + a12 = pack_frag_halves(a12_x0, a12_x1) + a13 = pack_frag_halves(a13_x0, a13_x1) + as10 = cur_scales[1][0] + as11 = cur_scales[1][1] + as12 = cur_scales[1][2] + as13 = cur_scales[1][3] + + rocdl.sched_barrier(0) + stage_b_subtile_pass(k_refill, 1, 0, refill_b) + mfma_2n(_acc_idx(1, 0, 0), a00, as00, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 0, refill_a) + mfma_2n(_acc_idx(1, 0, 2), a00, as00, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 1, refill_b) + mfma_2n(_acc_idx(1, 1, 0), a01, as01, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 1, refill_a) + mfma_2n(_acc_idx(1, 1, 2), a01, as01, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 2, refill_b) + mfma_2n(_acc_idx(1, 2, 0), a02, as02, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 2, refill_a) + mfma_2n(_acc_idx(1, 2, 2), a02, as02, b12, b13, bs12, bs13, 2) + + stage_b_subtile_pass(k_refill, 1, 3, refill_b) + mfma_2n(_acc_idx(1, 3, 0), a03, as03, b10, b11, bs10, bs11, 0) + + stage_a_subtile_pass(k_refill, 1, 3, refill_a) + mfma_2n(_acc_idx(1, 3, 2), a03, as03, b12, b13, bs12, bs13, 2) + hot_loop_scheduler_q_refill_2n() + + # Leave exactly the K+2 refill and scale loads outstanding. The following + # LDS reads consume the already-ready next page, not the page being refilled. + rocdl.sched_barrier(0) + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE + LOAD_PASSES_SCALES, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales_ready, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) + + return next_a0_regs, next_b0_regs, next_scales_ready, refill_scales + + def hk_one_k_tail_with_next(cur_a, cur_b, next_a, next_b, a0_regs, b0_regs, cur_scales, next_scales): + _barrier(vmcnt=2 * LOAD_PASSES_A_SUBTILE + 2 * LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + mfma_4n(_acc_idx(0, 0, 0), a00, as00, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 1, 0), a01, as01, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 2, 0), a02, as02, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + mfma_4n(_acc_idx(0, 3, 0), a03, as03, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + mfma_4n(_acc_idx(1, 0, 0), a00, as00, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 1, 0), a01, as01, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 2, 0), a02, as02, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + mfma_4n(_acc_idx(1, 3, 0), a03, as03, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + rocdl.sched_barrier(0) + _barrier(LOAD_PASSES_A_SUBTILE + LOAD_PASSES_B_SUBTILE, lgkmcnt=0) + rocdl.sched_barrier(0) + + next_a00, next_as00 = load_a_subtile_mi_regs(next_a, next_scales, 0, 0) + mfma_4n(_acc_idx(2, 0, 0), a10, as10, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a01, next_as01 = load_a_subtile_mi_regs(next_a, next_scales, 0, 1) + mfma_4n(_acc_idx(2, 1, 0), a11, as11, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a02, next_as02 = load_a_subtile_mi_regs(next_a, next_scales, 0, 2) + mfma_4n(_acc_idx(2, 2, 0), a12, as12, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_a03, next_as03 = load_a_subtile_mi_regs(next_a, next_scales, 0, 3) + mfma_4n(_acc_idx(2, 3, 0), a13, as13, b00, b01, b02, b03, bs00, bs01, bs02, bs03) + + next_b00, next_bs00 = load_b_subtile_ni_regs(next_b, next_scales, 0, 0) + mfma_4n(_acc_idx(3, 0, 0), a10, as10, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b01, next_bs01 = load_b_subtile_ni_regs(next_b, next_scales, 0, 1) + mfma_4n(_acc_idx(3, 1, 0), a11, as11, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b02, next_bs02 = load_b_subtile_ni_regs(next_b, next_scales, 0, 2) + mfma_4n(_acc_idx(3, 2, 0), a12, as12, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + next_b03, next_bs03 = load_b_subtile_ni_regs(next_b, next_scales, 0, 3) + mfma_4n(_acc_idx(3, 3, 0), a13, as13, b10, b11, b12, b13, bs10, bs11, bs12, bs13) + + hot_loop_scheduler_q_prefetch_4n() + + next_a0_regs = ( + next_a00, + next_a01, + next_a02, + next_a03, + next_as00, + next_as01, + next_as02, + next_as03, + ) + next_b0_regs = ( + next_b00, + next_b01, + next_b02, + next_b03, + next_bs00, + next_bs01, + next_bs02, + next_bs03, + ) + + return next_a0_regs, next_b0_regs + + def hk_one_k_final(cur_a, cur_b, a0_regs, b0_regs, cur_scales): + _barrier(vmcnt=0, lgkmcnt=0) + + a00, a01, a02, a03, as00, as01, as02, as03 = a0_regs + b00, b01, b02, b03, bs00, bs01, bs02, bs03 = b0_regs + + # Materialize the remaining final-page A/B fragments once. The + # subsequent schedule is entirely register/AGPR traffic. + b10, bs10 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 0) + b11, bs11 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 1) + b12, bs12 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 2) + b13, bs13 = load_b_subtile_ni_regs(cur_b, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a10, as10 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 0) + a11, as11 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 1) + a12, as12 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 2) + a13, as13 = load_a_subtile_mi_regs(cur_a, cur_scales, 1, 3) + + rocdl.sched_barrier(0) + _barrier(lgkmcnt=0) + rocdl.sched_barrier(0) + + a_frags = (a00, a01, a02, a03, a10, a11, a12, a13) + a_scales = (as00, as01, as02, as03, as10, as11, as12, as13) + b_frags = (b00, b01, b02, b03, b10, b11, b12, b13) + b_scales = (bs00, bs01, bs02, bs03, bs10, bs11, bs12, bs13) + + # Rolling final-page epilogue. + # + # Finalize accumulators in their own physical AGPR slots, but delay + # each AGPR read/store until several independent final MFMAs have + # been issued. + # + # MFMA 0, MFMA 1, MFMA 2, MFMA 3, drain 0, + # MFMA 4, drain 1, MFMA 5, drain 2, ... + # + # The buffer stores are only issued here; they may remain in flight + # while later MFMAs and accumulator drains continue. + FINAL_EPILOGUE_DEPTH = 4 + pending = [] + + for old_acc_idx in range_constexpr(ACCS_PER_WAVE): + subtile_id = old_acc_idx // (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + local_idx = old_acc_idx % (MFMA_M_PER_SUBTILE * MFMA_N_PER_SUBTILE) + sm = subtile_id // 2 + sn = subtile_id % 2 + mi = local_idx // MFMA_N_PER_SUBTILE + ni = local_idx % MFMA_N_PER_SUBTILE + + a_frag_idx = sm * MFMA_M_PER_SUBTILE + mi + b_frag_idx = sn * MFMA_N_PER_SUBTILE + ni + + # Final MFMA remains in-place. The logical accumulator's own + # AGPR slot is unique and cannot conflict with another pending + # result, so no ad-hoc physical-slot permutation is needed. + pinned_final_mfma( + old_acc_idx, + old_acc_idx, + a_frags[a_frag_idx], + b_frags[b_frag_idx], + a_scales[a_frag_idx], + b_scales[b_frag_idx], + mi, + ni, + ) + pending.append(old_acc_idx) + + # Drain the oldest completed result only after enough newer + # independent MFMAs have supplied the MFMA->AGPR-read spacing. + if len(pending) == FINAL_EPILOGUE_DEPTH: + drain_acc_idx = pending.pop(0) + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Flush the final results after all final-page MFMAs have issued. + for drain_acc_idx in pending: + acc = read_physical_accumulator_slot(drain_acc_idx) + store_acc_vector_for_logical_idx(drain_acc_idx, acc) + + # Prologue: stage K0/K1 data into ping-pong LDS pages. Scales are not staged in + # LDS: As/Bs are already MFMA-ready preshuffled packed uint32 [K128, row], + # and load_scale_tile returns the current wave's scale operands in VGPRs. + + # Load scales first, so that they become the oldest VMEM ops. + scales0 = load_scale_tile(fx.Index(0)) + scales1 = load_scale_tile(fx.Index(1)) + + stage_a_subtile(fx.Index(0), 0, lds_a0) + stage_b_subtile(fx.Index(0), 0, lds_b0) + stage_b_subtile(fx.Index(0), 1, lds_b0) + stage_a_subtile(fx.Index(0), 1, lds_a0) + + stage_a_subtile(fx.Index(BLOCK_K), 0, lds_a1) + stage_b_subtile(fx.Index(BLOCK_K), 0, lds_b1) + stage_b_subtile(fx.Index(BLOCK_K), 1, lds_b1) + stage_a_subtile(fx.Index(BLOCK_K), 1, lds_a1) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 4 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + # scales0 is already MFMA-ready; no byte extraction or broadcast is needed. + # Keep the hot loop consistent for k=0 and k>0: + # K0 is consumed directly. K1 MFMA-ready scales are carried as + # prev_refill_scales and become next_scales_ready at loop entry. + + # Seed the carried-register pipeline with K0 A-top. In later steady-state + # iterations, Q2/Q3 of the preceding iteration prefetch the next tile's + # A-top and B-left register tiles before their LDS half-pages are reused. + a0_regs = load_a_subtile_regs(lds_a0, scales0, 0) + + rocdl.sched_barrier(0) + _barrier(vmcnt=3 * LOAD_PASSES_A_SUBTILE + 3 * LOAD_PASSES_B_SUBTILE) + rocdl.sched_barrier(0) + + # Complete the K0 carried-register seed with B-left. + b0_regs = load_b_subtile_regs(lds_b0, scales0, 0) + + # Main HK loop: exactly one logical K128 per iteration. + # Even k consumes and refills LDS0; odd k does the same for LDS1. + # Scale tiles follow the same K128 progression but remain in VGPRs. + refill_scales = scales1 # K1 scales become the next ready scale tile at loop entry + for k128 in range_constexpr(NUM_K_TILES - 2): + if (k128 % 2) == 0: + a0_regs, b0_regs, scales1, refill_scales = hk_one_k_with_refill( + k128, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + scales0, + refill_scales, + ) + else: + a0_regs, b0_regs, scales0, refill_scales = hk_one_k_with_refill( + k128, + lds_a1, + lds_b1, + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + scales1, + refill_scales, + ) + + # Common two-page tail. The penultimate tile still uses the Q2/Q3 + # carry-prefetch scheduler to prepare A-top/B-left for the final tile, + # but it performs no K+2 data or scale refill. The final tile performs + # compute only. After the steady loop, a0_regs/b0_regs belong to the + # next tile to consume, while refill_scales belongs to the page most + # recently refilled; therefore tail page order depends on parity: + # even NUM_K_TILES: consume LDS0 then final LDS1 + # odd NUM_K_TILES: consume LDS1 then final LDS0 + if (NUM_K_TILES % 2) == 0: + scales1 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a0, + lds_b0, + lds_a1, + lds_b1, + a0_regs, + b0_regs, + scales0, + scales1, + ) + hk_one_k_final(lds_a1, lds_b1, a0_regs, b0_regs, scales1) + else: + scales0 = refill_scales + a0_regs, b0_regs = hk_one_k_tail_with_next( + lds_a1, + lds_b1, + lds_a0, + lds_b0, + a0_regs, + b0_regs, + scales1, + scales0, + ) + hk_one_k_final(lds_a0, lds_b0, a0_regs, b0_regs, scales0) + + @flyc.jit + def launch_gemm( + A: fx.Tensor, + As: fx.Tensor, + B: fx.Tensor, + Bs: fx.Tensor, + C: fx.Tensor, + c_m: fx.Int32, + c_n: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + # The integration only dispatches aligned shapes; no partial-tile masking exists. + grid_x = (c_m // BLOCK_M) * (c_n // BLOCK_N) + kernel_gemm( + A, + As, + B, + Bs, + C, + c_m, + c_n, + value_attrs={"rocdl.waves_per_eu": 1, "rocdl.flat_work_group_size": "256,256"}, + ).launch(grid=(grid_x, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch_gemm + +@functools.lru_cache(maxsize=None) +def _cached_launch( + K: int, + a_fp8_dtype: torch.dtype, + b_fp8_dtype: torch.dtype, + output_dtype: torch.dtype, +): + return _compile_kernel( + K, + a_fp8_dtype, + b_fp8_dtype, + output_dtype, + ) + + + +def do_gemm( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, + stream=None, +): + """Launch the K-specialized kernel with runtime M/N. + + A and B are shaped [M, K] and [N, K]. As/Bs are preshuffled packed + uint32 scale words shaped [K/128, M] and [K/128, N]. C is shaped [M, N]. + M and N are not hardcoded; K is used only to choose/cache the compile-time + specialized launch function. + """ + M_runtime, K_runtime = A.shape + N_runtime, Kb_runtime = B.shape + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + assert A.dtype in supported_fp8_dtypes, f"unsupported A MXFP8 dtype: {A.dtype}" + assert B.dtype in supported_fp8_dtypes, f"unsupported B MXFP8 dtype: {B.dtype}" + assert K_runtime == Kb_runtime, f"A.K={K_runtime} != B.K={Kb_runtime}" + if M_runtime % _BLOCK_M != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires M to be a multiple of {_BLOCK_M}, " + f"got M={M_runtime}" + ) + if N_runtime % _BLOCK_N != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires N to be a multiple of {_BLOCK_N}, " + f"got N={N_runtime}" + ) + if K_runtime % _BLOCK_K != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires K to be a multiple of {_BLOCK_K}, " + f"got K={K_runtime}" + ) + num_k_tiles = K_runtime // _BLOCK_K + if num_k_tiles < 4: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP8 GEMM requires at least 4 K{_BLOCK_K} tiles, " + f"got K={K_runtime} ({num_k_tiles} tiles)" + ) + expected_as = (K_runtime // _BLOCK_K, M_runtime) + expected_bs = (K_runtime // _BLOCK_K, N_runtime) + assert As.dtype == torch.int32, f"As dtype {As.dtype} != torch.int32 packed scales" + assert Bs.dtype == torch.int32, f"Bs dtype {Bs.dtype} != torch.int32 packed scales" + assert As.shape == expected_as, f"As shape {tuple(As.shape)} != {expected_as}" + assert Bs.shape == expected_bs, f"Bs shape {tuple(Bs.shape)} != {expected_bs}" + assert C.shape == (M_runtime, N_runtime), ( + f"C shape {tuple(C.shape)} != ({M_runtime}, {N_runtime})" + ) + assert C.dtype in (torch.float16, torch.bfloat16, torch.float32), ( + "C dtype must be torch.float16, torch.bfloat16, or torch.float32, " + f"got {C.dtype}" + ) + if stream is None: + stream = torch.cuda.current_stream() + # Match the Transformer Engine integration descriptor contract exactly. The optimized + # G2SLoader path consumes flat byte-addressed A/B tensors; scales and C are + # likewise passed as flat contiguous storage. Passing the original 2-D + # torch tensors changes the tensor descriptor/layout seen by + # make_fp8_buffer_tensor() and causes the loader's linear offsets to address + # the wrong elements. + A_arg = A.view(torch.uint8).contiguous().view(-1) + B_arg = B.view(torch.uint8).contiguous().view(-1) + As_arg = As.contiguous().view(-1) + Bs_arg = Bs.contiguous().view(-1) + C_arg = C.contiguous().view(-1) + + launch = _cached_launch( + int(K_runtime), + A.dtype, + B.dtype, + C.dtype, + ) + launch( + A_arg, + As_arg, + B_arg, + Bs_arg, + C_arg, + M_runtime, + N_runtime, + stream=stream, + ) + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "do_gemm", +] + + + +def mxfp8_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + D: torch.Tensor, + stream=None, +): + """Launch the fused MXFP8 kernel from canonical row-major operands. + + BLAS operand canonicalization, shape derivation, and output allocation are + intentionally owned by ``gemm_wrappers.py``. This function only validates + the MXFP8-specific scale contract, converts B to the HK [N, K] convention, + packs E8M0 scales, and launches the output-dtype-specialized 4-wave implementation. + """ + if a.ndim != 2 or b.ndim != 2: + raise ValueError( + f"FlyDSL MXFP8 expects rank-2 canonical operands, got " + f"a={tuple(a.shape)} and b={tuple(b.shape)}" + ) + + m, k = a.shape + kb, n = b.shape + if kb != k: + raise ValueError( + f"Incompatible canonical MXFP8 operands: " + f"{tuple(a.shape)} @ {tuple(b.shape)}" + ) + + supported_fp8_dtypes = ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + if a.dtype not in supported_fp8_dtypes or b.dtype not in supported_fp8_dtypes: + raise TypeError( + "FlyDSL MXFP8 expects E4M3 or E5M2 payloads independently, " + f"got a={a.dtype} and b={b.dtype}" + ) + + if a.device != b.device: + raise ValueError( + f"a and b must be on the same device, got {a.device} and {b.device}" + ) + if D.device != a.device: + raise ValueError(f"D must be on {a.device}, got {D.device}") + if tuple(D.shape) != (m, n): + raise ValueError( + f"D shape {tuple(D.shape)} does not match expected {(m, n)}" + ) + if D.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError( + "FlyDSL MXFP8 supports torch.float16, torch.bfloat16, or " + f"torch.float32 output, got {D.dtype}" + ) + if not D.is_contiguous(): + raise ValueError("FlyDSL MXFP8 requires contiguous output storage") + + if k % SCALE_GROUP_SIZE != 0: + raise ValueError( + f"K={k} must be divisible by MXFP8 scale group size " + f"{SCALE_GROUP_SIZE}" + ) + + # Canonical scale contract: + # a_scale [M, K/32] + # b_scale [K/32, N] + expected_a_scale = (m, k // SCALE_GROUP_SIZE) + expected_b_scale = (k // SCALE_GROUP_SIZE, n) + if tuple(a_scale.shape) != expected_a_scale: + raise ValueError( + f"a_scale shape {tuple(a_scale.shape)} != expected " + f"{expected_a_scale}" + ) + if tuple(b_scale.shape) != expected_b_scale: + raise ValueError( + f"b_scale shape {tuple(b_scale.shape)} != expected " + f"{expected_b_scale}" + ) + if a_scale.dtype != torch.uint8 or b_scale.dtype != torch.uint8: + raise TypeError("FlyDSL MXFP8 expects raw E8M0 scales as torch.uint8") + + # The HK core consumes B and its scales in row-oriented [N, K] form. + b_hk = b.transpose(0, 1).contiguous() + b_scale_rows = b_scale.transpose(0, 1).contiguous() + a_scale_hk = pack_mx32_scales_for_hk(a_scale) + b_scale_hk = pack_mx32_scales_for_hk(b_scale_rows) + + _debug( + f"private kernel inputs: a={tuple(a.shape)}, " + f"contiguous={a.is_contiguous()}; " + f"b_hk={tuple(b_hk.shape)}, contiguous={b_hk.is_contiguous()}; " + f"a_scale_hk={tuple(a_scale_hk.shape)}, " + f"b_scale_hk={tuple(b_scale_hk.shape)}, D={tuple(D.shape)}, " + f"D_dtype={D.dtype}" + ) + _debug("launching fused MXFP8 4-wave kernel") + + do_gemm( + a, + a_scale_hk, + b_hk, + b_scale_hk, + D.view(m, n), + stream=stream, + ) + + _debug("launch complete") + return D + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "SCALE_GROUP_SIZE", + "mxfp8_matmul", +]