tinygrad 0.10.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

tinygrad/runtime/ops_cloud.py

@@ -1,220 +0,0 @@
-# the CLOUD=1 device is a process boundary between the frontend/runtime
-# normally tinygrad is    frontend <-> middleware <-> runtime <-> hardware
-# with CLOUD tinygrad is  frontend <-> middleware <-> CloudDevice ///HTTP/// cloud_server <-> runtime <-> hardware
-# this client and server can be on the same machine, same network, or just same internet
-# it should be a secure (example: no use of pickle) boundary. HTTP is used for RPC
-
-from __future__ import annotations
-from typing import Optional, Any
-from collections import defaultdict
-from dataclasses import dataclass, field
-import multiprocessing, functools, http.client, hashlib, json, time, os, binascii, struct, ast, contextlib
-from http.server import HTTPServer, BaseHTTPRequestHandler
-from tinygrad.renderer import Renderer
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, Timing
-from tinygrad.device import Compiled, Allocator, Compiler, Device, BufferSpec
-
-# ***** API *****
-
-class CloudRequest: pass
-
-@dataclass(frozen=True)
-class BufferAlloc(CloudRequest): buffer_num: int; size: int; options: BufferSpec # noqa: E702
-
-@dataclass(frozen=True)
-class BufferFree(CloudRequest): buffer_num: int # noqa: E702
-
-@dataclass(frozen=True)
-class CopyIn(CloudRequest): buffer_num: int; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class CopyOut(CloudRequest): buffer_num: int
-
-@dataclass(frozen=True)
-class ProgramAlloc(CloudRequest): name: str; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class ProgramFree(CloudRequest): name: str; datahash: str # noqa: E702
-
-@dataclass(frozen=True)
-class ProgramExec(CloudRequest):
-  name: str; datahash: str; bufs: tuple[int, ...]; vals: tuple[int, ...] # noqa: E702
-  global_size: Optional[tuple[int, ...]]; local_size: Optional[tuple[int, ...]]; wait: bool # noqa: E702
-
-# for safe deserialization
-whitelist = {x.__name__:x for x in [BufferAlloc, BufferFree, CopyIn, CopyOut, ProgramAlloc, ProgramFree, ProgramExec, BufferSpec]}
-eval_fxns = {ast.Constant: lambda x: x.value, ast.Tuple: lambda x: tuple(map(safe_eval, x.elts)), ast.List: lambda x: list(map(safe_eval, x.elts)),
-  ast.Call: lambda x: safe_eval(x.func)(*[safe_eval(arg) for arg in x.args], **{kwarg.arg: safe_eval(kwarg.value) for kwarg in x.keywords}),
-  ast.Name: lambda x: whitelist[x.id], ast.Attribute: lambda x: {"imagef": dtypes.imagef, "imageh": dtypes.imageh}[x.attr]}
-def safe_eval(node): return eval_fxns[node.__class__](node)
-
-class BatchRequest:
-  def __init__(self):
-    self._q: list[CloudRequest] = []
-    self._h: dict[str, bytes] = {}
-  def h(self, d:bytes) -> str:
-    binhash = hashlib.sha256(d).digest()
-    self._h[datahash:=binascii.hexlify(binhash).decode()] = binhash+struct.pack("<Q", len(d))+d
-    return datahash
-  def q(self, x:CloudRequest): self._q.append(x)
-  def serialize(self) -> bytes:
-    self.h(repr(self._q).encode())
-    return b''.join(self._h.values())
-  def deserialize(self, dat:bytes) -> BatchRequest:
-    ptr = 0
-    while ptr < len(dat):
-      datahash, datalen = binascii.hexlify(dat[ptr:ptr+0x20]).decode(), struct.unpack("<Q", dat[ptr+0x20:ptr+0x28])[0]
-      self._h[datahash] = dat[ptr+0x28:ptr+0x28+datalen]
-      ptr += 0x28+datalen
-    self._q = safe_eval(ast.parse(self._h[datahash], mode="eval").body)
-    return self
-
-# ***** backend *****
-
-@dataclass
-class CloudSession:
-  programs: dict[tuple[str, str], Any] = field(default_factory=dict)
-  # TODO: the buffer should track this internally
-  buffers: dict[int, tuple[Any, int, Optional[BufferSpec]]] = field(default_factory=dict)
-
-class CloudHandler(BaseHTTPRequestHandler):
-  protocol_version = 'HTTP/1.1'
-  device: str
-  sessions: defaultdict[str, CloudSession] = defaultdict(CloudSession)
-
-  def setup(self):
-    super().setup()
-    print(f"connection established with {self.client_address}, socket: {self.connection.fileno()}")
-
-  def _do(self, method):
-    session = CloudHandler.sessions[unwrap(self.headers.get("Cookie")).split("session=")[1]]
-    ret, status_code = b"", 200
-    if self.path == "/batch" and method == "POST":
-      # TODO: streaming deserialize?
-      req = BatchRequest().deserialize(self.rfile.read(int(unwrap(self.headers.get('Content-Length')))))
-      # the cmds are always last (currently in datahash)
-      for c in req._q:
-        if DEBUG >= 1: print(c)
-        match c:
-          case BufferAlloc():
-            assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated"
-            session.buffers[c.buffer_num] = (Device[CloudHandler.device].allocator.alloc(c.size, c.options), c.size, c.options)
-          case BufferFree():
-            buf,sz,buffer_options = session.buffers[c.buffer_num]
-            Device[CloudHandler.device].allocator.free(buf,sz,buffer_options)
-            del session.buffers[c.buffer_num]
-          case CopyIn(): Device[CloudHandler.device].allocator._copyin(session.buffers[c.buffer_num][0], memoryview(bytearray(req._h[c.datahash])))
-          case CopyOut():
-            buf,sz,_ = session.buffers[c.buffer_num]
-            Device[CloudHandler.device].allocator._copyout(memoryview(ret:=bytearray(sz)), buf)
-          case ProgramAlloc():
-            lib = Device[CloudHandler.device].compiler.compile_cached(req._h[c.datahash].decode())
-            session.programs[(c.name, c.datahash)] = Device[CloudHandler.device].runtime(c.name, lib)
-          case ProgramFree(): del session.programs[(c.name, c.datahash)]
-          case ProgramExec():
-            bufs = [session.buffers[x][0] for x in c.bufs]
-            extra_args = {k:v for k,v in [("global_size", c.global_size), ("local_size", c.local_size)] if v is not None}
-            r = session.programs[(c.name, c.datahash)](*bufs, vals=c.vals, wait=c.wait, **extra_args)
-            if r is not None: ret = str(r).encode()
-    elif self.path == "/renderer" and method == "GET":
-      cls, args = Device[CloudHandler.device].renderer.__reduce__()
-      ret = json.dumps((cls.__module__, cls.__name__, args)).encode()
-    else: status_code = 404
-    self.send_response(status_code)
-    self.send_header('Content-Length', str(len(ret)))
-    self.end_headers()
-    return self.wfile.write(ret)
-
-  def do_GET(self): return self._do("GET")
-  def do_POST(self): return self._do("POST")
-
-def cloud_server(port:int):
-  multiprocessing.current_process().name = "MainProcess"
-  CloudHandler.device = getenv("CLOUDDEV", "METAL") if Device.DEFAULT == "CLOUD" else Device.DEFAULT
-  print(f"start cloud server on {port} with device {CloudHandler.device}")
-  server = HTTPServer(('', port), CloudHandler)
-  server.serve_forever()
-
-# ***** frontend *****
-
-class CloudAllocator(Allocator):
-  def __init__(self, dev:CloudDevice):
-    self.device = dev
-    super().__init__()
-  # TODO: ideally we shouldn't have to deal with images here
-  def _alloc(self, size:int, options:BufferSpec) -> int:
-    self.device.buffer_num += 1
-    self.device.req.q(BufferAlloc(self.device.buffer_num, size, options))
-    return self.device.buffer_num
-  # TODO: options should not be here in any Allocator
-  def _free(self, opaque:int, options): self.device.req.q(BufferFree(opaque))
-  def _copyin(self, dest:int, src:memoryview): self.device.req.q(CopyIn(dest, self.device.req.h(bytes(src))))
-  def _copyout(self, dest:memoryview, src:int):
-    self.device.req.q(CopyOut(src))
-    resp = self.device.batch_submit()
-    assert len(resp) == len(dest), f"buffer length mismatch {len(resp)} != {len(dest)}"
-    dest[:] = resp
-
-class CloudProgram:
-  def __init__(self, dev:CloudDevice, name:str, lib:bytes):
-    self.dev, self.name = dev, name
-    self.datahash = self.dev.req.h(lib)
-    self.dev.req.q(ProgramAlloc(self.name, self.datahash))
-    super().__init__()
-  def __del__(self): self.dev.req.q(ProgramFree(self.name, self.datahash))
-
-  def __call__(self, *bufs, global_size=None, local_size=None, vals:tuple[int, ...]=(), wait=False):
-    self.dev.req.q(ProgramExec(self.name, self.datahash, bufs, vals, global_size, local_size, wait))
-    if wait: return float(self.dev.batch_submit())
-
-class CloudDevice(Compiled):
-  def __init__(self, device:str):
-    if (host:=getenv("HOST", "")) != "": self.host = host
-    else:
-      p = multiprocessing.Process(target=cloud_server, args=(6667,))
-      p.daemon = True
-      p.start()
-      self.host = "127.0.0.1:6667"
-
-    # state for the connection
-    self.session = binascii.hexlify(os.urandom(0x10)).decode()
-    self.buffer_num = 0
-    self.req: BatchRequest = BatchRequest()
-
-    if DEBUG >= 1: print(f"cloud with host {self.host}")
-    while 1:
-      try:
-        self.conn = http.client.HTTPConnection(self.host, timeout=60.0)
-        clouddev = json.loads(self.send("GET", "renderer").decode())
-        break
-      except Exception as e:
-        print(e)
-        time.sleep(0.1)
-    if DEBUG >= 1: print(f"remote has device {clouddev}")
-    # TODO: how to we have BEAM be cached on the backend? this should just send a specification of the compute. rethink what goes in Renderer
-    if not clouddev[0].startswith("tinygrad.renderer.") or not clouddev[1].endswith("Renderer"): raise RuntimeError(f"bad renderer {clouddev}")
-    renderer_class = fromimport(clouddev[0], clouddev[1])  # TODO: is this secure?
-    if not issubclass(renderer_class, Renderer): raise RuntimeError(f"renderer isn't a Renderer {clouddev}")
-    super().__init__(device, CloudAllocator(self), renderer_class(*clouddev[2]), Compiler(), functools.partial(CloudProgram, self))
-
-  def __del__(self):
-    # TODO: this is never being called
-    # TODO: should close the whole session
-    with contextlib.suppress(ConnectionRefusedError, http.client.CannotSendRequest, http.client.RemoteDisconnected): self.batch_submit()
-
-  def batch_submit(self):
-    data = self.req.serialize()
-    with Timing(f"*** send {len(self.req._q):-3d} requests {len(self.req._h):-3d} hashes with len {len(data)/1024:.2f} kB in ", enabled=DEBUG>=1):
-      ret = self.send("POST", "batch", data)
-    self.req = BatchRequest()
-    return ret
-
-  def send(self, method, path, data:Optional[bytes]=None) -> bytes:
-    # TODO: retry logic
-    self.conn.request(method, "/"+path, data, headers={"Cookie": f"session={self.session}"})
-    response = self.conn.getresponse()
-    assert response.status == 200, f"failed on {method} {path}"
-    return response.read()
-
-if __name__ == "__main__": cloud_server(getenv("PORT", 6667))

