more

* move verify to spec / cleanup

* lil

* more explicit

test/test_uop_graph.py

@@ -4,7 +4,6 @@ from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
 from tinygrad.uop.symbolic import sym
-from tinygrad.codegen import full_rewrite_to_sink
 from tinygrad.codegen.late.expander import expander
 from test.test_uops import to_uops_list
 
@@ -722,7 +721,7 @@ class TestExpander(unittest.TestCase):
     self.assertTupleEqual(sink.src[0].arg, (0,2,1,3,4,6,5,7))
 
   def test_contract_no_expand(self):
-    e1 = UOp(Ops.DEFINE_VAR, dtypes.int)
+    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
     con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
     sink = expander_rewrite(con)
     assert sink.op is Ops.VECTORIZE and len(sink.src) == 2
@@ -811,54 +810,6 @@ class TestExpander(unittest.TestCase):
     sink = expander_rewrite(sink)
     print(sink)
 
-class TestIFUOps(unittest.TestCase):
-  def test_create_ifs(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=4, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    idx = UOp.const(dtypes.int, 0)
-    st = UOp(Ops.STORE, dtypes.void, (sbuf.index(idx), UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbuf = UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, 0)), barrier))
-    store = UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, 0), gate), lbuf))
-    sink = UOp(Ops.SINK, dtypes.void, (store,))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
-  def test_expand_ifs_one_gate(self):
-    gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(size=16, addrspace=AddrSpace.LOCAL), (), "smem")
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "gidx0")<1
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 16),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    st = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.float, 42)))
-    barrier = UOp(Ops.BARRIER, dtypes.void, (st,))
-    lbufs = [UOp(Ops.LOAD, dtypes.float, (sbuf.index(UOp.const(dtypes.int, i)), barrier)) for i in range(4)]
-    stores = [UOp(Ops.STORE, dtypes.void, (gbuf.index(UOp.const(dtypes.int, i), gate), lbufs[i])) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
-  # this will be fixed with the merge gated stores bounty
-  @unittest.expectedFailure
-  def test_expand_ifs_dumb(self):
-    buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0)
-    valid = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "gidx0")<5
-    lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), "lidx0")
-    gate = valid&(lidx.ne(2))
-    stores = [UOp(Ops.STORE, dtypes.void, (buf, UOp.const(dtypes.int, i), UOp.const(dtypes.float, i), gate)) for i in range(4)]
-    sink = UOp(Ops.SINK, dtypes.void, tuple(stores))
-    sink = full_rewrite_to_sink(sink)
-    if_uops = [u for u in sink.toposort() if u.op is Ops.IF]
-    self.assertEqual(len(if_uops), 1)
-    self.assertEqual(if_uops[0].src[0], gate)
-
 class TestUOpTags(unittest.TestCase):
   def test_inc_by_one(self):
     g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)

tinygrad/codegen/__init__.py

@@ -1,6 +1,6 @@
 from typing import cast
 from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, SPEC
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, test_pyrender, Ops, UPat
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer
 from tinygrad.dtype import dtypes
@@ -22,8 +22,7 @@ from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_c
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
   if ren is None: ren = Renderer()
 
-  if SPEC: type_verify(list(sink.toposort()), kernel_spec)
-  if SPEC > 1: test_pyrender(sink)
+  if SPEC: type_verify(sink, kernel_spec)
 
   # first we optimize
   if optimize:
@@ -90,7 +89,6 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
   # return the rewritten sink
-  if SPEC > 1: test_pyrender(sink)
   return sink
 
 # inject IF/ENDIF. only needed if device doesn't support gated stores

tinygrad/tensor.py

@@ -11,7 +11,6 @@ from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.mathtraits import MathTrait
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, srender
-from tinygrad.uop.ops import test_pyrender
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
@@ -230,8 +229,7 @@ class Tensor(MathTrait):
     big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
 
     # verify Tensors match the spec
-    if SPEC: type_verify(list(big_sink.toposort()), tensor_spec)
-    if SPEC > 1: test_pyrender(big_sink)
+    if SPEC: type_verify(big_sink, tensor_spec)
 
     if any(isinstance(x._device, tuple) for x in big_sink.toposort()):
       _apply_map_to_tensors(get_multi_map(big_sink), "Apply Multi Map")

tinygrad/uop/ops.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence
+from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence, Iterable
 import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref, collections
 from dataclasses import dataclass
 from enum import Enum, auto
@@ -42,6 +42,13 @@ def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_inf
 
 def range_str(u:UOp) -> str: return '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
 
