From 8b0dd870ce99a2e4f9f10dba1dffb17b8c6b1437 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:25:17 -0700 Subject: [PATCH] use wmma helper (#17038) --- extra/gemm/metal_uop_matmul.py | 7 ++----- extra/gemm/mi350x_uop_matmul.py | 8 +++----- extra/gemm/mi350x_uop_matmul_2.py | 5 ++--- extra/thunder/tiny/tk/group.py | 27 +++++++++++++-------------- tinygrad/uop/ops.py | 7 +++++++ 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/extra/gemm/metal_uop_matmul.py b/extra/gemm/metal_uop_matmul.py index 3ec9694ca5..fdbb6d04c7 100644 --- a/extra/gemm/metal_uop_matmul.py +++ b/extra/gemm/metal_uop_matmul.py @@ -1,5 +1,5 @@ from tinygrad import UOp, dtypes -from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace +from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace from extra.gemm.amd_uop_matmul import test_matmul N = 2048 @@ -27,11 +27,8 @@ def hand_spec_tc_cores(): acc = acc[0].set(0.0) acc = acc[1].set(0.0) - # TODO: make this simple - wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ()) - acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1]) - out = UOp(Ops.WMMA, dtypes.float, (a_tc, b_tc, acc_load), arg=wmma_arg) + out = UOp.wmma(a_tc, b_tc, acc_load, ((8, 8, 8), 'METAL', 32)) end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk) diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index c6fc56e625..17181e2448 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0" from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters from tinygrad.helpers import DEBUG, getenv from tinygrad.dtype import AddrSpace -from tinygrad.uop.ops import AxisType, KernelInfo, Ops +from tinygrad.uop.ops import AxisType, KernelInfo WARP_SIZE = 64 @@ -137,8 +137,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: acc_load = acc_after[N_inner_loop, M_inner_loop] # do WMMA - wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ()) - out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg) + out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, ((16, 16, 32), 'AMD', 64)) # store back the acc acc_store = acc[N_inner_loop, M_inner_loop].store(out) @@ -193,8 +192,7 @@ acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l) # do the wmma acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)]) -wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ()) -out = UOp(Ops.WMMA, dtypes.float, (A_in, B_in, acc_load), arg=wmma_arg) +out = UOp.wmma(A_in, B_in, acc_load, ((16, 16, 32), 'AMD', 64)) # store back the acc acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop)) diff --git a/extra/gemm/mi350x_uop_matmul_2.py b/extra/gemm/mi350x_uop_matmul_2.py index 03af1e547f..271b3de2e4 100644 --- a/extra/gemm/mi350x_uop_matmul_2.py +++ b/extra/gemm/mi350x_uop_matmul_2.py @@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0" from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters from tinygrad.helpers import DEBUG, getenv from tinygrad.dtype import AddrSpace -from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops +from tinygrad.uop.ops import AxisType, KernelInfo WARP_SIZE = 64 @@ -60,8 +60,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...] acc_load = acc_after[N_inner_loop, M_inner_loop] # do WMMA - wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ()) - out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg) + out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, ((16, 16, 32), 'AMD', 64)) # store back the acc acc_store = acc[N_inner_loop, M_inner_loop].store(out) diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index 394e73dfb4..c39dde61db 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -1,7 +1,6 @@ import math from typing import cast, Callable -from tinygrad import dtypes -from tinygrad.uop.ops import AxisType, UOp, Ops +from tinygrad.uop.ops import AxisType, UOp from tinygrad.dtype import AddrSpace from tinygrad.helpers import prod @@ -75,9 +74,9 @@ class Group: a_base_shape = cast(RT, a).base_shape if a_base_shape.cols == 16: - wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 16) elif a_base_shape.cols == 32: - wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 32) else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}") for height in self.ker.range(c.shape[-3], track=False): @@ -92,7 +91,7 @@ class Group: else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}") d_in = UOp.stack(*[c[height, width, i] for i in range(4)]) - out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg) + out = UOp.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64)) c_i = [c[height, width, i].store(out.index(i)) for i in range(4)] c_store = UOp.group(*c_i).end(height, width, inner) @@ -105,9 +104,9 @@ class Group: a_base_shape = cast(RT, a).base_shape if a_base_shape.cols == 16: - wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 16) elif a_base_shape.cols == 32: - wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 32) else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}") for height in self.ker.range(c.shape[-3], track=False): @@ -122,7 +121,7 @@ class Group: else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}") d_in = UOp.stack(*[c[height, width, i] for i in range(4)]) - out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg) + out = UOp.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64)) c_i = [c[height, width, i].store(out.index(i)) for i in range(4)] c_store = UOp.group(*c_i).end(height, width, inner) @@ -135,9 +134,9 @@ class Group: a_base_shape = cast(RT, a).base_shape if a_base_shape.cols == 16: - wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 16) elif a_base_shape.cols == 32: - wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 32) else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}") for height in self.ker.range(c.shape[-3], track=False): @@ -152,7 +151,7 @@ class Group: else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}") d_in = UOp.stack(*[c[height, width, i] for i in range(4)]) - out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg) + out = UOp.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64)) c_i = [c[height, width, i].store(out.index(i)) for i in range(4)] c_store = UOp.group(*c_i).end(height, width, inner) @@ -165,9 +164,9 @@ class Group: a_base_shape = cast(RT, a).base_shape if a_base_shape.cols == 16: - wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 16) elif a_base_shape.cols == 32: - wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore + wmma_dims = (16, 16, 32) else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}") for height in self.ker.range(c.shape[-3], track=False): @@ -182,7 +181,7 @@ class Group: else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}") d_in = UOp.stack(*[c[height, width, i] for i in range(4)]) - out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg) + out = UOp.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64)) c_i = [c[height, width, i].store(out.index(i)) for i in range(4)] c_store = UOp.group(*c_i).end(height, width, inner) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 49c53575b6..8d1eefc844 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -600,6 +600,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs) @staticmethod def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name) + @staticmethod + def wmma(a:UOp, b:UOp, acc:UOp, arg:tuple[tuple[int, int, int], str, int]): + dims, device, threads = arg + dtype_in, dtype_out = a.dtype, acc.dtype + tc_upcast_axes = tuple(((i, s.shape[-1]),) for i,s in enumerate((a, b, acc))) + name = f"WMMA_{'_'.join(map(str, dims))}_{dtype_in.name}_{dtype_out.name}" + return UOp(Ops.WMMA, src=(a, b, acc), arg=(name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ())) def _rop(self, op:Ops, axis:tuple[int, ...]): # NOTE: we don't allow reduce on 1s axis axis = tuple(sorted(axis))