Revert "fix scheduler copy (#12048)"

This reverts commit 51b7c40788.

tinygrad/codegen/opt/heuristic.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import ImageDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
 from tinygrad.codegen.opt.postrange import Scheduler
 
-def hand_coded_optimizations(k:Scheduler) -> Scheduler:
+def hand_coded_optimizations(k:Scheduler) -> list[Opt]:
   # first try the tensor cores
   """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
   Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
@@ -43,7 +43,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[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]))
-      return tk
+      return tk.applied_opts
 
   # make a copy so it does not mutate the input
   k = k.copy()
@@ -63,7 +63,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
           if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
           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))
-          return k
+          return k.applied_opts
 
   # are we grouping? (requires local shape support)
   if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
@@ -85,7 +85,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
           k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))
 
   # no more opt if we are grouping
-  if k.group_for_reduces: return k
+  if k.group_for_reduces: return k.applied_opts
 
   # **** below this line need to be optional and benchmarked ****
 
@@ -171,4 +171,4 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
         k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
         if will_delete_shape: deleted_shape += 1
 
-  return k
+  return k.applied_opts

tinygrad/codegen/opt/postrange.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import math, itertools
 from collections import defaultdict
-from typing import cast, Final
+from typing import cast, Final, Sequence
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
@@ -45,10 +45,8 @@ class Scheduler:
   def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
 
   def copy(self):
-    ret = Scheduler(self.ast, self.opts)
-    ret.dont_use_locals = self.dont_use_locals
-    ret.applied_opts = self.applied_opts[:]
-    return ret
+    # TODO: this is spamming the many ns on the names
+    return Scheduler(self.get_optimized_ast(), self.opts)
 
   kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
   def get_optimized_ast(self, name_override:str|None=None):
@@ -108,6 +106,10 @@ class Scheduler:
       return axis
     except IndexError as e: raise KernelOptError from e
 
+  def apply_opts(self, opts:Sequence[Opt]) -> Scheduler:
+    for opt in opts: self.apply_opt(opt)
+    return self
+
   def apply_opt(self, opt:Opt, append_opt:bool=True):
     if opt.op is OptOps.NOLOCALS:
       check(all(x not in {AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
@@ -312,10 +314,10 @@ def apply_opts(ctx:Renderer, ast:UOp):
   elif ast.arg is not None and ast.arg.opts_to_apply is not None:
     for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
   elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
+    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
     if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD):
-      # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
-      from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-      k = hand_coded_optimizations(k)
+      for opt in hand_coded_optimizations(k): k.apply_opt(opt)
   return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
 
 pm_postrange_opt = PatternMatcher([

tinygrad/helpers.py

@@ -227,7 +227,7 @@ def profile_marker(name:str, color="gray") -> None:
 cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches" if OSX else "~/.cache")), "tinygrad")
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
 
-VERSION = 23
+VERSION = 22
 _db_connection = None
 def db_connection():
   global _db_connection