tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

GROUPTOP heuristic on more axes (#13206)

Browse Source 
fixed dm speed
This commit is contained in:
chenyu
2025-11-10 23:30:37 -05:00
committed by GitHub
parent 829cdafccc
commit 58b7e4fab3
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/opt/heuristic.py
+3 -3
View File
@@ -81,10 +81,10 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
return k
# are we grouping? (requires local shape support)
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (128 if NOLOCALS else 2048), False):
for sz in [16]:
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
for axis, sz in itertools.product((0, 1, 2), (16,)):
try:
k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
break
except KernelOptError: pass