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

use UOp.wmma everywhere (#17045)

Browse Source 
* use UOp.wmma everywhere

* fix
This commit is contained in:
George Hotz
2026-07-16 10:40:48 -07:00
committed by GitHub
parent 5a4156c5d1
commit 61e104bdfb
11 changed files with 30 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/amd_copy_matmul.py
+1 -1
View File
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
# NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
a_frag = a_frag.reshape(2, 8)[lane_m, :]
b_frag = b_frag.reshape(2, 8)[lane_m, :]
wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), ((16, 16, 16), 'AMD', 32))
wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), (16, 16, 16), 'AMD', 32)
acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
else:
# registers for LOCAL -> REG
extra/gemm/amd_flash_attention.py
+3 -3
View File
@@ -13,7 +13,7 @@ WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
LDS_PAD = 4 # pad LDS rows to reduce bank conflicts
WMMA_ARG = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32)
WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
LOG2E = math.log2(math.e)
def warp_shfl_xor(val, offset, lane):
@@ -97,7 +97,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), WMMA_ARG)
qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
S_reg = S_reg.after(qk_done)
@@ -158,7 +158,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), WMMA_ARG)
pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
# end KV tile loop
n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
extra/gemm/metal_uop_matmul.py
+1 -1
View File
@@ -28,7 +28,7 @@ def hand_spec_tc_cores():
acc = acc[1].set(0.0)
acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
out = UOp.wmma(a_tc, b_tc, acc_load, ((8, 8, 8), 'METAL', 32))
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
+2 -2
View File
@@ -137,7 +137,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
acc_load = acc_after[N_inner_loop, M_inner_loop]
# do WMMA
out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, ((16, 16, 32), 'AMD', 64))
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)
@@ -192,7 +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)])
out = UOp.wmma(A_in, B_in, acc_load, ((16, 16, 32), 'AMD', 64))
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
+1 -1
View File
@@ -60,7 +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
out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, ((16, 16, 32), 'AMD', 64))
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
+4 -4
View File
@@ -91,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.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64))
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)
@@ -121,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.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64))
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)
@@ -151,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.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64))
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)
@@ -181,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.wmma(a_in, b_in, d_in, (wmma_dims, 'AMD', 64))
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)
test/null/test_uop_graph.py
-60
View File
@@ -309,66 +309,6 @@ class TestUOpGraph(unittest.TestCase):
for uop, const in zip(uops, consts):
self.assertEqual(uop, const)
@unittest.skip("no longer testable standalone")
def test_wmma_vectorize_fold(self):
for i in [2, 4, 8]:
vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
var = UOp.variable("var", 0, 1, dtypes.half)
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(vec, var, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[0], acc)
self.assertEqual(len(uops), 2) # +1 for SINK
for i in [2, 4, 8]:
var = UOp.variable("var", 0, 1, dtypes.half)
vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(var, vec, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[0], acc)
self.assertEqual(len(uops), 2) # +1 for SINK
@unittest.skip("wmma is wrong here, it needs an arg")
def test_wmma_vectorize_no_fold(self):
for i in [4, 8]:
vec = UOp(Ops.STACK, dtypes.half,
tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(vec, var, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[-2], wmma) # -2 to skip SINK
for i in [4, 8]:
var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
vec = UOp(Ops.STACK, dtypes.half,
tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(var, vec, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[-2], wmma) # -2 to skip SINK
for i in [2, 4, 8]:
vec = UOp(Ops.STACK, dtypes.half,
tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(vec, var, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[-2], wmma) # -2 to skip SINK
for i in [2, 4, 8]:
var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
vec = UOp(Ops.STACK, dtypes.half,
tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
acc = UOp.variable('acc', 0, 1, dtypes.half)
wmma = UOp(Ops.WMMA, src=(var, vec, acc))
uops = to_uops_list([wmma])
self.assertEqual(uops[-2], wmma) # -2 to skip SINK
def test_cast_alu_fold(self):
d0 = UOp.param(0, dtypes.bool, (1,))
d1 = UOp.param(1, dtypes.int, (1,))
tinygrad/codegen/opt/postrange.py
+2 -2
View File
@@ -302,8 +302,8 @@ class Scheduler:
# do the reduce_axes always disappear? i think they don't
# they need to be moved into the WMMA srcs
wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.target.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
tc_uop = UOp(Ops.WMMA, src=(
srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2])), arg=wmma_arg, tag=1)
tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]),
tc.dims, self.ren.target.device, tc.threads, tag=1).replace(arg=wmma_arg)
# preserve extra reduces
reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
tinygrad/renderer/cstyle.py
+2 -2
View File
@@ -512,8 +512,8 @@ class HIPRenderer(CStyleLanguage):
type_map = {dtypes.bfloat16: "hip_bfloat16", dtypes.fp8e4m3: "hip_fp8", dtypes.fp8e5m2: "hip_bf8"}
extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) + PatternMatcher([
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
x.src[2]), arg=(*x.arg,)) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
# bfloat16 constant casting
(UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
])
tinygrad/renderer/llvmir.py
+11 -10
View File
@@ -253,29 +253,30 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc
if self.is_cdna:
self.extra_matcher += PatternMatcher([
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
x.src[2]), arg=x.arg) if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
])
if target.arch in {"gfx1100", "gfx1151"}:
self.extra_matcher += PatternMatcher([
(UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(UOp(Ops.WMMA,
(UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0)
for j in range(x.max_numel()*2)))), arg=(*x.arg[:6], (*x.arg[6][:2], ((0, x.max_numel()*2),)), *x.arg[7:])).index(i*2)
for j in range(x.max_numel()*2)))),
arg=(*x.arg[:6], (*x.arg[6][:2], ((0, x.max_numel()*2),)), *x.arg[7:])).index(i*2)
for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
(UPat(Ops.WMMA, name="x"), lambda x: UOp(Ops.WMMA,
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
(UPat(Ops.WMMA, name="x"), lambda x: x.replace(
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
])
if target.arch in {"gfx1200", "gfx1201"}:
self.extra_matcher += PatternMatcher([
(UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: UOp(Ops.WMMA,
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)), arg=x.arg)
(UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
.bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
(UPat(Ops.WMMA, name="x", dtype=dtypes.float),
lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
])
tinygrad/uop/ops.py
+3 -4
View File
@@ -601,12 +601,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
@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
def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tag=None):
dtype_in, dtype_out = a.dtype, acc.dtype
tc_upcast_axes = tuple(((i, s.shape[-1]),) for i,s in enumerate((a, b, acc)))
tc_upcast_axes = tuple(((i, s.shape[-1]),) if s._shape else () 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, ()))
return UOp(Ops.WMMA, src=(a, b, acc), arg=(name, dims, dtype_in, dtype_out, device, threads, tc_upcast_axes, ()), tag=tag)
def _rop(self, op:Ops, axis:tuple[int, ...]):
# NOTE: we don't allow reduce on 1s axis
axis = tuple(sorted(axis))