not a part of regen

.github/workflows/test.yml

@@ -608,9 +608,8 @@ jobs:
           test/test_outerworld_range.py test/test_sample.py test/test_randomness.py
     - name: Test CPU=1 RANGEIFY=2
       run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
-    # slow (and still wrong on beautiful_mnist)
-    #- name: Test LLVM=1 RANGEIFY=1 (slow tests)
-    #  run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
+    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
+      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
 
   testdevectorize:
     name: Linux (devectorize)

tinygrad/codegen/__init__.py

@@ -18,7 +18,6 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
-from tinygrad.codegen.simplify import pm_simplify_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
 
 @dataclass
@@ -60,15 +59,11 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 
-  # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-  ret.append(RewriteStep(sym, name="initial symbolic"))
-
   # optimize (schedule) the AST
-  ret.append(RewriteStep(pm_simplify_ranges, name="simplify ranges"))
   ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
+  ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
   # expand
   ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))

tinygrad/codegen/late/expander.py

@@ -1,4 +1,5 @@
 # this converts a lowerer program into a vectorized program
+
 import functools, itertools, operator
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition

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,13 +1,13 @@
 from __future__ import annotations
 import math, itertools
 from collections import defaultdict
-from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad
+from typing import cast, Final, Sequence
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute, AxisType, ssimplify, can_pad
+from tinygrad.uop.symbolic import symbolic_flat
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod
 from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
-from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
 from tinygrad.schedule.rangeify import remove_tags
 
@@ -15,6 +15,20 @@ from tinygrad.schedule.rangeify import remove_tags
 axis_to_pos = {AxisType.LOOP: -1, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+
 class Scheduler:
   def __init__(self, ast:UOp, opts:Renderer):
     self.ast, self.opts = ast, opts
@@ -44,11 +58,17 @@ class Scheduler:
     return ret
   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
+  @property
+  def termination(self):
+    terminators = [u for u in self.ast.parents if u.op in {Ops.REDUCE, Ops.STORE}]
+    termination = {}
+    for t in terminators:
+      # works without pm_flatten_range
+      for u in UOp.sink(*t.src[1 if t.op is Ops.REDUCE else 2:]).parents:
+        if u.op is Ops.RANGE: termination[u] = t
+    return termination
+
+  def copy(self): 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):
@@ -76,6 +96,24 @@ class Scheduler:
 
     self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))
 
+  def simplify_merge_adjacent(self):
+    i = 0
+    while i < len(self.rngs)-1:
+      r0, r1 = self.rngs[i], self.rngs[i+1]
+      # same axistype and same termination
+      termination = self.termination
+      if r0.arg[1] == r1.arg[1] and r0 in termination and r1 in termination and termination[r0] == termination[r1]:
+        s0, s1 = r0.src[0], r1.src[0]
+        # do the merge
+        oidx = self.ast.simplify()
+        new_range = r0.replace(src=(s0*s1,))
+        nidx = graph_rewrite(oidx, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
+        # check if it simplifies
+        if count_divmod(nidx) <= count_divmod(oidx):
+          self.ast = nidx
+          continue
+      i += 1
+
   def colors(self) -> list[str]: return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
   def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
 
@@ -85,7 +123,7 @@ class Scheduler:
     new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
     replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
     sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
-    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount} {str(new_type).split('.')[1].lower()}")
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount}")
     return replaced_rng, new_rng
 
   def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
@@ -108,6 +146,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")
@@ -305,6 +347,7 @@ def apply_opts(ctx:Renderer, ast:UOp):
   if ast.tag is not None: return None
   k = Scheduler(ast, ctx)
   k.convert_loop_to_global()
+  k.simplify_merge_adjacent()
   if BEAM >= 1:
     from tinygrad.codegen.opt.search import beam_search
     rawbufs = bufs_from_ast(ast, ctx.device)
@@ -312,10 +355,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/codegen/opt/search.py

@@ -1,14 +1,17 @@
 from typing import cast
 import functools, math, time, multiprocessing, traceback, signal, atexit
+from collections import defaultdict
 from dataclasses import replace
-from tinygrad.uop.ops import Variable, sym_infer, AxisType, pyrender
+from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType, pyrender
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE
+from tinygrad.dtype import ImageDType, PtrDType
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.renderer import ProgramSpec
+
 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(8)]
@@ -91,6 +94,24 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
 
 # *** external API ***
 
+# get (scrap) buffers for timing the linearizer
+# NOTE: there's also bufs_from_ast in postrange
+def bufs_from_lin(lin:Scheduler, allocate:bool=True) -> list[Buffer]:
+  bufsts: defaultdict[int, list[UOp]] = defaultdict(list)
+  for x in lin.bufs:
+    if x.src[0].base.op is Ops.DEFINE_GLOBAL: bufsts[x.src[0].base.arg].append(x)
+  # TODO: Nones are staying in here if buffers are optimized out!
+  # TODO: add a test for this
+  rawbufs: list[Buffer|None] = [None]*(max(bufsts)+1)
+  for k,lx in bufsts.items():
+    buf_size = prod(dtype.shape) if isinstance(dtype:=lx[0].src[0].dtype, ImageDType) else max(y.st_arg.real_size() for y in lx)
+    assert isinstance(dtype, (PtrDType, ImageDType))
+    if buf_size == 0: buf_size = 1  # create a size 1 buffer if no cell is accessed in kernel. # TODO: remove from kernel input in this case.
+    buf_dtype = dtype if isinstance(dtype, ImageDType) else dtype.base
+    rawbufs[k] = Buffer(lin.opts.device, buf_size, buf_dtype).allocate() if allocate else Buffer(lin.opts.device, buf_size, buf_dtype)
+  #assert all(r is not None for r in rawbufs)
+  return cast(list[Buffer], rawbufs)
+
 # get dictionary of all possible actions
 def get_kernel_actions(lin:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
   acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)

tinygrad/codegen/simplify.py

@@ -1,37 +0,0 @@
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute
-from tinygrad.uop.symbolic import symbolic_flat
-
-def flatten_range(r:UOp):
-  off = 2 if r.op is Ops.STORE else 1
-  rngs = r.src[off:]
-  if not len(rngs): return None
-  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
-  return r.replace(src=r.src[:off]+tuple(new_rngs))
-
-pm_flatten_range = PatternMatcher([
-  # real ranges only
-  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
-])
-
-def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
-def simplify_merge_adjacent(u:UOp) -> UOp|None:
-  i = 2 if u.op is Ops.STORE else 1
-  while i < len(u.src)-1:
-    r0, r1 = u.src[i], u.src[i+1]
-    # check same type
-    if r0.arg[-1] == r1.arg[-1]:
-      s0, s1 = r0.src[0], r1.src[0]
-      # do the merge
-      new_range = r0.replace(src=(s0*s1,))
-      nidx = graph_rewrite(u, _substitute+symbolic_flat+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                           name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
-      # check if it simplifies
-      if count_divmod(nidx) <= count_divmod(u):
-        u = nidx
-        continue
-    i += 1
-  return u
-
-pm_simplify_ranges = PatternMatcher([
-  (UPat((Ops.STORE, Ops.REDUCE), name="u"), simplify_merge_adjacent),
-])

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