+def consumer_map_from_toposort(lst:Iterable[UOp]):
+  ret: dict[UOp, dict[UOp, None]] = {}
+  for u in lst:
+    ret[u] = {}
+    for s in u.src: ret[s][u] = None
+  return ret
+
 # used for UOp and UPat
 def pretty_print(x:Any, rep:Callable, srcfn=lambda x: x.src, cache=None, d=0)->str:
   def dfs(x:Any, cache:dict):
@@ -65,8 +72,8 @@ class UOpMetaClass(type):
       assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
       buffers[created] = _buffer
     if SPEC > 1:
+      from tinygrad.uop.spec import full_spec, test_pyrender
       if SPEC > 2: test_pyrender(created)
-      from tinygrad.uop.spec import full_spec
       with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
       if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
     return created
@@ -145,12 +152,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return ret
 
   # returns map of UOps to their consumers in the graph rooted by self
-  def get_consumer_map(self) -> dict[UOp, dict[UOp, None]]:
-    ret: dict[UOp, dict[UOp, None]] = {}
-    for u in self.toposort():
-      ret[u] = {}
-      for s in u.src: ret[s][u] = None
-    return ret
+  def get_consumer_map(self) -> dict[UOp, dict[UOp, None]]: return consumer_map_from_toposort(self.toposort())
 
   def reverse_toposort(self, consumer_map) -> dict[UOp, None]:
     ret: dict[UOp, None] = {}
@@ -1297,15 +1299,14 @@ pm_pyrender = pm_pyrender_extra+PatternMatcher([
 ])
 
 def pyrender(ast:UOp) -> str:
-  cmap = ast.get_consumer_map()
-  uops = list(ast.toposort())
-  ret: dict[str, str] = {}
-  r: dict[UOp, str] = {}
+  lst = list(ast.toposort())
 
+  cmap = consumer_map_from_toposort(lst)
   not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
   always_rendered = {Ops.DEFINE_GLOBAL, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.BUFFER, Ops.COPY, Ops.KERNEL, Ops.WHERE, Ops.END}
+
   to_render: set[UOp] = {ast}
-  for u in uops:
+  for u in lst:
     if u.op in {Ops.SINK}:
       for s in u.src: to_render.add(s)
     if u.op is Ops.STORE: to_render.add(u.src[1])
@@ -1316,7 +1317,9 @@ def pyrender(ast:UOp) -> str:
     to_render.add(u)
 
   kernels: dict[UOp, tuple[str, str]] = {}
-  for i,u in enumerate(uops):
+  r: dict[UOp, str] = {}
+  ret: dict[str, str] = {}
+  for i,u in enumerate(lst):
     if u.op is Ops.KERNEL:
       if u.arg.ast not in kernels:
         kernels[u.arg.ast] = (f"k{len(kernels)}", f"def k{len(kernels)}():\n  " + pyrender(u.arg.ast).replace('\n', '\n  ') + "\n  return ast\n\n")
@@ -1326,28 +1329,10 @@ def pyrender(ast:UOp) -> str:
     #if u.tag is not None: ren += f".rtag({u.tag})"
     if u not in to_render: r[u] = ren
     else:
-      r[u] = f"c{i}" if u is not uops[-1] else "ast"
+      r[u] = f"c{i}" if u is not lst[-1] else "ast"
       ret[r[u]] = ren
   return ''.join([v[1] for v in kernels.values()]) + '\n'.join([f"{k} = {v}" for k,v in ret.items()])
 
-def eval_pyrender(code:str) -> UOp:
-  from tinygrad.dtype import AddrSpace
-  from tinygrad.codegen.opt import Opt, OptOps
-  from tinygrad.schedule.rangeify import BufferizeOpts, Kernel
-  lcls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel,
-                         "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}
-  exec(code, None, lcls)
-  return lcls['ast']
-
-def test_pyrender(test_ast:UOp, check_parents=True):
-  code = pyrender(test_ast)
-  ast:UOp = eval_pyrender(code)
-  if ast is not test_ast:
-    if check_parents:
-      for u in test_ast.toposort(): test_pyrender(u, check_parents=False)
-    raise RuntimeError(f"PYRENDER ISSUE:\nSTR MATCH: {str(test_ast) == str(ast)}\nUOP:\n{test_ast}\nPRODUCED:\n{ast}\nCODE:\n{code}")
-  return code
-
 # *** what was symbolic.py ***
 
 sint = int|UOp

tinygrad/uop/spec.py

@@ -1,7 +1,8 @@
-from typing import cast
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType
+import math
+from typing import cast, Any
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
-from tinygrad.helpers import DEBUG, Context, prod
+from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata
 from tinygrad.uop.validate import validate_index
 
 # four specs:
@@ -161,7 +162,7 @@ kernel_spec = PatternMatcher([
   (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
 
   # END can end multiple axes here
-  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), allow_any_len=True, dtype=dtypes.void), lambda: True),
+  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
 
   # bufferize (must be on ranges)
   (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.op in {Ops.RANGE, Ops.CONST} for y in x.src[1:])),
@@ -174,24 +175,14 @@ kernel_spec = PatternMatcher([
 # *** this spec should match all UOps ever created ***
 
 full_spec = PatternMatcher([
-  # any END
-  (UPat(Ops.END), lambda: True),
-
   # NOOP in the full spec
   (UPat(Ops.NOOP), lambda: True),
 
-  # Invalid must have type Index
-  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),
-  # where on index in rhs position is fine
-  (UPat(Ops.WHERE, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
-
   # all rewrite error are okay
   (UPat(Ops.REWRITE_ERROR), lambda: True),
 
   # rangeify: buffer view with index or load is okay
   (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),)), lambda: True),
-  # copy on index
-  (UPat(Ops.COPY, src=(UPat(Ops.INDEX), UPat())), lambda: True),
   # assign on index. the third op is the shape
   (UPat(Ops.ASSIGN, src=(UPat(), UPat(), UPat())), lambda: True),
 
@@ -209,33 +200,56 @@ full_spec = PatternMatcher([
   # linearizer: outputs + intermediate KERNELs
   (UPat(Ops.KERNEL, dtype=dtypes.void), lambda: True),
 
+  # Invalid must have type Index
+  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),
+  # where on index in rhs position is fine
+  (UPat(Ops.WHERE, dtype=dtypes.index, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True),
   # allow index dtype on a restricted set of UOps
-  (UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.MAX, Ops.WHERE,
+  (UPat((Ops.ADD, Ops.MUL, Ops.MOD, Ops.IDIV, Ops.MAX,
          Ops.SPECIAL, Ops.CAST, Ops.RANGE, Ops.VCONST, Ops.VECTORIZE), dtype=dtypes.index), lambda: True),
 
   # while BIND is being casted
-  (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(), UPat()), arg=None), lambda: True),
+  (UPat(Ops.BIND, (dtypes.int, dtypes.index), (UPat(), UPat()), arg=None), lambda: True),
 
   # in progress MSTACK may lose device
   (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),
 
   # all loads/stores
   (UPat((Ops.LOAD, Ops.STORE)), lambda: True),
-  # all ifs
-  (UPat(Ops.IF), lambda: True),
-  # all DEFINE_VAR to deal with the floats used in reduce collapse
-  (UPat(Ops.DEFINE_VAR), lambda: True),
-  # reshape on STORE
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.STORE),)), lambda: True),
+  # DEFINE_VAR to deal with the floats used in reduce collapse
+  (UPat(Ops.DEFINE_VAR, dtype=dtypes.floats), lambda: True),
   # allow any AFTER
   (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True),
 ])+tensor_spec+kernel_spec+program_spec+shared_spec
 
 # ***** uop helpers *****
 
-def type_verify(uops:list[UOp], check_spec:PatternMatcher):
-  for i,u in enumerate(uops):
+def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
+  lst = list(ast.toposort()) if isinstance(ast, UOp) else ast
+  if SPEC > 1: test_pyrender(lst[-1])  # assume this is the sink
+
+  for i,u in enumerate(lst):
     with Context(TRACK_MATCH_STATS=0): ret = check_spec.rewrite(u)
     if cast(bool|None, ret) is not True:
-      if DEBUG >= 3: print_uops(uops)
+      if DEBUG >= 3: print_uops(lst)
       raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
+
+# late imports to avoid circular import
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.schedule.rangeify import BufferizeOpts, Kernel
+glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel, "Metadata": Metadata,
+                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
+                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace}
+def eval_pyrender(code:str) -> UOp:
+  lcls:dict[str, Any] = {}
+  exec(code, glbls, lcls)
+  return lcls['ast']
+
+def test_pyrender(test_ast:UOp, assert_parents=True):
+  code = pyrender(test_ast)
+  ast:UOp = eval_pyrender(code)
+  if ast is not test_ast:
+    if assert_parents:
+      for u in test_ast.toposort(): test_pyrender(u, assert_parents=False)
+    raise RuntimeError(f"PYRENDER ISSUE:\nSTR MATCH: {str(test_ast) == str(ast)}\nUOP:\n{test_ast}\nPRODUCED:\n{ast}\nCODE:\n{code}")
+  return code