unused

.github/workflows/test.yml

@@ -264,8 +264,8 @@ jobs:
       run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
     - name: Run unit tests
       run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Check SPEC=2
-      run: SPEC=2 python3 test/test_tiny.py
+    - name: Check SPEC=3
+      run: SPEC=3 python3 test/test_tiny.py
     - name: Run targetted tests on NULL backend
       run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
     # TODO: too slow

tinygrad/codegen/__init__.py

@@ -1,6 +1,6 @@
 from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, SPEC
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype
-from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
+from tinygrad.uop.spec import type_verify, program_spec, kernel_spec, validate_pyrender
 from tinygrad.renderer import Renderer
 
 # import all pattern matchers here
@@ -20,6 +20,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   if ren is None: ren = Renderer()
 
   if SPEC: type_verify(list(sink.toposort()), kernel_spec)
+  if SPEC > 2: validate_pyrender(sink)
 
   # first we optimize
   if optimize:
@@ -105,4 +106,5 @@ def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   assert len(full_sink.ranges) == 0, "all ranges must end by the sink"
   lst = linearize(full_sink)
   if SPEC: type_verify(lst, program_spec)
+  if SPEC > 2: validate_pyrender(sink)
   return lst

tinygrad/dtype.py

@@ -47,7 +47,7 @@ class DType(metaclass=DTypeMetaClass):
   @staticmethod
   def new(priority:int, itemsize:int, name:str, fmt:FmtStr|None): return DType(priority, itemsize, name, fmt, 1, None)
   def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
-  def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count > 1 else "")
+  def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count != 1 else "")
   def __lt__(self, o:DType): return (self.priority, self.itemsize, self.name, self.fmt, self.count) < (o.priority, o.itemsize, o.name, o.fmt, o.count)
   @property
   def base(self): return self

tinygrad/tensor.py

@@ -11,7 +11,7 @@ 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.spec import type_verify, tensor_spec
+from tinygrad.uop.spec import type_verify, tensor_spec, validate_pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
@@ -230,6 +230,7 @@ class Tensor(MathTrait):
 
     # verify Tensors match the spec
     if SPEC: type_verify(list(big_sink.toposort()), tensor_spec)
+    if SPEC > 2: validate_pyrender(big_sink)
 
     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

@@ -118,7 +118,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def argstr(self): return f'({", ".join(map(str, self.arg))})' if self.op is Ops.REDUCE_AXIS else repr(self.arg)
   def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""
 
-  def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
+  def f(self, op, src=(), **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,)+src, **kwargs)
 
   @functools.cached_property
   def backward_slice(self:UOp) -> dict[UOp, None]:
@@ -335,6 +335,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
     return UOp.const(self.dtype, b, device=self._device, shape=self._shape)
+  def vectorize(self, *src, **kwargs): return UOp(Ops.VECTORIZE, self.dtype.vec(1+len(src)), (self,)+src, **kwargs)
   def broadcast(self, count:int):
     assert self.dtype.count == 1
     if count == 1: return self
@@ -371,15 +372,17 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None, src=None):
     if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
     if isinstance(b, tuple) and all_same(b): b = b[0]  # doesn't have to be a VCONST if they are all the same
+    # NOTE: float('nan') != float('nan'), so we canonicalize here
+    if isinstance(b, float) and math.isnan(b): b = math.nan
     ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtypes.as_const(b, dtype), src=() if src is None else (src,))
     if device is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
     if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape)
     return ret
   @staticmethod
-  def range(end:sint, *arg, dtype=dtypes.index):
+  def range(end:sint, *arg, dtype=dtypes.index, **kwargs):
     if len(arg) == 0: raise RuntimeError("range needs an arg")
     if len(arg) == 1: arg = arg+(AxisType.LOOP,)
-    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=arg)
+    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=arg, **kwargs)
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name)
   def r(self, op:Ops, axis:tuple[int, ...]):
@@ -527,12 +530,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   # TODO: use this in Buffer
   unique_num = itertools.count(0)
   @staticmethod
-  def unique(): return UOp(Ops.UNIQUE, arg=next(UOp.unique_num))
+  def unique(num:int|None=None): return UOp(Ops.UNIQUE, arg=next(UOp.unique_num) if num is None else num)
 
   # *** uop Buffer stuff ***
 
   @staticmethod
-  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType): return UOp(Ops.BUFFER, dtype, (UOp.unique(), UOp(Ops.DEVICE, arg=device)), size)
+  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
+    return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
   @property
   def device(self) -> str|tuple[str, ...]: return cast(str|tuple[str, ...], unwrap(self._device))
   @recursive_property
@@ -1231,17 +1235,21 @@ renderer_infer = PatternMatcher([
 ])
 
 sugar = { Ops.SINK: "sink", Ops.STORE: "store", Ops.LOAD: "load", Ops.SQRT: "sqrt", Ops.INDEX: "index", Ops.REDUCE: "reduce",
-          Ops.WHERE: "where", Ops.RECIP: "reciprocal", Ops.EXP2: "exp2", Ops.LOG2: "log2", Ops.SIN: "sin"}
+          Ops.BIND: "bind", Ops.ASSIGN: "assign", Ops.DETACH: "detach", Ops.TRUNC: "trunc",
+          Ops.WHERE: "where", Ops.RECIP: "reciprocal", Ops.EXP2: "exp2", Ops.LOG2: "log2", Ops.SIN: "sin", Ops.CONTIGUOUS: "contiguous"}
 pm_pyrender = PatternMatcher([
+  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg}, device=\"{x.src[0].arg}\")")),
   (UPat(Ops.CONST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg}, src={x.src[0].arg})")),
