From c2aada33ffdcc39bcef2ce674412ab25440cbcde Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Tue, 1 Sep 2026 00:10:17 -0400
Subject: [PATCH] GROUPTOP is LOCAL with top=True (#17874)

---
 test/backend/test_linearizer.py      |  2 +-
 test/external/speed_v_theoretical.py |  2 +-
 test/opt/test_kernel_opts.py         | 72 +++++++++++++++-------------
 tinygrad/codegen/opt/__init__.py     |  5 +-
 tinygrad/codegen/opt/heuristic.py    |  2 +-
 tinygrad/codegen/opt/postrange.py    |  9 ++--
 tinygrad/codegen/opt/search.py       |  4 +-
 7 files changed, 50 insertions(+), 46 deletions(-)

diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py
index 99ac143df7..1d70d786e2 100644
--- a/test/backend/test_linearizer.py
+++ b/test/backend/test_linearizer.py
@@ -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])
diff --git a/test/external/speed_v_theoretical.py b/test/external/speed_v_theoretical.py
index 432d7adef2..1ee30ccbd8 100644
--- a/test/external/speed_v_theoretical.py
+++ b/test/external/speed_v_theoretical.py
@@ -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
diff --git a/test/opt/test_kernel_opts.py b/test/opt/test_kernel_opts.py
index 2aed363f26..fddc1da18f 100644
--- a/test/opt/test_kernel_opts.py
+++ b/test/opt/test_kernel_opts.py
@@ -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()
diff --git a/tinygrad/codegen/opt/__init__.py b/tinygrad/codegen/opt/__init__.py
index 61737f2184..a408a28cf8 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(); 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=""):
diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py
index 7c7fb65494..1a5ee07854 100644
--- a/tinygrad/codegen/opt/heuristic.py
+++ b/tinygrad/codegen/opt/heuristic.py
@@ -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
 
diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py
index 5ee2ed5501..595600ba7d 100644
--- a/tinygrad/codegen/opt/postrange.py
+++ b/tinygrad/codegen/opt/postrange.py
@@ -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")
diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py
index 4e505e5dfe..e65f019672 100644
--- a/tinygrad/codegen/opt/search.py
+++ b/tinygrad/codegen/opt/search.py
@@ -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)