From 9dbcddc64910ecccfa8adc44897c623e39e654eb Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 1 Sep 2026 10:48:59 -0400 Subject: [PATCH] UPCAST and LOCAL are SPLIT (#17889) --- extra/gemm/max_matmul.py | 1 - extra/gemm/tinygrad_nv_matmul.py | 15 +- extra/optimization/test_beam_search.py | 7 +- test/backend/test_linearizer.py | 23 +-- test/backend/test_linearizer_dumb.py | 2 +- test/backend/test_opt_gemm.py | 11 +- test/backend/test_quantize_onnx.py | 18 ++- test/backend/test_uops.py | 2 +- test/device/test_hcq.py | 3 +- test/null/test_linearizer_rewrite.py | 8 +- test/null/test_process_replay.py | 3 +- test/null/test_uops_stats.py | 18 ++- test/opt/test_gen_float4.py | 30 ++-- test/opt/test_kernel_opts.py | 215 ++++++++++++++----------- test/opt/test_tensor_cores.py | 8 +- tinygrad/codegen/opt/__init__.py | 5 +- tinygrad/codegen/opt/heuristic.py | 31 ++-- tinygrad/codegen/opt/postrange.py | 35 ++-- tinygrad/codegen/opt/search.py | 12 +- 19 files changed, 245 insertions(+), 202 deletions(-) diff --git a/extra/gemm/max_matmul.py b/extra/gemm/max_matmul.py index 5a41fe17f8..9e3be02a6f 100644 --- a/extra/gemm/max_matmul.py +++ b/extra/gemm/max_matmul.py @@ -113,7 +113,6 @@ if __name__ == "__main__": } elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float: print("Using CUDA and generated hcopt") - # [Opt(op=OptOps.TC, axis=0, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.LOCAL, axis=1, amt=4)] prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read())) args = (c, a, b) kwargs = { diff --git a/extra/gemm/tinygrad_nv_matmul.py b/extra/gemm/tinygrad_nv_matmul.py index afb9d99408..a3826ee76d 100644 --- a/extra/gemm/tinygrad_nv_matmul.py +++ b/extra/gemm/tinygrad_nv_matmul.py @@ -1,6 +1,7 @@ from tinygrad import Tensor, dtypes, Context from tinygrad.helpers import getenv from tinygrad.codegen.opt import Opt, OptOps +from tinygrad.uop.ops import AxisType from tinygrad.engine.realize import run_linear from dataclasses import replace @@ -13,17 +14,17 @@ if __name__ == "__main__": C = A.matmul(B) if getenv("GEMV"): opts = [ - Opt(op=OptOps.UPCAST, axis=1, amt=8), - Opt(op=OptOps.LOCAL, axis=1, amt=32), + Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)), + Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)), ] else: opts = [ Opt(op=OptOps.TC, axis=0, amt=0), - Opt(op=OptOps.UPCAST, axis=0, amt=4), - Opt(op=OptOps.UPCAST, axis=1, amt=8), - Opt(op=OptOps.LOCAL, axis=0, amt=2), - Opt(op=OptOps.LOCAL, axis=1, amt=2), - Opt(op=OptOps.LOCAL, axis=0, amt=2), + Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)), + Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.LOCAL)), + Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)), ] linear = C.schedule_linear() call = linear.src[-1] diff --git a/extra/optimization/test_beam_search.py b/extra/optimization/test_beam_search.py index c3527bd45a..382fd29005 100644 --- a/extra/optimization/test_beam_search.py +++ b/extra/optimization/test_beam_search.py @@ -89,7 +89,8 @@ class TestBeamSearch(unittest.TestCase): s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1))) up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)]) actions = get_kernel_actions(s, include_0=False, max_up=int(up)) - upcasted = [s for s in actions.values() if any(opt.op is OptOps.UPCAST for opt in s.applied_opts)] + upcasted = [s for s in actions.values() if any(o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL) + for o in s.applied_opts)] assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none" def test_max_up(self): @@ -98,8 +99,8 @@ class TestBeamSearch(unittest.TestCase): s = Scheduler(ast, Device[Device.DEFAULT].renderer) for max_up in (2, 4): actions = get_kernel_actions(s, include_0=False, max_up=max_up) - for up_opts in [s.applied_opts for s in actions.values() if any(opt.op is OptOps.UPCAST for opt in s.applied_opts)]: - assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0 + up_opts = [o for s in actions.values() for o in s.applied_opts if o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)] + assert len([opt for opt in up_opts if opt.arg[0] > max_up]) == 0 and len([op for op in up_opts if op.arg[0] <= max_up]) > 0 if __name__ == '__main__': unittest.main() diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 93ac7bb783..82e79d8607 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -127,7 +127,7 @@ class TestLinearizer(unittest.TestCase): # these are of size 3 to avoid float4 coalesce r = a[:-1] + a[1:] - uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) num_loads = len([uop for uop in uops if uop.op is Ops.LOAD]) assert num_loads <= 4, "more load uops than needed" @@ -140,7 +140,7 @@ class TestLinearizer(unittest.TestCase): a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize() r = a.expand([2]) + b.expand([2]) - uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU]) assert num_ops <= 1, "more alu uops than needed" @@ -151,7 +151,8 @@ class TestLinearizer(unittest.TestCase): r = Tensor.conv2d(x,w,padding=1).relu() uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], - [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=1, arg=(0, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) accs = [u for u in uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG] stores = [u for u in uops if u.op is Ops.STORE] assert len(accs) == 0 # it's removed now @@ -162,7 +163,7 @@ class TestLinearizer(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU") def test_upcast_with_locals_cpu(self): out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous() - prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]), + prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.SPLIT, axis=0, arg=(4, AxisType.LOCAL))]), renderer=Device[Device.DEFAULT].renderer) self.assertEqual(len(prg.src[2].arg.split("for")), 5) @@ -173,7 +174,8 @@ class TestLinearizer(unittest.TestCase): def test_upcast_with_locals(self): x, y = Tensor.rand(1,128), Tensor.rand(128, 128) r = (x@y).relu() - opts_to_apply = [Opt(op=OptOps.LOCAL, axis=1, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)] + opts_to_apply = [Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.GROUP_REDUCE)), Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.LOCAL)), + Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))] program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer) stores = [u for u in tuple(program.src[1].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG] @@ -189,7 +191,7 @@ class TestLinearizer(unittest.TestCase): def test_zero_fold(self): a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize() r = Tensor.stack(a, b) - uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), + uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU]) assert num_ops == 0, "more alu uops than needed" @@ -232,7 +234,7 @@ class TestLinearizer(unittest.TestCase): def test_simple_unroll_no_between_phi_dependencies(self): x, y = Tensor.empty(64, 64), Tensor.empty(64, 64) r = (x@y).relu() - opt = [Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UPCAST, 0, 4)] + opt = [Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))] ast = helper_linearizer_opt(r, [opt]) # the uops graph is reg BUFFER -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src) @@ -342,8 +344,9 @@ class TestLinearizer(unittest.TestCase): def test_grouped_store_locals_and_globals(self): x, y = Tensor.empty(64, 64), Tensor.empty(64, 64) out = x@y - opt = [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 3, 8, top=True), - Opt(OptOps.UPCAST, 3, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces + opt = [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 3, (8, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 3, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST))] # upcast accs in both reduces ast = helper_linearizer_opt(out, opts=[opt]) def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src]) uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src) @@ -383,7 +386,7 @@ class TestLinearizer(unittest.TestCase): def test_two_grouped_stores_local(self): # GROUP_REDUCE on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier a = Tensor.rand(32, 32).realize() - opts = [Opt(OptOps.LOCAL, 3, 4), Opt(OptOps.LOCAL, 5, 4)] + opts = [Opt(OptOps.SPLIT, 3, (4, AxisType.GROUP_REDUCE)), Opt(OptOps.SPLIT, 5, (4, AxisType.GROUP_REDUCE))] ast = helper_linearizer_opt(single_kernel_softmax(a), [opts]) uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2) diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index 63908f24af..164f9df193 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -24,7 +24,7 @@ class TestLinearizerFailure(unittest.TestCase): c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True))) c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD) c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3) - ast = c12.sink(arg=KernelInfo(name='test', applied_opts=(Opt(op=OptOps.LOCAL, axis=4, arg=16),), opts_to_apply=None)) + ast = c12.sink(arg=KernelInfo(name='test', applied_opts=(Opt(op=OptOps.SPLIT, axis=4, arg=(16, AxisType.GROUP_REDUCE)),), opts_to_apply=None)) _ = to_program(ast, Device["METAL"].renderer) if __name__ == '__main__': diff --git a/test/backend/test_opt_gemm.py b/test/backend/test_opt_gemm.py index 244e3df889..c248254554 100644 --- a/test/backend/test_opt_gemm.py +++ b/test/backend/test_opt_gemm.py @@ -4,7 +4,7 @@ from tinygrad import Tensor from tinygrad.helpers import get_single_element from tinygrad.codegen.opt import Opt, OptOps from tinygrad.engine.realize import run_linear -from tinygrad.uop.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp, AxisType from test.helpers import replace_opts class TestOptGemm(unittest.TestCase): @@ -26,20 +26,21 @@ class TestOptGemm(unittest.TestCase): np.testing.assert_allclose(self.res, test, atol=1e-4) def test_gemm_unrolled_permute_l_44(self): - opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(4, AxisType.UPCAST))] self._test_gemm_unrolled_permute_l(opts) def test_gemm_unrolled_permute_l_424(self): # was failing with LLVM - opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))] self._test_gemm_unrolled_permute_l(opts) def test_gemm_unrolled_permute_l_42(self): - opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)] + opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST))] self._test_gemm_unrolled_permute_l(opts) def test_gemm_unrolled_permute_l_22(self): - opts = [Opt(op=OptOps.UPCAST, axis=0, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2)] + opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST))] self._test_gemm_unrolled_permute_l(opts) if __name__ == '__main__': diff --git a/test/backend/test_quantize_onnx.py b/test/backend/test_quantize_onnx.py index f743ee425b..6848190b59 100644 --- a/test/backend/test_quantize_onnx.py +++ b/test/backend/test_quantize_onnx.py @@ -2,7 +2,7 @@ import numpy as np import tempfile, unittest from tinygrad import Tensor, Context, Device, dtypes, UOp -from tinygrad.uop.ops import Ops +from tinygrad.uop.ops import Ops, AxisType from tinygrad.dtype import AddrSpace from tinygrad.codegen.opt import Opt, OptOps from tinygrad.engine.realize import run_linear @@ -98,7 +98,7 @@ class TestQuantizeOnnx(unittest.TestCase): X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8)) W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8)) out = X.conv2d(W, dtype=X.dtype) - opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UPCAST, axis=3, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))] sexec(out, opts) def test_prequant_gemm(self): @@ -106,7 +106,7 @@ class TestQuantizeOnnx(unittest.TestCase): X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)) W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)) out = X.matmul(W, dtype=X.dtype) - opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UPCAST, axis=3, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))] sexec(out, opts) # TODO: this has to work @@ -116,7 +116,7 @@ class TestQuantizeOnnx(unittest.TestCase): W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi)) # this divide is interesting and forces the accumulator to actually be an int out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8") - opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UPCAST, axis=3, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))] sexec(out, opts) def test_prequant_gemm_handcode(self): @@ -200,10 +200,11 @@ class TestQuantizeOnnx(unittest.TestCase): self.test_prequant_gemm_intacc(np.uint8, np.int8, src) def test_prequant_gemm_intacc_32(self): - opts = [Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=0)] + opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(0, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=3, arg=(0, AxisType.UNROLL))] self.test_prequant_gemm_intacc(np.uint8, np.int8, N=32, opts=opts) def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128, - opts=[Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UPCAST, axis=2, arg=4)]) + opts=[Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=2, arg=(4, AxisType.UNROLL))]) def test_prequant_gemm_intacc_256(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=256) def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None): X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize() @@ -212,7 +213,8 @@ class TestQuantizeOnnx(unittest.TestCase): out = (X.int().matmul(W.int())//1000) if clip: out = out.clip(tg_dtype.min, tg_dtype.max) out = out.cast(tg_dtype) - opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UPCAST, axis=3, arg=4)] if opts is None else opts + opts = [Opt(op=OptOps.SPLIT, axis=1, arg=(128, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=3, arg=(4, AxisType.UNROLL))] if opts is None else opts sexec(out, opts, replace_src, run_count=1) tout = out.numpy() mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000) @@ -233,7 +235,7 @@ class TestQuantizeOnnx(unittest.TestCase): #out = X.cast(dtypes.int) @ W.cast(dtypes.int) #out = X @ W out = X.matmul(W, dtype=X.dtype) - opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UPCAST, axis=2, arg=4)] + opts = [Opt(op=OptOps.SPLIT, axis=0, arg=(128, AxisType.UPCAST)), Opt(op=OptOps.SPLIT, axis=2, arg=(4, AxisType.UNROLL))] sexec(out, opts) if __name__ == "__main__": diff --git a/test/backend/test_uops.py b/test/backend/test_uops.py index 48a2c98b7a..e5d0c8c9ab 100644 --- a/test/backend/test_uops.py +++ b/test/backend/test_uops.py @@ -271,7 +271,7 @@ class TestAssembly(unittest.TestCase): b = Tensor.empty(1024) c = (a*b).sum() ast = c.schedule_linear().src[-1].src[0] - opts_to_apply = [Opt(OptOps.UPCAST, 0, 4)] + opts_to_apply = [Opt(OptOps.SPLIT, 0, (4, AxisType.UNROLL))] ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) program = to_program(ast, Device[Device.DEFAULT].renderer) uops = tuple(program.src[1].src) diff --git a/test/device/test_hcq.py b/test/device/test_hcq.py index b8390d0b86..55c4812669 100644 --- a/test/device/test_hcq.py +++ b/test/device/test_hcq.py @@ -9,6 +9,7 @@ from tinygrad.runtime.support.system import PCIIfaceBase from tinygrad.engine.realize import get_runtime from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps +from tinygrad.uop.ops import AxisType from tinygrad import Variable MOCKGPU = DEV.interface.startswith("MOCK") @@ -167,7 +168,7 @@ class TestHCQ(unittest.TestCase): b = a + 1 si = b.schedule_linear().src[-1] - prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer) + prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(3, AxisType.LOCAL)) for _ in range(3)]), TestHCQ.d0.renderer) runtime = get_runtime(Device.DEFAULT, prg) zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated() diff --git a/test/null/test_linearizer_rewrite.py b/test/null/test_linearizer_rewrite.py index 1946629912..12890cf6ec 100644 --- a/test/null/test_linearizer_rewrite.py +++ b/test/null/test_linearizer_rewrite.py @@ -2,7 +2,7 @@ import unittest from tinygrad import Tensor, Context, Device from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps -from tinygrad.uop.ops import KernelInfo +from tinygrad.uop.ops import KernelInfo, AxisType class TestLinearizerRewrite(unittest.TestCase): def test_reduction(self): @@ -11,8 +11,8 @@ class TestLinearizerRewrite(unittest.TestCase): with Context(SPLIT_REDUCEOP=0): si = out.schedule_linear().src[-1] opts_to_apply = [] - opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) - opts_to_apply.append(Opt(OptOps.UPCAST, 2, 4)) + opts_to_apply.append(Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))) + opts_to_apply.append(Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL))) ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) prg = to_program(ast, Device["CPU"].renderer) print(prg.src[2].arg) @@ -22,7 +22,7 @@ class TestLinearizerRewrite(unittest.TestCase): with Context(SPLIT_REDUCEOP=0): si = out.schedule_linear().src[-1] opts_to_apply = [] - opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4)) + opts_to_apply.append(Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))) ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply))) prg = to_program(ast, Device["CPU"].renderer) print(prg.src[2].arg) diff --git a/test/null/test_process_replay.py b/test/null/test_process_replay.py index b75e4b09cf..182c39b0a9 100644 --- a/test/null/test_process_replay.py +++ b/test/null/test_process_replay.py @@ -2,6 +2,7 @@ import unittest from tinygrad import Tensor, Device, Context from tinygrad.codegen import do_to_program from tinygrad.codegen.opt import Opt, OptOps +from tinygrad.uop.ops import AxisType from test.external.process_replay.process_replay import replay_to_program from test.helpers import replace_opts @@ -27,7 +28,7 @@ class TestProcessReplay(unittest.TestCase): def test_replay_with_opt(self): # opts=[Opt(...)] means apply a specific opt - opts = [Opt(OptOps.UPCAST, 0, 4)] + opts = [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))] ast = replace_opts(self.ast, opts) p = do_to_program(ast, self.renderer) good, compare, _ = replay_to_program(p, ast, self.renderer) diff --git a/test/null/test_uops_stats.py b/test/null/test_uops_stats.py index 315f2a657a..831ed67db2 100644 --- a/test/null/test_uops_stats.py +++ b/test/null/test_uops_stats.py @@ -4,7 +4,7 @@ from tinygrad.helpers import GlobalCounters from tinygrad.engine.realize import compile_linear, estimate_uop from tinygrad.codegen import to_program from tinygrad.renderer import Estimates -from tinygrad.uop.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp, AxisType from tinygrad.dtype import dtypes from tinygrad.codegen.opt import Opt, OptOps, KernelOptError from tinygrad.device import Device @@ -190,7 +190,7 @@ class TestStatsOptimized(unittest.TestCase): @unittest.skip("fails locally on AMD") def test_gemm_tc_unroll_half(self): try: - p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UPCAST, 4, 2)]), + p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no tensor cores") @@ -199,7 +199,7 @@ class TestStatsOptimized(unittest.TestCase): def test_gemm_tc_unroll(self): try: - p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UPCAST, 4, 2)]), + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no tensor cores") @@ -209,20 +209,22 @@ class TestStatsOptimized(unittest.TestCase): # this is a good lesson about why UPCASTing is a good idea def test_gemm_one_upcasted(self): - p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer) self.check_gemm(p) self.assertEqual(p.src[0].arg.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N) def test_gemm_upcasted(self): - p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 4, 4)]), + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer) self.check_gemm(p) self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N) def test_gemm_upcasted_locals(self): try: - p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4), - Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL))]), + renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no locals") self.check_gemm(p) @@ -230,7 +232,7 @@ class TestStatsOptimized(unittest.TestCase): def test_gemm_group(self): try: - p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.LOCAL, 2, 4)]), renderer=Device[Device.DEFAULT].renderer) + p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.SPLIT, 2, (4, AxisType.GROUP_REDUCE))]), renderer=Device[Device.DEFAULT].renderer) except KernelOptError: raise unittest.SkipTest("no locals") SZ = N*N*4 diff --git a/test/opt/test_gen_float4.py b/test/opt/test_gen_float4.py index 7bdcd6e1f3..6306f5ef70 100644 --- a/test/opt/test_gen_float4.py +++ b/test/opt/test_gen_float4.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Device, Tensor, Variable, dtypes -from tinygrad.uop.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops, AxisType from tinygrad.codegen import to_program from tinygrad.codegen.opt import Opt, OptOps @@ -24,7 +24,7 @@ class TestFloat4(unittest.TestCase): s = c.schedule_linear().src[0] realized_ast = s.src[0] - opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)] + opts_to_apply = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))] program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) assert TestFloat4.count_float4(tuple(program.src[1].src)) == (2, 1) @@ -35,7 +35,8 @@ class TestFloat4(unittest.TestCase): c = a + b s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]), + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (4, 2) @@ -46,7 +47,7 @@ class TestFloat4(unittest.TestCase): s = c.schedule_linear().src[0] realized_ast = s.src[0] - opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)] + opts_to_apply = [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))] program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer) assert TestFloat4.count_float4(tuple(program.src[1].src)) == (0, 1) @@ -57,7 +58,8 @@ class TestFloat4(unittest.TestCase): c = a + b s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]), + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=1, arg=(4, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.UPCAST))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (0, 2) @@ -70,7 +72,8 @@ class TestFloat4(unittest.TestCase): # float4 should be emitted (the reduce axis of size 4 is the float4 axis here) s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=1, arg=(4, AxisType.UNROLL))]), + renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (0, 0) @@ -84,7 +87,8 @@ class TestFloat4(unittest.TestCase): # UPDATE: now we do this fusion s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=1, arg=0)]), + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST)), + Opt(op=OptOps.SPLIT, axis=1, arg=(0, AxisType.UNROLL))]), renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) in {(0,1), (1,1)} @@ -98,7 +102,8 @@ class TestFloat4(unittest.TestCase): # since the top axis is not contiguous. s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]), + renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (0, 1) @@ -110,7 +115,8 @@ class TestFloat4(unittest.TestCase): # should float4 b but not a s = c.schedule_linear().src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]), + renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (1, 1) @@ -123,7 +129,8 @@ class TestFloat4(unittest.TestCase): # should float4 both s = c.linear_with_vars()[0].src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]), + renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (2, 1) @@ -136,7 +143,8 @@ class TestFloat4(unittest.TestCase): # should float4 a but not b s = c.linear_with_vars()[0].src[0] - uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src) + uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))]), + renderer=Device[Device.DEFAULT].renderer).src[1].src) assert TestFloat4.count_float4(uops) == (1, 1) diff --git a/test/opt/test_kernel_opts.py b/test/opt/test_kernel_opts.py index 096f3fa0ef..c40f079853 100644 --- a/test/opt/test_kernel_opts.py +++ b/test/opt/test_kernel_opts.py @@ -1,6 +1,7 @@ import unittest from tinygrad import Device, Tensor, dtypes from tinygrad.codegen.opt import Opt, OptOps, KernelOptError +from tinygrad.uop.ops import AxisType # TODO: write a clean version of this from test.backend.test_linearizer import helper_linearizer_opt @@ -15,24 +16,28 @@ class TestKernelOpts(unittest.TestCase): b = Tensor.rand(4, 4, N) r = (b.sqrt() + ((a+1).sum(axis=3).exp())) helper_linearizer_opt(r, [ - [Opt(OptOps.LOCAL, 0, 2)], - [Opt(OptOps.LOCAL, 0, 8)], - [Opt(OptOps.LOCAL, 0, 16)], # Checking how it works with locals - [Opt(OptOps.LOCAL, 1, 2, top=True)], - [Opt(OptOps.LOCAL, 1, 32, top=True)], - [Opt(OptOps.LOCAL, 1, 64, top=True)], # Checking how it works with grouped reduce - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 2, top=True)], - [Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.LOCAL, 2, 16, top=True)], - [Opt(OptOps.LOCAL, 0, 32), Opt(OptOps.LOCAL, 2, 2, top=True)], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL))], # Checking how it works with locals + [Opt(OptOps.SPLIT, 1, (2, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 1, (32, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 1, (64, AxisType.GROUP_REDUCE, True))], # Checking how it works with grouped reduce + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (16, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 0, (32, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))], # Checking how it works with locals + grouped reduce - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 64, top=True)], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (64, AxisType.GROUP_REDUCE, True))], # Checking how it works with locals + grouped reduce + upcasts - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 2, top=True), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 4, 4)], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL))], # many local + many group - [Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.LOCAL, 2, 2), Opt(OptOps.LOCAL, 3, 2), Opt(OptOps.LOCAL, 4, 2)], - [Opt(OptOps.LOCAL, 0, 2)] * 4, - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 2), Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 4, 2), - Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 6, 2), Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 8, 2)], + [Opt(OptOps.SPLIT, 1, (2, AxisType.GROUP_REDUCE)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE)), + Opt(OptOps.SPLIT, 3, (2, AxisType.GROUP_REDUCE)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE))], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))] * 4, + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE)), + Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE)), + Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 6, (2, AxisType.GROUP_REDUCE)), + Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 8, (2, AxisType.GROUP_REDUCE))], ]) def test_upcasts(self): @@ -42,9 +47,9 @@ class TestKernelOpts(unittest.TestCase): b = Tensor.rand(N, N) r = (a+b).sqrt() * ((a+1).exp()) helper_linearizer_opt(r, [ - [Opt(OptOps.UPCAST, 0, 2)], - [Opt(OptOps.UPCAST, 0, 4)], - [Opt(OptOps.UPCAST, 0, 8)], # Checking how it works with upcasts + [Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST))], # Checking how it works with upcasts ]) def test_full_upcast(self): @@ -53,7 +58,7 @@ class TestKernelOpts(unittest.TestCase): b = Tensor.rand(4) r = (a+b).sqrt() * ((a+1).exp()) helper_linearizer_opt(r, [ - [Opt(OptOps.UPCAST, 0, 4)], # Checking how it works with upcasts + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], # Checking how it works with upcasts ]) @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @@ -65,24 +70,28 @@ class TestKernelOpts(unittest.TestCase): b = Tensor.rand(N, N) r = a@b helper_linearizer_opt(r, [ - [Opt(OptOps.UPCAST, 0, 2)], - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)], # Checking how it works with upcasts - [Opt(OptOps.LOCAL, 0, 2)], - [Opt(OptOps.LOCAL, 1, 32)], - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4)], - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 32)], - [Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.LOCAL, 1, 8)], # Checking how it works with locals - [Opt(OptOps.LOCAL, 2, 2, top=True)], - [Opt(OptOps.LOCAL, 2, 32, top=True)], - [Opt(OptOps.LOCAL, 2, 32, top=True), Opt(OptOps.UPCAST, 2, 4)], # Checking how it works with grouped_reduce - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.LOCAL, 4, 32, top=True)], - [Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.LOCAL, 3, 32, top=True)], - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.LOCAL, 4, 4, top=True)], # Checking how it works with local+grouped_reduce + [Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], # Checking how it works with upcasts + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 1, (32, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (32, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (8, AxisType.LOCAL))], # Checking how it works with locals + [Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 2, (32, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 2, (32, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL))], # Checking how it works with grouped_reduce + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (32, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL)), Opt(OptOps.SPLIT, 3, (32, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL)), + Opt(OptOps.SPLIT, 4, (4, AxisType.GROUP_REDUCE, True))], # Checking how it works with local+grouped_reduce # Checking all together - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 4, 8, top=True), Opt(OptOps.UPCAST, 4, 4), Opt(OptOps.UPCAST, 0, 4), - Opt(OptOps.UPCAST, 1, 2)], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST))], # Full global upcast + local - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 4, 8, top=True), Opt(OptOps.UPCAST, 4, 4), Opt(OptOps.UPCAST, 0, 8)], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST))], ]) @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @@ -94,25 +103,34 @@ class TestKernelOpts(unittest.TestCase): r = a.sum(axis=(1,3)) helper_linearizer_opt(r, [ # openCL / DEV=CL is 256 max threads - [Opt(OptOps.LOCAL, 2, 2, top=True)], [Opt(OptOps.LOCAL, 2, 32, top=True)], - [Opt(OptOps.LOCAL, 3, 2, top=True)], [Opt(OptOps.LOCAL, 3, 32, top=True)], # Checking how it works with 1 grouped_reduce. - [Opt(OptOps.LOCAL, 2, 2, top=True), Opt(OptOps.LOCAL, 4, 2, top=True)], - [Opt(OptOps.LOCAL, 2, 16, top=True), Opt(OptOps.LOCAL, 4, 2, top=True)], - [Opt(OptOps.LOCAL, 2, 4, top=True), Opt(OptOps.LOCAL, 4, 64, top=True)], # Checking how it works with 2 grouped_reduces. - [Opt(OptOps.LOCAL, 2, 16, top=True), Opt(OptOps.LOCAL, 4, 2, top=True), Opt(OptOps.UPCAST, 2, 4)], + [Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True))], [Opt(OptOps.SPLIT, 2, (32, AxisType.GROUP_REDUCE, True))], + # Checking how it works with 1 grouped_reduce. + [Opt(OptOps.SPLIT, 3, (2, AxisType.GROUP_REDUCE, True))], [Opt(OptOps.SPLIT, 3, (32, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 2, (16, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE, True))], + [Opt(OptOps.SPLIT, 2, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 4, (64, AxisType.GROUP_REDUCE, True))], # Checking how it works with 2 grouped_reduces. + [Opt(OptOps.SPLIT, 2, (16, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 2, (4, AxisType.UNROLL))], # Checking how it works with 2 grouped_reduces + upcasts. - [Opt(OptOps.LOCAL, 2, 2, top=True), Opt(OptOps.LOCAL, 4, 32, top=True), Opt(OptOps.UPCAST, 4, 4)], - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.LOCAL, 4, 4, top=True), Opt(OptOps.LOCAL, 6, 4, top=True)], + [Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 4, (32, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 6, (4, AxisType.GROUP_REDUCE, True))], # Checking how it works with 2 grouped_reduces + upcasts + locals. - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.LOCAL, 4, 2, top=True), Opt(OptOps.LOCAL, 6, 32, top=True), - Opt(OptOps.UPCAST, 5, 4)], - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.LOCAL, 4, 8, top=True), Opt(OptOps.LOCAL, 6, 4, top=True), - Opt(OptOps.UPCAST, 0, 2)], - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.LOCAL, 4, 8, top=True), Opt(OptOps.LOCAL, 6, 4, top=True), - Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 4, 4), - Opt(OptOps.UPCAST, 5, 4)], # Checking how it works with 2 grouped_reduces + upcasts + locals. - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.LOCAL, 4, 4, top=True), Opt(OptOps.LOCAL, 6, 4, top=True), - Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 0, 2)], # No globals + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (2, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 6, (32, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 5, (4, AxisType.UNROLL))], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 6, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 6, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 4, (4, AxisType.UNROLL)), + Opt(OptOps.SPLIT, 5, (4, AxisType.UNROLL))], # Checking how it works with 2 grouped_reduces + upcasts + locals. + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 6, (4, AxisType.GROUP_REDUCE, True)), + Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], # No globals ]) @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") @@ -126,17 +144,18 @@ class TestKernelOpts(unittest.TestCase): atol, rtol = 0.25, 0.01 helper_linearizer_opt(r, [ [], - [Opt(OptOps.UPCAST, 0, 4)], - [Opt(OptOps.UPCAST, 1, 4)], - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)], # check upcasts - [Opt(OptOps.UPCAST, 4, 2)], # check unroll - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 5, 2)], # check combo of unroll and upcast - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 6, 2)], - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 6, 4)], - [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 0, 4)], # check permutations - [Opt(OptOps.UPCAST, 4, 2), Opt(OptOps.UPCAST, 0, 4)], - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 5, 2), Opt(OptOps.UPCAST, 1, 4)], - [Opt(OptOps.UPCAST, 4, 2), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 6, 4)], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], # check upcasts + [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))], # check unroll + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 5, (2, AxisType.UNROLL))], # check combo of unroll and upcast + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (2, AxisType.UNROLL))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL))], + [Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], # check permutations + [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 5, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), + Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL))], ], apply_tc=True, atol=atol, rtol=rtol) @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") @@ -150,10 +169,12 @@ class TestKernelOpts(unittest.TestCase): r = a.matmul(b, dtype=dtypes.half) atol, rtol = 0.25, 0.01 helper_linearizer_opt(r, [ - [Opt(OptOps.UPCAST, 4, 0)], # check full unroll of reduce with locals - [Opt(OptOps.LOCAL, 0, 4)], # check local - [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 6, 4), Opt(OptOps.LOCAL, 0, 2)], - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 6, 2), Opt(OptOps.UPCAST, 0, 4)], + [Opt(OptOps.SPLIT, 4, (0, AxisType.UNROLL))], # check full unroll of reduce with locals + [Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL))], # check local + [Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (4, AxisType.UNROLL)), + Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], + [Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 6, (2, AxisType.UNROLL)), + Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], ], apply_tc=True, atol=atol, rtol=rtol) def test_padto_matmul(self): @@ -168,7 +189,7 @@ class TestKernelOpts(unittest.TestCase): [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32)], [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.PADTO, 2, 32)], # can optimize further post PADTO - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST)),], ]) def test_padto_upcasted_not_ok(self): @@ -176,19 +197,19 @@ class TestKernelOpts(unittest.TestCase): a = Tensor.rand(N, N) b = Tensor.rand(N, N) helper_linearizer_opt(a@b, [ - [Opt(OptOps.UPCAST, 0, 0)], - [Opt(OptOps.UPCAST, 1, 0)], - [Opt(OptOps.UPCAST, 2, 0)], + [Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST))], + [Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL))], [Opt(OptOps.PADTO, 0, 8)], [Opt(OptOps.PADTO, 1, 8)], [Opt(OptOps.PADTO, 2, 8)], ]) with self.assertRaises(KernelOptError): - helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 0, 0), Opt(OptOps.PADTO, 1, 8)]]) + helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 1, 8)]]) with self.assertRaises(KernelOptError): - helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 1, 0), Opt(OptOps.PADTO, 1, 8)]]) + helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 1, 8)]]) with self.assertRaises(KernelOptError): - helper_linearizer_opt(a@b, [[Opt(OptOps.UPCAST, 2, 0), Opt(OptOps.PADTO, 2, 8)]]) + helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL)), Opt(OptOps.PADTO, 2, 8)]]) def test_padto_sum_ok(self): N = 18 @@ -198,11 +219,11 @@ class TestKernelOpts(unittest.TestCase): helper_linearizer_opt(a.sum(0), [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) helper_linearizer_opt(a.sum(1), [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) for axis in (0, 1): @@ -226,7 +247,8 @@ class TestKernelOpts(unittest.TestCase): def test_padto_group_full_unroll_sum(self): a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize() out = ((a * 0.5).float().square()).sum(axis=(0, 2)) - opts_to_apply = [Opt(OptOps.LOCAL, 2, 256, top=True), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UPCAST, 3, 0), Opt(OptOps.UPCAST, 0, 7)] + opts_to_apply = [Opt(OptOps.SPLIT, 2, (256, AxisType.GROUP_REDUCE, True)), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.SPLIT, 3, (0, AxisType.UNROLL)), + Opt(OptOps.SPLIT, 0, (7, AxisType.UPCAST))] helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False) def test_padto_sum(self): @@ -247,11 +269,11 @@ class TestKernelOpts(unittest.TestCase): helper_linearizer_opt(a.max(0), [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) helper_linearizer_opt(a.max(1), [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) helper_linearizer_opt(a.max(), [[Opt(OptOps.PADTO, 0, 32)],]) @@ -263,7 +285,7 @@ class TestKernelOpts(unittest.TestCase): a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0).int() helper_linearizer_opt(a.max(0), [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) def test_padto_where_multioutput(self): @@ -274,7 +296,7 @@ class TestKernelOpts(unittest.TestCase): a1 = r.where(2, 0).int() helper_linearizer_opt([a0.max(0), a1.max(0)], [ [Opt(OptOps.PADTO, 0, 32)], - [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),], + [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),], ]) @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") @@ -286,16 +308,21 @@ class TestKernelOpts(unittest.TestCase): b = Tensor.rand(N, N) r = a@b opts_shapes = [ - ([Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("red",32)]), - ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.LOCAL, 3, 2)], [("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]), + ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], [("blue",16),("blue",32),("cyan",2),("red",32)]), + ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (2, AxisType.GROUP_REDUCE))], + [("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]), # check to ensure local_dims are stable for full UNROLL of the first reduce - ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.UPCAST, 3, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), - ([Opt(OptOps.UPCAST, 2, 0),Opt(OptOps.LOCAL, 0, 2)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), + ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.UNROLL))], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), + ([Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL)),Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), # check behavior for full UNROLL on an existing GROUP - ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.LOCAL, 3, 0),Opt(OptOps.UPCAST, 3, 2)], [("blue",16),("blue",32),("cyan",2),("green",16),("magenta",2)]), - ([Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.LOCAL, 3, 0),Opt(OptOps.UPCAST, 3, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), - ([Opt(OptOps.LOCAL, 2, 0),Opt(OptOps.LOCAL, 0, 2),Opt(OptOps.UPCAST, 2, 0)], [("blue",16),("blue",32),("cyan",2),("magenta",32)]), - ([Opt(OptOps.LOCAL, 2, 2),Opt(OptOps.UPCAST, 2, 0)], [("blue",32),("blue",32),("red",16),("magenta",2)]), + ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 3, (2, AxisType.UNROLL))], + [("blue",16),("blue",32),("cyan",2),("green",16),("magenta",2)]), + ([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 3, (0, AxisType.UNROLL))], + [("blue",16),("blue",32),("cyan",2),("magenta",32)]), + ([Opt(OptOps.SPLIT, 2, (0, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL))], + [("blue",16),("blue",32),("cyan",2),("magenta",32)]), + ([Opt(OptOps.SPLIT, 2, (2, AxisType.GROUP_REDUCE)),Opt(OptOps.SPLIT, 2, (0, AxisType.UNROLL))], + [("blue",32),("blue",32),("red",16),("magenta",2)]), ] helper_linearizer_opt(r, [x[0] for x in opts_shapes], color_sizes=[x[1] for x in opts_shapes]) @@ -306,21 +333,21 @@ class TestKernelOpts(unittest.TestCase): a = Tensor.arange(128).clone() # NOTE: arange no longer has reduce ops available for opt helper_linearizer_opt(a, [ - [Opt(op=OptOps.LOCAL, axis=0, arg=8)], - [Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=0, arg=0)], + [Opt(op=OptOps.SPLIT, axis=0, arg=(8, AxisType.LOCAL))], + [Opt(op=OptOps.SPLIT, axis=0, arg=(8, AxisType.LOCAL)), Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))], ]) def test_double_sum_group(self): a = Tensor.rand(4, 4, 4) r = a.sum((1, 2)).sum() with self.assertRaises(KernelOptError): - helper_linearizer_opt(r, [[Opt(OptOps.LOCAL, 0, 16, top=True)],]) + helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 0, (16, AxisType.GROUP_REDUCE, True))],]) r = a.sum((1, 2)).sum() with self.assertRaises(KernelOptError): - helper_linearizer_opt(r, [[Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 16, top=True)],]) + helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 1, (4, AxisType.UNROLL)), Opt(OptOps.SPLIT, 0, (16, AxisType.GROUP_REDUCE, True))],]) r = a.sum((1, 2)).sum() with self.assertRaises(KernelOptError): - helper_linearizer_opt(r, [[Opt(OptOps.LOCAL, 1, 4, top=True), Opt(OptOps.LOCAL, 1, 16, top=True)],]) + helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 1, (4, AxisType.GROUP_REDUCE, True)), Opt(OptOps.SPLIT, 1, (16, AxisType.GROUP_REDUCE, True))],]) if __name__ == '__main__': unittest.main() diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index b0c6d5705f..d45a262383 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -3,7 +3,7 @@ import unittest from tinygrad import Device, Tensor, dtypes from tinygrad.tensor import _to_np_dtype -from tinygrad.uop.ops import Ops, UOp, buffers +from tinygrad.uop.ops import Ops, UOp, AxisType, buffers from tinygrad.dtype import DType from tinygrad.device import Buffer from tinygrad.helpers import DEV, Context @@ -178,7 +178,7 @@ class TestTensorCores(unittest.TestCase): tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s) x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) - opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UPCAST, 4, 2)] + opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))] ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA] self.assertGreater(len(wmmas), 0) @@ -192,7 +192,7 @@ class TestTensorCores(unittest.TestCase): tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) - opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UPCAST, 4, 2)] + opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))] ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA] self.assertGreater(len(wmmas), 0) @@ -207,7 +207,7 @@ class TestTensorCores(unittest.TestCase): tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out).relu() - opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UPCAST, 4, 2)] + opts = [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.SPLIT, 4, (2, AxisType.UNROLL))] ast = helper_linearizer_opt(r, [opts[1:]], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) wmmas = [u for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.WMMA] self.assertGreater(len(wmmas), 0) diff --git a/tinygrad/codegen/opt/__init__.py b/tinygrad/codegen/opt/__init__.py index a408a28cf8..ec637b2034 100644 --- a/tinygrad/codegen/opt/__init__.py +++ b/tinygrad/codegen/opt/__init__.py @@ -4,7 +4,7 @@ from enum import Enum, auto from dataclasses import dataclass class OptOps(Enum): - TC = auto(); UPCAST = auto(); LOCAL = auto(); PADTO = auto(); SWAP = auto() # noqa: E702 + TC = auto(); SPLIT = auto(); PADTO = auto(); SWAP = auto() # noqa: E702 def __lt__(self, x:OptOps): return self.value < x.value @dataclass(frozen=True, order=True) @@ -12,8 +12,7 @@ class Opt: op: OptOps axis: int|None = None arg: int|tuple|None = None - top: bool = False - def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg}{', top=True' if self.top else ''})" + def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})" class KernelOptError(Exception): pass def check(cond:bool, msg:str=""): diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 1a5ee07854..de94a2356e 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -35,9 +35,9 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None] if szs: # set it to the replaced range - rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0] + rngs[tc_dim] = tk.apply_opt(Opt(OptOps.SPLIT, tk.rngs.index(rngs[tc_dim]), (szs[0], AxisType.UPCAST)))[0] if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N - tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0])) + tk.apply_opt(Opt(OptOps.SPLIT, tk.rngs.index(rngs[0]), (szs[0], AxisType.LOCAL))) return tk # make a copy so it does not mutate the input @@ -52,7 +52,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: unit_stride_axes_mul_4 = [k.rngs.index(c) for c in idx.get_idx().split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0 and c not in idx.get_valid().backward_slice] if len(unit_stride_axes_mul_4): - if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims+k.unrollable_dims: k.apply_opt(Opt(OptOps.UPCAST, axis, 4)) + if (axis:=unit_stride_axes_mul_4[0]) in (upd:=k.upcastable_dims)+k.unrollable_dims: + k.apply_opt(Opt(OptOps.SPLIT, axis, (4, AxisType.UPCAST if axis in upd else AxisType.UNROLL))) # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4) @@ -68,17 +69,17 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if DEBUG >= 3: print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}") try: - if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.LOCAL, k.axes_of(AxisType.REDUCE)[0], MV_THREADS_PER_ROW)) + if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.SPLIT, k.axes_of(AxisType.REDUCE)[0], (MV_THREADS_PER_ROW, AxisType.GROUP_REDUCE))) except KernelOptError: pass - if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE)) - if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD)) + if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.SPLIT, global_idx, (MV_BLOCKSIZE, AxisType.LOCAL))) + if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.SPLIT, global_idx, (MV_ROWS_PER_THREAD, AxisType.UPCAST))) return k # are we grouping? (requires local shape support) if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if k.ren.target.device == "QCOM" else 2048), False): for axis, sz in itertools.product(k.axes_of(AxisType.REDUCE)[:3], (16,)): try: - k.apply_opt(Opt(OptOps.LOCAL, axis, sz, top=True)) + k.apply_opt(Opt(OptOps.SPLIT, axis, (sz, AxisType.GROUP_REDUCE, True))) break except KernelOptError: pass @@ -103,7 +104,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if resolve(global_items_after < getenv("OCCUPANCY_FLOOR", 4096), False): continue if DEBUG >= 4: print(f"upcasting masked axis : {axis}") to_upcast.append(axis) - for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0)) + for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.SPLIT, axis, (0, AxisType.UPCAST))) # potentially do more upcasts of non reduce axes based on a heuristic is_dsp = k.ren is not None and k.ren.target.device == "DSP" @@ -129,7 +130,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if xb_choices: xb_choices = sorted(xb_choices) if DEBUG >= 4: print(f"more upcast axis : {xb_choices}") - k.apply_opt(Opt(OptOps.UPCAST, xb_choices[0][2], xb_choices[0][3])) + k.apply_opt(Opt(OptOps.SPLIT, xb_choices[0][2], (xb_choices[0][3], AxisType.UPCAST))) upcasted_axis.add(xb_choices[0][2]) else: break @@ -138,21 +139,21 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: try: if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64): if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32: - k.apply_opt(Opt(OptOps.UPCAST, k.unrollable_dims[-1], 0)) + k.apply_opt(Opt(OptOps.SPLIT, k.unrollable_dims[-1], (0, AxisType.UNROLL))) # if it's small, upcast a second reduce dimension too if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3: - k.apply_opt(Opt(OptOps.UPCAST, k.unrollable_dims[-1], 0)) + k.apply_opt(Opt(OptOps.SPLIT, k.unrollable_dims[-1], (0, AxisType.UNROLL))) else: for splits in [4]: if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0: - k.apply_opt(Opt(OptOps.UPCAST, axis, splits)) + k.apply_opt(Opt(OptOps.SPLIT, axis, (splits, AxisType.UNROLL))) break except KernelOptError: pass # if nothing at all is upcasted and it's easy to, do an upcast for splits in [4]: if not k.upcasted and k.upcastable_dims and k.full_shape[k.upcastable_dims[-1]] % splits == 0: - k.apply_opt(Opt(OptOps.UPCAST, k.upcastable_dims[-1], splits)) + k.apply_opt(Opt(OptOps.SPLIT, k.upcastable_dims[-1], (splits, AxisType.UPCAST))) # **** local groups **** @@ -169,7 +170,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if opts and workgroup < 32: # fill at least one wave: grow the innermost local as much as possible axis, sz = opts[0] opts[0] = axis, max(x for x in range(1, min(int(k.full_shape[axis]), 128 * sz // workgroup) + 1) if int(k.full_shape[axis]) % x == 0) - for axis, sz in opts: k.apply_opt(Opt(OptOps.LOCAL, axis, sz)) + for axis, sz in opts: k.apply_opt(Opt(OptOps.SPLIT, axis, (sz, AxisType.LOCAL))) else: # prioritize making expand axes local local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \ @@ -183,7 +184,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: for axis, local_sz in sorted(to_local[:3]): axis = axis - deleted_shape will_delete_shape = local_sz == k.full_shape[axis] - k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz)) + k.apply_opt(Opt(OptOps.SPLIT, axis, (local_sz, AxisType.LOCAL))) if will_delete_shape: deleted_shape += 1 return k diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index a97157c396..66e96a0968 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -11,9 +11,8 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check from tinygrad.codegen.simplify import pm_flatten_range from tinygrad.renderer import Renderer -upcast_to = {AxisType.GLOBAL: AxisType.UPCAST, AxisType.LOCAL: AxisType.UPCAST, AxisType.WEAK: AxisType.UPCAST, - AxisType.GROUP_REDUCE: AxisType.UNROLL, AxisType.REDUCE: AxisType.UNROLL} -local_to = {AxisType.GLOBAL: AxisType.LOCAL, AxisType.WEAK: AxisType.LOCAL, AxisType.REDUCE: AxisType.GROUP_REDUCE} +split_targets = {AxisType.UPCAST: (AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK), AxisType.UNROLL: (AxisType.REDUCE, AxisType.GROUP_REDUCE), + AxisType.LOCAL: (AxisType.GLOBAL, AxisType.WEAK), AxisType.GROUP_REDUCE: (AxisType.REDUCE,)} class Scheduler: def __init__(self, ast:UOp, ren:Renderer): @@ -115,30 +114,26 @@ class Scheduler: return axis def apply_opt(self, opt:Opt, append_opt:bool=True): - if opt.op is OptOps.LOCAL: check(self.ren.has_local, "locals needed for opt") - rng = self.rngs[real_axis] if (real_axis:=self.real_axis(opt.op, opt.axis)) >= 0 else UOp(Ops.NOOP) - check(not opt.top or (opt.op is OptOps.LOCAL and rng.arg[-1] is AxisType.REDUCE), "top is only for group reduce") - - opt_to_at = {OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST} ret = None - if opt.op in opt_to_at: - amt:int = int(rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg) - new_type = opt_to_at[opt.op] + if opt.op is OptOps.SPLIT: + check(isinstance(opt.arg, tuple) and len(opt.arg) in (2, 3), f"split arg is (amt, target) or (amt, target, top), not {opt.arg}") + amt, new_type, top = (*cast(tuple, opt.arg), False)[0:3] + check(type(amt) is int and (amt == 0 or amt > 1) and isinstance(new_type, AxisType) and new_type in split_targets and isinstance(top, bool), + f"invalid split arg {opt.arg}") + check(not top or new_type is AxisType.GROUP_REDUCE, "top is only for group reduce") + if new_type in (AxisType.LOCAL, AxisType.GROUP_REDUCE): check(self.ren.has_local, "locals needed for opt") + check(rng.arg[-1] in split_targets[new_type], f"{new_type} is from {split_targets[new_type]}, not {rng.arg[-1]}") - if opt.op is OptOps.UPCAST: - check(rng.arg[-1] in upcast_to, f"upcast is for GLOBAL/LOCAL/LOOP/REDUCE, not {rng.arg[-1]}") - if (new_type:=upcast_to[rng.arg[-1]]) is AxisType.UNROLL: check(amt <= 32, "don't unroll more than 32") - else: check((self.ren is not None and self.ren.target.device == "DSP") or amt <= 16, "don't upcast more than 16") - if opt.op is OptOps.LOCAL: - check(rng.arg[-1] in local_to, f"local is for GLOBAL/LOOP/REDUCE, not {rng.arg[-1]}") - new_type = local_to[rng.arg[-1]] + if amt == 0: amt = int(rng.vmax+1) + if new_type is AxisType.UNROLL: check(amt <= 32, "don't unroll more than 32") + if new_type is AxisType.UPCAST: check(self.ren.target.device == "DSP" or amt <= 16, "don't upcast more than 16") if new_type is AxisType.GROUP_REDUCE: check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores") # TODO: why is this wrong? # prevents METAL compiler hangs - if self.reduceop is not None and (new_type is AxisType.GROUP_REDUCE or (self.group_for_reduces and opt.op != OptOps.PADTO)): + if self.reduceop is not None and (new_type is AxisType.GROUP_REDUCE or self.group_for_reduces): upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]) smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize check(smem_sz <= self.ren.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.ren.shared_max}") @@ -147,7 +142,7 @@ class Scheduler: reduce = [u for u in self.ast.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0] check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges), "cannot have a GROUP_REDUCE inside another reduce") - ret = self.shift_to(rng, amt, new_type, top=opt.top) + ret = self.shift_to(rng, amt, new_type, top=top) elif opt.op is OptOps.TC: check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps check(opt.axis is not None, "tensor core opts must have an axis") diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index e65f019672..c53e0a2a9f 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -11,11 +11,12 @@ from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool from tinygrad.codegen import to_program from tinygrad.codegen.opt.postrange import Scheduler -actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(10)] -actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt) for amt in [0,2,3,4,8,13,16,29] for axis in range(8)] -actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt, top=True) for amt in [13,16,28,29,32,49,64,256] for axis in range(8)] +actions = [Opt(op=OptOps.SPLIT, axis=axis, arg=(amt, at)) for at in (AxisType.UPCAST, AxisType.UNROLL) for amt in [0,2,3,4,5,7] for axis in range(10)] +actions += [Opt(op=OptOps.SPLIT, axis=axis, arg=(amt, at)) for at in (AxisType.LOCAL, AxisType.GROUP_REDUCE) + for amt in [0,2,3,4,8,13,16,29] for axis in range(8)] +actions += [Opt(op=OptOps.SPLIT, axis=axis, arg=(amt, AxisType.GROUP_REDUCE, True)) for amt in [13,16,28,29,32,49,64,256] for axis in range(8)] if getenv("BEAM_PADTO", 0): actions += [Opt(op=OptOps.PADTO, axis=axis, arg=amt) for amt in [32] for axis in range(7)] -actions += [Opt(op=OptOps.LOCAL, axis=0, arg=32)] +actions += [Opt(op=OptOps.SPLIT, axis=0, arg=(32, at)) for at in (AxisType.LOCAL, AxisType.GROUP_REDUCE)] actions += [Opt(op=OptOps.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))] # covers resnet kernels (3 global * 3 reduce) actions += [Opt(op=OptOps.TC, axis=axis, arg=(-1, getenv("TC_OPT", 2), getenv("TC", 1))) for axis in range(9)] @@ -88,7 +89,8 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic if a.axis is not None and a.op is not OptOps.TC: try: ax = s.real_axis(a.op, a.axis) except KernelOptError: continue - if (ax >= s.shape_len) or (s.full_shape[ax] == a.arg and Opt(a.op, a.axis, 0, a.top) in kernel_actions): continue + if (ax >= s.shape_len) or (a.op is OptOps.SPLIT and isinstance(arg:=a.arg, tuple) and s.full_shape[ax] == arg[0] + and replace(a, arg=(0,)+arg[1:]) in kernel_actions): continue s2 = s.copy() try: s2.apply_opt(a)