From 478e7587557b2c65ed288255c00664e0efc6742e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 5 Sep 2025 18:21:55 -0700 Subject: [PATCH] Revert "fix scheduler copy (#12048)" This reverts commit 51b7c407887d9e7ded6f1d792e5745a1a59b8b9c. --- tinygrad/codegen/opt/heuristic.py | 10 +++++----- tinygrad/codegen/opt/postrange.py | 18 ++++++++++-------- tinygrad/helpers.py | 2 +- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 2ca2c8213b..d95a73f864 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/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 diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index c2cca9cd1e..8c0154c2b2 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/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([ diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index e79998ac6d..37467c5eea 100644 --- a/tinygrad/helpers.py +++ b/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