tinygrad/runtime/support/allocator.py

@@ -1,94 +0,0 @@
-import collections
-from tinygrad.helpers import round_up
-
-class TLSFAllocator:
-  """
-  The allocator is based on the Two-Level Segregated Fit (TLSF) algorithm. The allocator maintains 2 level of buckets:
-    * 1st level is determined by the most significant bit of the size.
-    * 2nd level splits the covered memory of 1st level into @lv2_cnt entries.
-
-  For each allocation request, the allocator searches for the smallest block that can fit the requested size.
-  For each deallocation request, the allocator merges the block with its neighbors if they are free.
-  """
-
-  def __init__(self, size:int, base:int=0, block_size:int=16, lv2_cnt:int=16):
-    self.size, self.base, self.block_size, self.l2_cnt = size, base, block_size, lv2_cnt.bit_length()
-    self.storage:list = [collections.defaultdict(list) for _ in range(size.bit_length() + 1)]
-    self.lv1_entries:list[int] = [0] * len(self.storage)
-
-    # self.blocks is more like a linked list, where each entry is a contigous block.
-    self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free
-    self._insert_block(0, size)
-
-  def lv1(self, size): return size.bit_length()
-  def lv2(self, size): return (size - (1 << (size.bit_length() - 1))) // (1 << max(0, size.bit_length() - self.l2_cnt))
-
-  def _insert_block(self, start:int, size:int, prev:int|None=None):
-    if prev is None: prev = self.blocks[start][2]
-    self.storage[self.lv1(size)][self.lv2(size)].append(start)
-    self.lv1_entries[self.lv1(size)] += 1
-    self.blocks[start] = (size, start + size, prev, True)
-    return self
-
-  def _remove_block(self, start:int, size:int, prev:int|None=None):
-    if prev is None: prev = self.blocks[start][2]
-    self.storage[self.lv1(size)][self.lv2(size)].remove(start)
-    self.lv1_entries[self.lv1(size)] -= 1
-    self.blocks[start] = (size, start + size, prev, False)
-    return self
-
-  def _split_block(self, start:int, size:int, new_size:int):
-    nxt = self.blocks[start][1]
-    assert self.blocks[start][3], "block must be free"
-    self._remove_block(start, size)._insert_block(start, new_size)._insert_block(start + new_size, size - new_size, prev=start)
-    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start + new_size, self.blocks[nxt][3])
-    return self
-
-  def _merge_right(self, start:int):
-    size, nxt, _, is_free = self.blocks[start]
-    assert is_free, "block must be free"
-
-    while is_free and nxt in self.blocks:
-      if (blk:=self.blocks[nxt])[3] is False: break
-      self._remove_block(start, size)._remove_block(nxt, blk[0])._insert_block(start, size:=size + blk[0])
-      assert self.blocks[start][1] == blk[1]
-      _, nxt, _, _ = self.blocks.pop(nxt)
-
-    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start, self.blocks[nxt][3])
-
-  def _merge_block(self, start:int):
-    # Go left while blocks are free. Then merge all them right.
-    while (x:=self.blocks[start][2]) is not None and self.blocks[x][3] is True: start = x
-    self._merge_right(start)
-
-  def alloc(self, req_size:int, align:int=1) -> int:
-    req_size = max(self.block_size, req_size) # at least block size.
-    size = max(self.block_size, req_size + align - 1)
-
-    # Round up the allocation size to the next bucket, so any entry there can fit the requested size.
-    size = round_up(size, (1 << size.bit_length() - self.l2_cnt))
-
-    # Search for the smallest block that can fit the requested size. Start with the it's bucket and go up until any block is found.
-    for l1 in range(self.lv1(size), len(self.storage)):
-      if self.lv1_entries[l1] == 0: continue
-      for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
-        if len(self.storage[l1][l2]) > 0:
-          nsize = self.blocks[self.storage[l1][l2][0]][0]
-          assert nsize >= size, "block must be larger"
-
-          # Block start address.
-          start = self.storage[l1][l2][0]
-
-          # If request contains alignment, split the block into two parts.
-          if (new_start:=round_up(start, align)) != start:
-            self._split_block(start, nsize, new_start - start)
-            start, nsize = new_start, self.blocks[new_start][0]
-
-          # If the block is larger than the requested size, split it into two parts.
-          if nsize > req_size: self._split_block(start, nsize, req_size)
-          self._remove_block(start, req_size) # Mark the block as allocated.
-          return start + self.base
-    raise MemoryError(f"Can't allocate {req_size} bytes")
-
-  def free(self, start:int):
-    self._insert_block(start - self.base, self.blocks[start - self.base][0])._merge_block(start - self.base)

