diff --git a/extra/gemm/amd_copy_matmul.py b/extra/gemm/amd_copy_matmul.py index af0ccb58a5..5d9b4c701a 100644 --- a/extra/gemm/amd_copy_matmul.py +++ b/extra/gemm/amd_copy_matmul.py @@ -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 diff --git a/extra/gemm/amd_flash_attention.py b/extra/gemm/amd_flash_attention.py index 6f34687b98..73740d8165 100644 --- a/extra/gemm/amd_flash_attention.py +++ b/extra/gemm/amd_flash_attention.py @@ -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) diff --git a/extra/gemm/metal_uop_matmul.py b/extra/gemm/metal_uop_matmul.py index fdbb6d04c7..cfbb48a4ac 100644 --- a/extra/gemm/metal_uop_matmul.py +++ b/extra/gemm/metal_uop_matmul.py @@ -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) diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index 17181e2448..cc3c039e10 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -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)) diff --git a/extra/gemm/mi350x_uop_matmul_2.py b/extra/gemm/mi350x_uop_matmul_2.py index 271b3de2e4..26bb43ff4d 100644 --- a/extra/gemm/mi350x_uop_matmul_2.py +++ b/extra/gemm/mi350x_uop_matmul_2.py @@ -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) diff --git a/extra/thunder/tiny/tk/group.py b/extra/thunder/tiny/tk/group.py index c39dde61db..08f847dbd3 100644 --- a/extra/thunder/tiny/tk/group.py +++ b/extra/thunder/tiny/tk/group.py @@ -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) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 6080e78f35..4da248df28 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -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,)) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 0dfec14aec..92fb81d209 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -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] diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 9d8dbab035..c4074f484d 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -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))), ]) diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index d38bedab35..9057906f68 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -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) ]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8d1eefc844..c1bbc09772 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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))