diff --git a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td index 01f3e13c7..0ac2a5719 100644 --- a/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td +++ b/include/flydsl/Dialect/FlyROCDL/IR/MmaAtom.td @@ -85,17 +85,37 @@ def FlyROCDL_MmaOpGFX1250_WMMA : FlyROCDL_MmaOp<"MmaOpGFX1250_WMMA", "gfx1250.wm // false = unsigned, no clamp. "bool":$signA, "bool":$signB, - "bool":$clamp + "bool":$clamp, + // Intrinsic attributes forwarded to the ROCDL WMMA op: modC (I16 + // C-operand modifier) and reuseA/reuseB (I1 operand-reuse scheduler + // hints). Default 0/false; elided from the assembly when at defaults. + DefaultValuedParameter<"int32_t", "0">:$modC, + DefaultValuedParameter<"bool", "false">:$reuseA, + DefaultValuedParameter<"bool", "false">:$reuseB ); - let assemblyFormat = "`<` custom($m, $n, $k) `,` `(` $elemTyA `,` $elemTyB `)` `->` $elemTyAcc `,` `signA` `=` $signA `,` `signB` `=` $signB `,` `clamp` `=` $clamp `>`"; + let assemblyFormat = [{ + `<` custom($m, $n, $k) `,` `(` $elemTyA `,` $elemTyB `)` `->` $elemTyAcc + `,` `signA` `=` $signA `,` `signB` `=` $signB `,` `clamp` `=` $clamp + (`,` `modC` `=` $modC^)? + (`,` `reuseA` `=` $reuseA^)? + (`,` `reuseB` `=` $reuseB^)? `>` + }]; let builders = [ - // Back-compat: default sign/clamp to false (unsigned, no clamp). + // Back-compat: default sign/clamp/modC/reuse to false/0. TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc), [{ - return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, /*signA=*/false, /*signB=*/false, /*clamp=*/false); + return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, + /*signA=*/false, /*signB=*/false, /*clamp=*/false, + /*modC=*/0, /*reuseA=*/false, /*reuseB=*/false); }]>, TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc, "bool":$signA, "bool":$signB, "bool":$clamp), [{ - return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, signA, signB, clamp); + return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, + signA, signB, clamp, + /*modC=*/0, /*reuseA=*/false, /*reuseB=*/false); + }]>, + TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc, "bool":$signA, "bool":$signB, "bool":$clamp, "int32_t":$modC, "bool":$reuseA, "bool":$reuseB), [{ + return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc, + signA, signB, clamp, modC, reuseA, reuseB); }]> ]; let genVerifyDecl = 1; diff --git a/lib/Bindings/Python/FlyROCDLExtension.cpp b/lib/Bindings/Python/FlyROCDLExtension.cpp index 6263830bd..384c45bd1 100644 --- a/lib/Bindings/Python/FlyROCDLExtension.cpp +++ b/lib/Bindings/Python/FlyROCDLExtension.cpp @@ -66,17 +66,21 @@ struct PyMmaOpGFX1250_WMMAType : PyConcreteType { c.def_static( "get", [](int32_t m, int32_t n, int32_t k, PyType &elemTyA, PyType &elemTyB, PyType &elemTyAcc, - bool signA, bool signB, bool clamp, DefaultingPyMlirContext context) { + bool signA, bool signB, bool clamp, int32_t modC, bool reuseA, bool reuseB, + DefaultingPyMlirContext context) { return PyMmaOpGFX1250_WMMAType( - context->getRef(), - wrap(MmaOpGFX1250_WMMAType::get(m, n, k, unwrap(elemTyA), unwrap(elemTyB), - unwrap(elemTyAcc), signA, signB, clamp))); + context->getRef(), wrap(MmaOpGFX1250_WMMAType::get( + m, n, k, unwrap(elemTyA), unwrap(elemTyB), unwrap(elemTyAcc), + signA, signB, clamp, modC, reuseA, reuseB))); }, "m"_a, "n"_a, "k"_a, "elem_ty_a"_a, "elem_ty_b"_a, "elem_ty_acc"_a, "sign_a"_a = false, - "sign_b"_a = false, "clamp"_a = false, nb::kw_only(), "context"_a = nb::none(), + "sign_b"_a = false, "clamp"_a = false, "mod_c"_a = 0, "reuse_a"_a = false, + "reuse_b"_a = false, nb::kw_only(), "context"_a = nb::none(), "Create a MmaOpGFX1250_WMMAType with m, n, k dimensions and element types. " "sign_a / sign_b / clamp are integer-only (iu4 / iu8) controls (signed operands / " - "accumulator saturation); they must be false on the float paths."); + "accumulator saturation); they must be false on the float paths. " + "mod_c (I16 C-operand modifier, default 0), reuse_a / reuse_b (operand-reuse " + "scheduler hints, default false) are forwarded to the ROCDL intrinsic."); } }; diff --git a/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp b/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp index ca7d87a03..02b3eed8c 100644 --- a/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp +++ b/lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp @@ -116,7 +116,8 @@ Attribute MmaOpGFX1250_WMMAType::getThrValLayoutC() const { LogicalResult MmaOpGFX1250_WMMAType::verify(function_ref emitError, int32_t m, int32_t n, int32_t k, Type elemTyA, Type elemTyB, - Type elemTyAcc, bool signA, bool signB, bool clamp) { + Type elemTyAcc, bool signA, bool signB, bool clamp, + int32_t modC, bool reuseA, bool reuseB) { if (m != 16 || n != 16) return emitError() << "GFX1250 WMMA requires M=N=16, got " << m << "x" << n; @@ -235,19 +236,19 @@ enum class WmmaVariant { ModsAllReuse, ModsC, ModsABClamp, ModsIUClamp }; template static FailureOr emitWmmaSSA(OpBuilder &builder, Location loc, VectorType accTy, Value a, Value b, Value c, bool signA = false, bool signB = false, - bool clamp = false) { + bool clamp = false, int32_t modC = 0, bool reuseA = false, + bool reuseB = false) { Value res; if constexpr (Variant == WmmaVariant::ModsAllReuse) { - res = WmmaOp::create(builder, loc, accTy, a, b, ROCDL::WMMACModifier::none, c, - /*reuseA=*/false, /*reuseB=*/false) + res = WmmaOp::create(builder, loc, accTy, a, b, static_cast(modC), c, + reuseA, reuseB) .getResult(); } else if constexpr (Variant == WmmaVariant::ModsC) { - res = WmmaOp::create(builder, loc, accTy, a, b, ROCDL::WMMACModifier::none, c, - /*reuseA=*/false, /*reuseB=*/false) + res = WmmaOp::create(builder, loc, accTy, a, b, static_cast(modC), c, + reuseA, reuseB) .getResult(); } else if constexpr (Variant == WmmaVariant::ModsABClamp) { - res = WmmaOp::create(builder, loc, accTy, signA, a, signB, b, c, - /*reuseA=*/false, /*reuseB=*/false, clamp) + res = WmmaOp::create(builder, loc, accTy, signA, a, signB, b, c, reuseA, reuseB, clamp) .getResult(); } else { static_assert(Variant == WmmaVariant::ModsIUClamp); @@ -291,11 +292,14 @@ FailureOr MmaOpGFX1250_WMMAType::emitAtomCallSSA(OpBuilder &builder, Loca bool signA = getSignA(); bool signB = getSignB(); bool clamp = getClamp(); + int32_t modC = getModC(); + bool reuseA = getReuseA(); + bool reuseB = getReuseB(); #define DISPATCH_WMMA_SSA(M_, K_, PRED, OP, VARIANT) \ if (m == M_ && n == M_ && k == K_ && (PRED)) { \ return emitWmmaSSA(builder, loc, accTy, a, b, c, signA, \ - signB, clamp); \ + signB, clamp, modC, reuseA, reuseB); \ } #define DISPATCH_WMMA_SSA_FP8(K_, ACC_PRED, ACC_PREFIX) \ diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index cc3818081..e07b5b458 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -117,15 +117,19 @@ def MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None): def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): """Create an arch-appropriate WMMA atom. - Supported kwargs (integer paths only — iu8 / iu4): - sign_a (bool, default False): treat A operand as signed. - sign_b (bool, default False): treat B operand as signed. - clamp (bool, default False): saturate integer accumulator. + Supported kwargs: + sign_a (bool, default False): treat A operand as signed (iu8/iu4 only). + sign_b (bool, default False): treat B operand as signed (iu8/iu4 only). + clamp (bool, default False): saturate integer accumulator (iu8/iu4 only). + mod_c (int, default 0): I16 C-operand modifier (gfx1250 only). + reuse_a (bool, default False): operand-reuse scheduler hint (gfx1250 only). + reuse_b (bool, default False): operand-reuse scheduler hint (gfx1250 only). Forwarded to the arch-specific WMMA atom (MmaOpGFX11_WMMAType on gfx11, MmaOpGFX120X_WMMAType on gfx120x, MmaOpGFX1250_WMMAType on gfx1250); the - atom's verify() rejects them on the float (fp16/bf16/fp8) paths, where the - intrinsic has no such operands. Future WMMA ops for new architectures - should extend kwargs here rather than growing the positional signature. + atom's verify() rejects sign_a/sign_b/clamp on the float (fp16/bf16/fp8) + paths, where the intrinsic has no such operands. Future WMMA ops for new + architectures should extend kwargs here rather than growing the positional + signature. """ ty_ab = elem_ty_ab.ir_type if hasattr(elem_ty_ab, "ir_type") else elem_ty_ab if elem_ty_acc is None: @@ -155,6 +159,9 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): sign_a=bool(kwargs.get("sign_a", False)), sign_b=bool(kwargs.get("sign_b", False)), clamp=bool(kwargs.get("clamp", False)), + mod_c=int(kwargs.get("mod_c", 0)), + reuse_a=bool(kwargs.get("reuse_a", False)), + reuse_b=bool(kwargs.get("reuse_b", False)), ) if arch.startswith("gfx120"): return MmaOpGFX120X_WMMAType.get( diff --git a/tests/mlir/Conversion/wmma_gfx1250.mlir b/tests/mlir/Conversion/wmma_gfx1250.mlir index 984fa8eca..d01b316d4 100644 --- a/tests/mlir/Conversion/wmma_gfx1250.mlir +++ b/tests/mlir/Conversion/wmma_gfx1250.mlir @@ -49,3 +49,26 @@ func.func @test_wmma_iu4_signed_clamp( fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom i32, signA = true, signB = true, clamp = true>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () return } + +// ----- + +// bf16 WMMA with modC and reuse controls: verifies that modC / reuseA / +// reuseB are forwarded to the emitted rocdl.wmma op on the bf16 path. + +// CHECK-LABEL: @test_wmma_bf16_modc_reuse +func.func @test_wmma_bf16_modc_reuse( + %atom: !fly.mma_atom f32, signA = false, signB = false, clamp = false, modC = 1, reuseA = true, reuseB = true>>) { + %lay_ab = fly.static : !fly.layout<16:1> + %lay_cd = fly.static : !fly.layout<8:1> + %d = fly.memref.alloca(%lay_cd) : (!fly.layout<8:1>) -> !fly.memref + %a = fly.memref.alloca(%lay_ab) : (!fly.layout<16:1>) -> !fly.memref + %b = fly.memref.alloca(%lay_ab) : (!fly.layout<16:1>) -> !fly.memref + %c = fly.memref.alloca(%lay_cd) : (!fly.layout<8:1>) -> !fly.memref + + // CHECK: rocdl.wmma.f32.16x16x32.bf16 + // CHECK-SAME: modC = neg + // CHECK-SAME: reuseA = true + // CHECK-SAME: reuseB = true + fly.mma_atom_call(%atom, %d, %a, %b, %c) : (!fly.mma_atom f32, signA = false, signB = false, clamp = false, modC = 1, reuseA = true, reuseB = true>>, !fly.memref, !fly.memref, !fly.memref, !fly.memref) -> () + return +} diff --git a/tests/unit/test_gfx1250_atoms.py b/tests/unit/test_gfx1250_atoms.py index b8478c7c5..3b8e25c71 100644 --- a/tests/unit/test_gfx1250_atoms.py +++ b/tests/unit/test_gfx1250_atoms.py @@ -61,6 +61,26 @@ def test_wmma_scale_type_roundtrip(): assert ir.Type.parse(str(t_reuse)) == t_reuse +def test_wmma_type_modc_reuse_roundtrip(): + with _ctx(), ir.Location.unknown(): + from flydsl._mlir._mlir_libs._mlirDialectsFlyROCDL import MmaOpGFX1250_WMMAType + from flydsl._mlir.dialects import fly_rocdl # noqa: F401 + + bf16 = ir.BF16Type.get() + f32 = ir.F32Type.get() + + t_default = MmaOpGFX1250_WMMAType.get(16, 16, 32, bf16, bf16, f32) + assert "gfx1250.wmma<" in str(t_default) + # Defaults (modC=0, reuseA/reuseB=false) are elided from the printed form. + assert "modC" not in str(t_default) + assert ir.Type.parse(str(t_default)) == t_default + + t_modc = MmaOpGFX1250_WMMAType.get(16, 16, 32, bf16, bf16, f32, mod_c=1, reuse_a=True, reuse_b=True) + assert "modC = 1, reuseA = true, reuseB = true" in str(t_modc) + assert ir.Type.parse(str(t_modc)) == t_modc + assert t_modc != t_default + + def test_tdm2d_type_roundtrip(): with _ctx(), ir.Location.unknown(): from flydsl._mlir.dialects import fly_rocdl # noqa: F401 diff --git a/tests/unit/test_gfx1250_mma_atom_precision_check.py b/tests/unit/test_gfx1250_mma_atom_precision_check.py new file mode 100644 index 000000000..c131de86c --- /dev/null +++ b/tests/unit/test_gfx1250_mma_atom_precision_check.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Verify gfx1250 bf16 WMMA mma_atom_call correctness with different modC values. + +Single-wave (32 threads), single 16x16x32 WMMA: D = A @ B^T + modC(C). + modC=0 (none): D = A*B + C + modC=1 (neg): D = A*B - C + modC=2 (abs): D = A*B + |C| + modC=3 (neg_abs): D = A*B - |C| + +Each test compares GPU output against torch reference. +""" + +import pytest +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr +from flydsl.runtime.device import get_rocm_arch + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_arch = get_rocm_arch() or "" +_skip_not_gfx1250 = pytest.mark.skipif(not _arch.startswith("gfx1250"), reason=f"requires gfx1250, got {_arch}") + +WAVE_SIZE = 32 +M, N, K = 16, 16, 32 + + +def _make_wmma_kernel(mod_c): + """Create a WMMA kernel with a specific modC value (compile-time constant).""" + + @flyc.kernel(known_block_size=[WAVE_SIZE, 1, 1]) + def wmma_kernel(A_frag: fx.Tensor, B_frag: fx.Tensor, C_frag: fx.Tensor, D_frag: fx.Tensor): + tid = fx.thread_idx.x + + a_rmem = fx.make_rmem_tensor(16, fx.BFloat16) + b_rmem = fx.make_rmem_tensor(16, fx.BFloat16) + c_rmem = fx.make_rmem_tensor(8, fx.Float32) + + for i in fx.range_constexpr(16): + a_rmem[i] = A_frag[tid, i] + for i in fx.range_constexpr(16): + b_rmem[i] = B_frag[tid, i] + for i in fx.range_constexpr(8): + c_rmem[i] = C_frag[tid, i] + + atom = fx.make_mma_atom(fx.rocdl.WMMA(M, N, K, fx.BFloat16, fx.Float32, mod_c=const_expr(mod_c))) + fx.mma_atom_call(atom, c_rmem, a_rmem, b_rmem, c_rmem) + + for i in fx.range_constexpr(8): + D_frag[tid, i] = c_rmem[i] + + @flyc.jit + def launch( + A_frag: fx.Tensor, + B_frag: fx.Tensor, + C_frag: fx.Tensor, + D_frag: fx.Tensor, + stream: fx.Stream = fx.Stream(None), + ): + wmma_kernel(A_frag, B_frag, C_frag, D_frag).launch(grid=(1, 1, 1), block=(WAVE_SIZE, 1, 1), stream=stream) + + return launch + + +def _build_ab_fragments(A, B): + """Build per-thread A/B fragments matching gfx1250 WMMA bf16 register layout. + + Layout ((16,2),(8,2)):((1,128),(16,256)): + m = l % 16, k = (v // 8) * 16 + (l // 16) * 8 + (v % 8) + """ + A_frag = torch.zeros(WAVE_SIZE, 16, dtype=torch.bfloat16, device=A.device) + B_frag = torch.zeros(WAVE_SIZE, 16, dtype=torch.bfloat16, device=B.device) + for lane in range(WAVE_SIZE): + g = lane // 16 + m = lane % 16 + for v in range(16): + k = (v // 8) * 16 + g * 8 + (v % 8) + A_frag[lane, v] = A[m, k] + B_frag[lane, v] = B[m, k] + return A_frag, B_frag + + +def _pack_c_fragments(C): + """Pack M=16 x N=16 f32 matrix into per-thread C fragments. + + Layout ((16,2),(8)):((16,8),(1)): + m = (l // 16) * 8 + v, n = l % 16 + """ + C_frag = torch.zeros(WAVE_SIZE, 8, dtype=torch.float32, device=C.device) + for lane in range(WAVE_SIZE): + g = lane // 16 + n = lane % 16 + for v in range(8): + m = g * 8 + v + C_frag[lane, v] = C[m, n] + return C_frag + + +def _unpack_c_fragments(D_frag): + """Unpack per-thread D fragments back to M=16 x N=16 matrix.""" + D = torch.zeros(M, N, dtype=torch.float32, device=D_frag.device) + for lane in range(WAVE_SIZE): + g = lane // 16 + n = lane % 16 + for v in range(8): + m = g * 8 + v + D[m, n] = D_frag[lane, v] + return D + + +def _run_wmma(A, B, C_init, mod_c): + """Run WMMA with given modC and return the unpacked D matrix.""" + A_frag, B_frag = _build_ab_fragments(A, B) + C_frag = _pack_c_fragments(C_init) + D_frag = torch.zeros(WAVE_SIZE, 8, dtype=torch.float32, device=A.device) + + launch_fn = _make_wmma_kernel(mod_c) + launch_fn(A_frag, B_frag, C_frag, D_frag, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + + return _unpack_c_fragments(D_frag) + + +# modC=0 (none): D = A*B + C +@_skip_not_gfx1250 +def test_wmma_modc_none(): + torch.manual_seed(42) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + B = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + C = torch.randn(M, N, dtype=torch.float32, device="cuda") + + D_gpu = _run_wmma(A, B, C, mod_c=0) + D_ref = A.float() @ B.float().T + C + + max_diff = (D_gpu - D_ref).abs().max().item() + print(f"[modC=none] Max abs diff: {max_diff:.6e}") + assert torch.allclose(D_gpu, D_ref, atol=0.05, rtol=1e-3), f"modC=none mismatch, max diff = {max_diff}" + + +# modC=1 (neg): D = A*B + (-C) = A*B - C +@_skip_not_gfx1250 +def test_wmma_modc_neg(): + torch.manual_seed(42) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + B = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + C = torch.randn(M, N, dtype=torch.float32, device="cuda") + + D_gpu = _run_wmma(A, B, C, mod_c=1) + D_ref = A.float() @ B.float().T + (-C) + + max_diff = (D_gpu - D_ref).abs().max().item() + print(f"[modC=neg] Max abs diff: {max_diff:.6e}") + assert torch.allclose(D_gpu, D_ref, atol=0.05, rtol=1e-3), f"modC=neg mismatch, max diff = {max_diff}" + + +# modC=2 (abs): D = A*B + |C| +@_skip_not_gfx1250 +def test_wmma_modc_abs(): + torch.manual_seed(42) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + B = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + C = torch.randn(M, N, dtype=torch.float32, device="cuda") + + D_gpu = _run_wmma(A, B, C, mod_c=2) + D_ref = A.float() @ B.float().T + C.abs() + + max_diff = (D_gpu - D_ref).abs().max().item() + print(f"[modC=abs] Max abs diff: {max_diff:.6e}") + assert torch.allclose(D_gpu, D_ref, atol=0.05, rtol=1e-3), f"modC=abs mismatch, max diff = {max_diff}" + + +# modC=3 (neg_abs): D = A*B + (-(|C|)) = A*B - |C| +@_skip_not_gfx1250 +def test_wmma_modc_neg_abs(): + torch.manual_seed(42) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + B = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + C = torch.randn(M, N, dtype=torch.float32, device="cuda") + + D_gpu = _run_wmma(A, B, C, mod_c=3) + D_ref = A.float() @ B.float().T - C.abs() + + max_diff = (D_gpu - D_ref).abs().max().item() + print(f"[modC=neg_abs] Max abs diff: {max_diff:.6e}") + assert torch.allclose(D_gpu, D_ref, atol=0.05, rtol=1e-3), f"modC=neg_abs mismatch, max diff = {max_diff}" + + +if __name__ == "__main__": + test_wmma_modc_none() + test_wmma_modc_neg() + test_wmma_modc_abs() + test_wmma_modc_neg_abs() + print("ALL PASS")