sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 6 Activity

UPCAST and LOCAL are SPLIT (#17889)

Browse Source
This commit is contained in:
chenyu
2026-09-01 10:48:59 -04:00
committed by GitHub
parent fd10521869
commit 9dbcddc649
19 changed files with 245 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/max_matmul.py
-1
View File
@@ -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 = {
extra/gemm/tinygrad_nv_matmul.py
+8 -7
View File
@@ -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]
extra/optimization/test_beam_search.py
+4 -3
View File
@@ -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()
test/backend/test_linearizer.py
+13 -10
View File
@@ -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)
test/backend/test_linearizer_dumb.py
+1 -1
View File
@@ -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__':
test/backend/test_opt_gemm.py
+6 -5
View File
@@ -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__':
test/backend/test_quantize_onnx.py
+10 -8
View File
@@ -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__":
test/backend/test_uops.py
+1 -1
View File
@@ -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)
test/device/test_hcq.py
+2 -1
View File
@@ -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()
test/null/test_linearizer_rewrite.py
+4 -4
View File
@@ -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)
test/null/test_process_replay.py
+2 -1
View File
@@ -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)
test/null/test_uops_stats.py
+10 -8
View File
@@ -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
test/opt/test_gen_float4.py
+19 -11
View File
@@ -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)
test/opt/test_kernel_opts.py
+121 -94
View File
@@ -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()
test/opt/test_tensor_cores.py
+4 -4
View File
@@ -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)
tinygrad/codegen/opt/__init__.py
+2 -3
View File
@@ -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=""):
tinygrad/codegen/opt/heuristic.py
+16 -15
View File
@@ -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
tinygrad/codegen/opt/postrange.py
+15 -20
View File
@@ -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")
tinygrad/codegen/opt/search.py
+7 -5
View File
@@ -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)