geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions 7 Activity

use wmma helper (#17038)

Browse Source
This commit is contained in:
George Hotz
2026-07-15 23:25:17 -07:00
committed by GitHub
parent 783042d216
commit 8b0dd870ce
5 changed files with 27 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/metal_uop_matmul.py
+2 -5
View File
@@ -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)
extra/gemm/mi350x_uop_matmul.py
+3 -5
View File
@@ -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))
extra/gemm/mi350x_uop_matmul_2.py
+2 -3
View File
@@ -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)
extra/thunder/tiny/tk/group.py
+13 -14
View File
@@ -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)
tinygrad/uop/ops.py
+7
View File
@@ -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))