tinygrad/spec.py

@@ -1,155 +0,0 @@
-from typing import cast
-from tinygrad.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops
-from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType
-from tinygrad.helpers import all_same, dedup, prod
-
-buffer_spec = PatternMatcher([
-  (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True),
-  (UPat(Ops.DEVICE, dtypes.void, (), name="device"), lambda device: isinstance(device.arg, str)),
-  (UPat(Ops.BUFFER, src=(UPat(Ops.DEVICE), UPat(Ops.UNIQUE)), name="buf"),
-   lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
-])
-
-# *** this is the spec of a Tensor in UOp ***
-
-tensor_uop_spec = buffer_spec+PatternMatcher([
-  (UPat(GroupOp.Movement, name="mv", src=(UPat.var("x"),)),
-   # naturally correct
-   lambda mv,x: (isinstance(mv.arg, tuple) and mv.dtype == x.dtype) or
-   # "make things that can't be images not images" can change the buffer dtype
-   # this is fine as long as it's a realized buffer and base dtypes match.
-   ((isinstance(mv.dtype, ImageDType) or isinstance(x.dtype, ImageDType)) and x.dtype.base == mv.dtype.base and x.is_realized)),
-  (UPat(Ops.VIEW, src=(UPat(GroupOp.All-{Ops.CONST, Ops.DEVICE}),)), lambda: False),
-
-  # Tensor variable bindings
-  (UPat(Ops.BIND, dtypes.int, (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=dtypes.int)), arg=None), lambda: True),
-
-  # Tensor const has a device and an unmasked ShapeTracker of stride 0
-  (UPat(Ops.CONST, src=(UPat(Ops.VIEW, name="st", src=(UPat(Ops.DEVICE),)),)),
-   lambda st: st.st.views[0].mask is None and len(st.st.views) == 1 and all(s == 0 for s in st.st.views[0].strides)),
-
-  # DETACH and CONTIGUOUS change how we interpret the source UOp
-  # CONTIGUOUS ensures the source UOp realizes
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None), lambda root,x: root.dtype == x.dtype),
-
-  # COPY
-  # NOTE: the arg here specifies clone=True, which prevents folding same device copy
-  (UPat(Ops.COPY, name="copy", src=(UPat(Ops.DEVICE), UPat.var("x"))), lambda copy,x: isinstance(copy.arg, bool) and copy.dtype == x.dtype),
-
-  # ASSIGN changes the value of a realized buffer
-  (UPat(Ops.ASSIGN, name="assign", src=(UPat.var("target"), UPat.var("new_val"))),
-   lambda assign,target,new_val: target.is_realized and (assign.dtype == target.dtype == new_val.dtype)),
-])
-
-# ***** uop type spec *****
-
-# this is the matcher for the final rendered UOps
-# matcher functions returns True or False (or None to not match)
-spec = PatternMatcher([
-  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and not x.dtype.local),
-  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.local),
-  (UPat(Ops.DEFINE_ACC, src=(UPat.var("c"),), name="x", allow_any_len=True),
-   lambda x,c: all(y.op is Ops.RANGE for y in x.src[1:]) and c.dtype == x.dtype),
-  (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)),
-
-  (UPat(Ops.RANGE, src=(UPat.var("x"), UPat.var("y")), name="rng"), lambda rng,x,y: rng.dtype == x.dtype == y.dtype and isinstance(rng.arg, int)),
-  (UPat(Ops.SPECIAL, src=()), lambda: True),
-
-  # TODO: confirm the args of both of these are shapetrackers
-  (UPat(Ops.VIEW, dtypes.void, src=()), lambda: True),
-  (UPat(Ops.VIEW, src=(UPat.var("src"),), name="x"), lambda x,src: src.op is not Ops.STORE and x.dtype.base == src.dtype.base),
-
-  (UPat(Ops.VALID, dtypes.bool, (UPat(Ops.VIEW),)), lambda: True),
-  (UPat(Ops.CONST, name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),
-
-  # early LOAD has a <buf, shapetracker, store?>
-  (UPat(Ops.LOAD, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat(Ops.VIEW))), lambda: True),
-  (UPat(Ops.LOAD, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat(Ops.VIEW), UPat(Ops.STORE))), lambda: True),
-
-  # early STORE has a <buf, shapetracker, val>
-  (UPat(Ops.STORE, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat(Ops.VIEW), UPat())), lambda: True),
-
-  # **** new style load/store ****
-
-  # INDEX is used in new style load/store
-  (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat())), lambda: True),
-
-  # LOAD takes a <bufidx, alt?, gate?, barrier?>
-  (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)),)), lambda: True),
-  (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)), UPat((Ops.IF, Ops.BARRIER)))), lambda: True),
-  (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)), UPat.var("alt"), UPat(dtype=dtypes.bool)), name="ld"), lambda ld,alt: ld.dtype == alt.dtype),
-
-  # STORE takes a <bufidx, val, gate?>
-  (UPat(Ops.STORE, dtype=dtypes.void, src=(UPat((Ops.INDEX, Ops.CAST)), UPat())), lambda: True),
-  (UPat(Ops.STORE, dtype=dtypes.void, src=(UPat((Ops.INDEX, Ops.CAST)), UPat(), UPat(dtype=dtypes.bool))), lambda: True),
-  (UPat(Ops.STORE, dtype=dtypes.void, src=(UPat((Ops.INDEX, Ops.CAST)), UPat(), UPat(Ops.IF))), lambda: True),
-
-  # most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
-  (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))), lambda w,x,y: w.dtype == x.dtype == y.dtype),
-  (UPat((Ops.CMPLT, Ops.CMPNE), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
-  # and SHL/SHR, the shift distance can be an int
-  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
-  (UPat(Ops.IDIV, name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
-  (UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
-
-  (UPat(Ops.ASSIGN, src=(UPat((Ops.DEFINE_ACC, Ops.DEFINE_GLOBAL)), UPat())), lambda: True),
-  (UPat(Ops.ENDRANGE, dtype=dtypes.void, src=(UPat(Ops.RANGE),)), lambda: True),
-
-  # WMMA has a <a, b, acc>
-  (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
-  (UPat(Ops.CONTRACT, name="x"), lambda x: x.dtype.count == prod(y[1] for y in x.arg)),
-  (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
-
-  # if has a <gate, barrier?>
-  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(),)), lambda: True),
-  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(), UPat(Ops.BARRIER))), lambda: True),
-  (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
-
-  (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
-  (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
-  (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.count and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
-  (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None),
-  (UPat(Ops.BARRIER, dtypes.void, src=UPat(Ops.STORE, allow_any_len=True)), lambda: True), # NOTE: all pointers must be local
-
-  # NOTE: for testing, we let sinks be anything
-  #(UPat(Ops.SINK, src=UPat(Ops.STORE)), lambda: True),
-  (UPat((Ops.NAME, Ops.SINK), dtypes.void), lambda: True),
-  (UPat((Ops.NOOP, Ops.CUSTOM)), lambda: True),
-
-  # PTX LOAD/STORE
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(dtype=dtypes.int64),), allow_any_len=True), lambda: True),
-])
-
-# *** this is the spec of a Kernel in UOp ***
-
-kernel_spec = buffer_spec+PatternMatcher([
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.ASSIGN))), lambda: True),
-  # assign has a buffer view and kernel source, it can optionally depend on other assigns
-  (UPat(Ops.ASSIGN, src=UPat((Ops.BUFFER, Ops.VIEW, Ops.KERNEL, Ops.ASSIGN))), lambda: True),
-  # view/sink/const can also exist in the kernel graph
-  (UPat((Ops.VIEW, Ops.SINK, Ops.CONST)), lambda: True),
-  (UPat(GroupOp.All), lambda: False),
-])
-
-# *** this is the UOp shape spec ***
-
-def verify_sink_dims(sink:UOp):
-  shape_dims = [sorted(dedup(dims)) for dims in zip(*[x.shape for x in sink.toposort if x.op is not Ops.SINK and x.st is not None])]
-  return all_same([x.st_arg.size for x in sink.src]) and all(len(x) == 1 or (len(x) == 2 and x[0] == 1) for x in shape_dims)
-
-shape_spec = PatternMatcher([
-  # shapes must have either 1 or n in each dimension
-  (UPat(Ops.SINK, src=UPat(Ops.STORE), allow_any_len=True, name="sink"), verify_sink_dims),
-  # all parent UOps must have the same shape
-  (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: all_same([x.shape for x in root.src if x.st is not None])),
-])
-
-# ***** uop helpers *****
-
-def type_verify(uops:list[UOp], *extra_specs:PatternMatcher):
-  specs = [spec, *extra_specs]
-  for i,u in enumerate(uops):
-    spec_ret = [cast(bool|None, s.rewrite(u)) for s in specs]
-    if any(ret is False for ret in spec_ret) or all(ret is None for ret in spec_ret):
-      print_uops(uops)
-      raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[x.op for x in u.src]} {u.arg}")