From 58b7e4fab332fc5a2c86db32bdc445a3f84c288f Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 10 Nov 2025 20:30:37 -0800 Subject: [PATCH] GROUPTOP heuristic on more axes (#13206) fixed dm speed --- tinygrad/codegen/opt/heuristic.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 44ae569508..2bbcffd4f3 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -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