-  (UPat(Ops.CONST, name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg})")),
+  (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.const({x.dtype}, {x.arg})")),
   (UPat(Ops.END, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.end({', '.join([y.arg for y in x.src[1:]])})")),
   (UPat(Ops.CAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.cast({x.dtype})")),
   (UPat(Ops.BITCAST, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.bitcast({x.dtype})")),
   (UPat({Ops.MAX, Ops.THREEFRY, Ops.CMPLT, Ops.CMPNE, Ops.POW}, src=UPat(Ops.NOOP), name="x"),
    lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.alu({x.op}, {x.src[1].arg})")),
   (UPat(Ops.RANGE, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg=
-    f"UOp.range({x.src[0].arg}, {str(x.arg[0])}, {str(x.arg[1])}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})")),
+    f"UOp.range({x.src[0].arg}, {str(x.arg[0])}, {str(x.arg[1])}"+\
+      (', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+\
+      (', tag='+str(x.tag) if x.tag is not None else '')+")")),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.NOOP),), name="x"), lambda x: UOp(Ops.NOOP, arg= f"UOp.special({x.src[0].arg}, \"{x.arg}\", dtype={x.dtype})")),
   (UPat(Ops.DEFINE_VAR, name="x"), lambda x: UOp(Ops.NOOP, arg=
     f"UOp.variable(\"{x.arg[0]}\", {x.arg[1]}, {x.arg[2]}{', dtype='+str(x.dtype) if x.dtype is not dtypes.index else ''})")),
@@ -1249,15 +1257,28 @@ pm_pyrender = PatternMatcher([
     arg=f"{x.src[0].arg}.{sugar[x.op]}({', '.join([y.arg for y in x.src[1:]] + ([f'arg={str(x.arg)}'] if x.arg is not None else []))})")),
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.NOOP),), name="x"),
    lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, arg=({', '.join([str(y) for y in x.arg])}))")),
+  # UNIQUE/DEVICE aren't rendered
+  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"), lambda x,u,d: UOp(Ops.NOOP, arg=
+    f"UOp.new_buffer(\"{d.arg}\", {x.size}, {x.dtype}, {u.arg})")),
+  (UPat(Ops.COPY, src=(UPat(Ops.NOOP, name="x"), UPat(Ops.DEVICE, name="d"))), lambda x,d: UOp(Ops.NOOP, arg=
+    f"{x.arg}.copy_to_device(\"{d.arg}\")")),
+  # MovementOp render is short circuited
+  (UPat({Ops.PERMUTE, Ops.FLIP}, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.{x.op.name.lower()}({x.arg})")),
+  (UPat({Ops.RESHAPE, Ops.EXPAND, Ops.SHRINK, Ops.PAD}, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=
+    f"{x.src[0].arg}.f({x.op}, src=({', '.join([y.arg for y in x.src[1:]])},))")),
+  (UPat(Ops.VECTORIZE, src=(), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp({x.op}, dtype={x.dtype})")),
+  (UPat(Ops.VECTORIZE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"UOp.vectorize({', '.join([y.arg for y in x.src])})")),
 ])
 
 @Context(SPEC=0)
 def pyrender(ast:UOp) -> list[str]:
   cmap = ast.get_consumer_map()
-  to_render = set()
+  to_render = set({ast})
+  always_rendered = {Ops.DEFINE_GLOBAL, Ops.LOAD, Ops.BUFFER, Ops.COPY, Ops.CONTIGUOUS} | GroupOp.Movement
+  not_rendered = {Ops.VCONST, Ops.CONST, Ops.DEVICE, Ops.BUFFER, Ops.VECTORIZE}
   for u in ast.toposort():
     if u.op is Ops.STORE: to_render.add(u.src[1])
-    if len(cmap[u]) == 1 and u.op not in {Ops.DEFINE_GLOBAL, Ops.LOAD} or u.op in {Ops.CONST}: continue
+    if len(cmap[u]) == 1 and u.op not in always_rendered or u.op in not_rendered: continue
     if u.op in {Ops.SINK}:
       for s in u.src: to_render.add(s)
     to_render.add(u)

tinygrad/uop/spec.py

@@ -1,5 +1,6 @@
-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, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid
 from tinygrad.helpers import DEBUG, Context, prod
 from tinygrad.uop.validate import validate_index
@@ -239,3 +240,17 @@ def type_verify(uops:list[UOp], check_spec:PatternMatcher):
     if cast(bool|None, ret) is not True:
       if DEBUG >= 3: print_uops(uops)
       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}")
+
+@Context(SPEC=0)
+def validate_pyrender(test_ast:UOp):
+  from tinygrad.codegen.opt import Opt, OptOps
+  from tinygrad.uop.ops import KernelInfo
+  code = '\n'.join(pyrender(test_ast))
+  lcls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Opt": Opt, "OptOps": OptOps}
+  exec(code, None, lcls)
+  if lcls['ast'] is not test_ast:
+    if str(test_ast) == str(lcls['ast']):
+      for u1,u2 in zip(list(test_ast.toposort()), list(lcls['ast'].toposort())):
+        if u1 is not u2:
+          raise RuntimeError("STRING SAME, UOP MISMATCH", u1, u2, id(u1), id(u2), id(u1.arg), id(u2.arg))
+    raise RuntimeError(f"PYRENDER ISSUE:\nCODE:\n{code}\nSTR MATCH: {str(test_ast) == str(lcls['ast'])}\nUOP:\n{test_ast}\nPRODUCED:\n{lcls['ast']}")