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

GROUPTOP is LOCAL with top=True (#17874)

Browse Source
This commit is contained in:
chenyu
2026-09-01 00:10:17 -04:00
committed by GitHub
parent cf6fff13aa
commit c2aada33ff
7 changed files with 50 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_linearizer.py
+1 -1
View File
@@ -345,7 +345,7 @@ 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.GROUPTOP, 3, 8),
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
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])
test/external/speed_v_theoretical.py
Vendored
+1 -1
View File
@@ -89,7 +89,7 @@ class TestKernelSpeed(unittest.TestCase):
def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=13)
# theoretical is nv_tflops=165, amd_tflops=123
def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=110, amd_tflops=65)
def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=109, amd_tflops=65)
def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=115, amd_tflops=60)
# theoretical is nv_gbs=1008, amd_gbs=960
test/opt/test_kernel_opts.py
+38 -34
View File
@@ -18,16 +18,16 @@ class TestKernelOpts(unittest.TestCase):
[Opt(OptOps.LOCAL, 0, 2)],
[Opt(OptOps.LOCAL, 0, 8)],
[Opt(OptOps.LOCAL, 0, 16)], # Checking how it works with locals
[Opt(OptOps.GROUPTOP, 1, 2)],
[Opt(OptOps.GROUPTOP, 1, 32)],
[Opt(OptOps.GROUPTOP, 1, 64)], # Checking how it works with grouped reduce
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 2, 2)],
[Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.GROUPTOP, 2, 16)],
[Opt(OptOps.LOCAL, 0, 32), Opt(OptOps.GROUPTOP, 2, 2)],
[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)],
# Checking how it works with locals + grouped reduce
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 2, 64)],
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 64, top=True)],
# Checking how it works with locals + grouped reduce + upcasts
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 2, 2), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 4, 4)],
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 2, 2, top=True), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 4, 4)],
# 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,
@@ -72,17 +72,17 @@ class TestKernelOpts(unittest.TestCase):
[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.GROUPTOP, 2, 2)],
[Opt(OptOps.GROUPTOP, 2, 32)],
[Opt(OptOps.GROUPTOP, 2, 32), Opt(OptOps.UPCAST, 2, 4)], # Checking how it works with grouped_reduce
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 4, 32)],
[Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 3, 32)],
[Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 4, 4)], # Checking how it works with local+grouped_reduce
[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
# Checking all together
[Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 4, 8), Opt(OptOps.UPCAST, 4, 4), Opt(OptOps.UPCAST, 0, 4),
[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)],
# Full global upcast + local
[Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 4, 8), Opt(OptOps.UPCAST, 4, 4), Opt(OptOps.UPCAST, 0, 8)],
[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)],
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -94,21 +94,25 @@ class TestKernelOpts(unittest.TestCase):
r = a.sum(axis=(1,3))
helper_linearizer_opt(r, [
# openCL / DEV=CL is 256 max threads
[Opt(OptOps.GROUPTOP, 2, 2)], [Opt(OptOps.GROUPTOP, 2, 32)],
[Opt(OptOps.GROUPTOP, 3, 2)], [Opt(OptOps.GROUPTOP, 3, 32)], # Checking how it works with 1 grouped_reduce.
[Opt(OptOps.GROUPTOP, 2, 2), Opt(OptOps.GROUPTOP, 4, 2)],
[Opt(OptOps.GROUPTOP, 2, 16), Opt(OptOps.GROUPTOP, 4, 2)],
[Opt(OptOps.GROUPTOP, 2, 4), Opt(OptOps.GROUPTOP, 4, 64)], # Checking how it works with 2 grouped_reduces.
[Opt(OptOps.GROUPTOP, 2, 16), Opt(OptOps.GROUPTOP, 4, 2), Opt(OptOps.UPCAST, 2, 4)],
[Opt(OptOps.GROUPTOP, 2, 2), Opt(OptOps.GROUPTOP, 4, 32), Opt(OptOps.UPCAST, 4, 4)], # Checking how it works with 2 grouped_reduces + upcasts.
[Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 4, 4), Opt(OptOps.GROUPTOP, 6, 4)],
[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)],
# 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)],
# Checking how it works with 2 grouped_reduces + upcasts + locals.
[Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 4, 2), Opt(OptOps.GROUPTOP, 6, 32), Opt(OptOps.UPCAST, 5, 4)],
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 4, 8), Opt(OptOps.GROUPTOP, 6, 4), Opt(OptOps.UPCAST, 0, 2)],
[Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 4, 8), Opt(OptOps.GROUPTOP, 6, 4), 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.GROUPTOP, 4, 4), Opt(OptOps.GROUPTOP, 6, 4), Opt(OptOps.UPCAST, 0, 2),
Opt(OptOps.UPCAST, 0, 2)], # No globals
[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
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
@@ -222,7 +226,7 @@ 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.GROUPTOP, 2, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UPCAST, 3, 0), Opt(OptOps.UPCAST, 0, 7)]
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)]
helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
def test_padto_sum(self):
@@ -310,13 +314,13 @@ class TestKernelOpts(unittest.TestCase):
a = Tensor.rand(4, 4, 4)
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 0, 16)],])
helper_linearizer_opt(r, [[Opt(OptOps.LOCAL, 0, 16, top=True)],])
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.GROUPTOP, 0, 16)],])
helper_linearizer_opt(r, [[Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 16, top=True)],])
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.GROUPTOP, 1, 16)],])
helper_linearizer_opt(r, [[Opt(OptOps.LOCAL, 1, 4, top=True), Opt(OptOps.LOCAL, 1, 16, top=True)],])
if __name__ == '__main__':
unittest.main()
tinygrad/codegen/opt/__init__.py
+3 -2
View File
@@ -4,7 +4,7 @@ from enum import Enum, auto
from dataclasses import dataclass
class OptOps(Enum):
TC = auto(); UPCAST = auto(); LOCAL = auto(); GROUPTOP = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
TC = auto(); UPCAST = auto(); LOCAL = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
def __lt__(self, x:OptOps): return self.value < x.value
@dataclass(frozen=True, order=True)
@@ -12,7 +12,8 @@ class Opt:
op: OptOps
axis: int|None = None
arg: int|tuple|None = None
def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
top: bool = False
def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg}{', top=True' if self.top else ''})"
class KernelOptError(Exception): pass
def check(cond:bool, msg:str=""):
tinygrad/codegen/opt/heuristic.py
+1 -1
View File
@@ -78,7 +78,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
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.GROUPTOP, axis, sz))
k.apply_opt(Opt(OptOps.LOCAL, axis, sz, top=True))
break
except KernelOptError: pass
tinygrad/codegen/opt/postrange.py
+4 -5
View File
@@ -116,12 +116,12 @@ class Scheduler:
return axis
def apply_opt(self, opt:Opt, append_opt:bool=True):
if opt.op in {OptOps.LOCAL, OptOps.GROUPTOP}:
check(self.ren.has_local, "locals needed for opt")
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, OptOps.GROUPTOP: AxisType.GROUP_REDUCE}
opt_to_at = {OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST}
ret = None
if opt.op in opt_to_at:
@@ -135,7 +135,6 @@ class Scheduler:
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 opt.op is OptOps.GROUPTOP: check(rng.arg[-1] is AxisType.REDUCE, "grouptop is for reduce")
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?
@@ -149,7 +148,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.op is OptOps.GROUPTOP)
ret = self.shift_to(rng, amt, new_type, top=opt.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
+2 -2
View File
@@ -13,7 +13,7 @@ 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.GROUPTOP, axis=axis, arg=amt) for amt in [13,16,28,29,32,49,64,256] 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)]
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.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))]
@@ -88,7 +88,7 @@ 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) in kernel_actions): 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
s2 = s.copy()
try:
s2.apply_opt(a)