tinygrad 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

tinygrad/nn/image.py

@@ -1,100 +0,0 @@
-import numpy as np
-from tinygrad.helpers import prod, IMAGE, ImageDType, getenv, dtypes
-from tinygrad.lazy import get_single_root
-
-FLOAT16 = getenv("FLOAT16", 0)
-base_image_type = (100, 2, "imageh", np.float16) if FLOAT16 else (100, 4, "imagef", np.float32)
-
-def image_dot(self, w):
-  # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1)
-  n1, n2 = len(self.shape), len(w.shape)
-  assert n1 != 0 and n2 != 0, f"both arguments to matmul need to be at least 1D, but they are {n1}D and {n2}D"
-  assert self.shape[-1] == w.shape[-min(n2, 2)], f"Input Tensor shapes {self.shape} and {w.shape} cannot be multiplied ({self.shape[-1]} != {w.shape[-min(n2, 2)]})"
-  bs, groups = prod(self.shape[0:-2]), prod(w.shape[0:-2])
-  cin, cout = w.shape[-2], w.shape[-1]
-  out_shape_t = self.shape[0:-2] + (cout,-1)
-  if len(self.shape) > 1:
-    order = tuple(range(len(self.shape)-2)) + (len(self.shape)-1, len(self.shape)-2)
-  else:
-    order, out_shape_t = (0,), (cout, )
-  worder = tuple(range(len(w.shape)-2)) + (len(w.shape)-1, len(w.shape)-2)
-
-  # NOTE: with NHWC we can remove the transposes
-  # bs x groups*cin x H x W
-  cx = self.permute(order=order).reshape(shape=(bs//groups, groups*cin, -1, 1))
-  # groups*cout x cin x H, W
-  cw = w.permute(order=worder).reshape(shape=(groups*cout, cin, 1, 1))
-  return cx.conv2d(cw, groups=groups).reshape(shape=out_shape_t).permute(order=order)
-
-def image_conv2d(self, weight, bias=None, groups=1, stride=1, dilation=1, padding=0):
-  (bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
-  rcout = cout//groups
-  x, w = self, weight.reshape(groups, rcout, cin, H, W)
-
-  # hack for non multiples of 4 on cin
-  if cin % 4 != 0 and not (cin == 1 and groups%4 == 0):
-    x = x.reshape(bs, groups, cin, iy, ix)   # do this always?
-    added_input_channels = 4 - (cin % 4)
-    w = w.pad(tuple((0, added_input_channels) if i == 2 else (0, 0) for i in range(len(w.shape))))
-    x = x.pad(tuple((0, added_input_channels) if i == 2 else (0, 0) for i in range(len(x.shape))))
-    cin = cin + added_input_channels
-    x = x.reshape(bs, groups*cin, iy, ix)
-
-  # hack for non multiples of 4 on rcout
-  added_output_channels = 0
-  if rcout % 4 != 0 and not (rcout == 1 and groups%4 == 0):
-    added_output_channels = 4 - (rcout % 4)
-    rcout += added_output_channels
-    cout = groups * rcout
-    w = w.slice(tuple((0, rcout) if i == 1 else (0, w.shape[i]) for i in range(len(w.shape))))
-
-  # packed (note: flipping bs and iy would make the auto-padding work)
-  x = x.permute(0,2,3,1).reshape(bs * iy, ix * groups * cin//4, 4)
-  cin_last = iy == 1 and ix == 1
-  if cin == 1: w = w.reshape(cout//4,4,H*W).permute(0,2,1)
-  elif cin_last: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,1,3).reshape(cout//4, H*cin//4*W*4, 4)
-  else: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,3,1).reshape(cout//4, H*cin//4*W*4, 4)
-
-  # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
-  if IMAGE >= 2: x,w = x.cast(ImageDType(*base_image_type, shape=x.shape)), w.cast(ImageDType(*base_image_type, shape=w.shape))
-  x, w = x.contiguous(), w.contiguous()
-  if get_single_root(w.lazydata).realized: w.realize()
-
-  # expand out
-  rcin_hi, rcin_lo = cin//4 if cin >= 4 else 1, 4 if cin >= 4 else 1
-  cout_expand = [groups//4 if cin == 1 else groups, 4 if cin == 1 else 1, rcout//4 if rcout >= 4 else 1, 4 if rcout >= 4 else 1]
-  x = x.reshape(bs, iy, ix, groups, rcin_hi, rcin_lo)
-  if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo)
-  else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)
-
-  # padding
-  padding_ = [padding]*4 if isinstance(padding, int) else (padding if len(padding) == 4 else [padding[1], padding[1], padding[0], padding[0]])
-  x = x.slice((None, (-padding_[2], x.shape[1]+padding_[3]), (-padding_[0], x.shape[2]+padding_[1]), None, None, None))
-
-  # prepare input
-  x = x.permute(0,3,4,5,1,2)._pool((H, W), stride, dilation) # -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
-  oy, ox = x.shape[4:6]
-  x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, oy, ox, *cout_expand[0:2], 1, 1, rcin_hi, rcin_lo, H, W)
-  x = x.expand(bs, oy, ox, *cout_expand, rcin_hi, rcin_lo, H, W)
-
-  # prepare weights
-  w = w.permute(0,4,2,5,1,3)
-  w = w.reshape((1, 1, 1, *cout_expand, rcin_hi, rcin_lo, H, W)).expand(x.shape)
-
-  # the conv! (+ the bias)
-  ret = (x*w).cast(dtypes.float32).sum((-4, -3, -2, -1))
-
-  # reshape to image and cast back to image
-  ret = ret.reshape(bs*oy, ox*cout//4, 4)
-  if IMAGE >= 2: ret = ret.cast(ImageDType(*base_image_type, shape=ret.shape))
-  if IMAGE >= 3: ret = ret.contiguous()
-
-  # undo hack for non multiples of 4 on C.rcout
-  if added_output_channels != 0:
-    ret = ret.reshape(bs, oy, ox, groups, rcout)[:, :, :, :, :-added_output_channels]
-    rcout -= added_output_channels
-    cout = groups * rcout
-
-  # NCHW output
-  ret = ret.reshape(bs, oy, ox, cout).permute(0,3,1,2)
-  return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1))

tinygrad/renderer/assembly_arm64.py

@@ -1,169 +0,0 @@
-import struct
-from platform import system
-from typing import Tuple, Dict, List, Optional
-from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.linearizer import UOps, UOp
-from tinygrad.helpers import dtypes, CI
-from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-
-def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
-def compute_offsets(total):
-  quotient, remainder = divmod(total, 4096)
-  return [4096]*quotient + [remainder] if remainder else [4096]*quotient
-
-#NOTE: Darwin needs names to start with a "_"
-def get_name(name): return ('_' if system() == 'Darwin' else '') + name
-
-class ARM64Language(AssemblyLanguage): pass
-
-def specialize_to_arm64(fn_nm, asm):
-  var_size = 16
-  prev_uop:Optional[UOps] = None
-  ins = []
-  x_regs = ['x' + str(i) for i in reversed(range(12))]
-  s_regs = ['s' + str(i) for i in reversed(range(3,32)) if i <= 7 or i >= 16]
-  type_to_reg = {dtypes.double: "d", dtypes.half: 'h', dtypes.float32: 's', dtypes.bool: 'w', dtypes.int8:'w', dtypes.int32: 'w', dtypes.int64: 'x', dtypes.uint8:'w', dtypes.uint32: 'w', dtypes.uint64: 'x'}
-  alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
-          BinaryOps.MOD: "", BinaryOps.CMPLT: "subs",
-          UnaryOps.SIN:'bl ' + get_name('sinf'), UnaryOps.LOG2: 'bl ' + get_name("log2f"), UnaryOps.EXP2: 'bl ' + get_name("exp2f"), UnaryOps.SQRT: 'bl ' + get_name("sqrtf"),
-          TernaryOps.MULACC: "madd", TernaryOps.WHERE: "fcsel"}
-
-  def mov_imm(value, reg):
-    # Manually move value into reg if value can't fit
-    if value.__class__ is not float and abs(value) > abs(65535):
-      ins.append(f"movz w15, #{value & 0xffff}")
-      ins.append(f"movk w15, #{(value >> 16) & 0xffff}, lsl #16")
-      ins.append(f"sxtw {reg}, w15")
-    elif reg[0] == 's':
-      ins.append(f"movz x15, 0x{float_to_hex(value)[4:]}")
-      ins.append(f"movk x15, 0x{float_to_hex(value)[:4]}, lsl #16")
-      ins.append("str x15, [sp, 16]")
-      ins.append(f"ldr {reg}, [sp, 16]")
-    else:
-      ins.append(f"mov {reg}, #{value}")
-
-  # Get variables intervals
-  live_range:Dict[str, List[int]] = {}
-  for i, (uop, out, vin, arg) in enumerate(asm):
-    for var in ([v for v in [out] + vin if v is not None and v.__class__ is not int]):
-      live_range[var.nm] = [i,i] if var.nm not in live_range else [live_range[var.nm][0], i]
-
-  mem_vars:Dict[str, int] = {}
-  rtor:Dict[str, str] = {}
-  def allocate_regs(mvars):
-    nonlocal var_size
-    for v in [v for v in mvars if v is not None and v.__class__ is not int and v.nm not in rtor]:
-      available_regs = s_regs if dtypes.is_float(v[1]) else x_regs
-      #NOTE: Very simple spill, everything that don't fit in regs goes to mem
-      if not available_regs:
-        # ARM needs the stack 16-byte aligned
-        var_size += 16
-        available_regs.append('s0' if dtypes.is_float(out[1]) else 'x12')
-        mem_vars[v.nm] = var_size
-      rtor[v.nm] = available_regs.pop()
-
-  temp_floats = ['s0', 's1', 's2']
-  temp_ints = ['x12', 'x13', 'x16']
-  for i, (uop, out, vin, arg) in enumerate(asm):
-    # Clear regs out of interval
-    for var, reg in list(rtor.items()):
-      available_regs = s_regs if reg[0] == 's' else x_regs
-      if var[1] not in 'B' and var not in mem_vars and i > live_range[var][1]:
-        available_regs.append(rtor.pop(var))
-    # Assign a registers to the variables using live ranges.
-    allocate_regs([out] + vin)
-    # Assign temp regs to vin and load them before direct use
-    for i, v in enumerate([v for v in vin if v.__class__ is not int and v.nm in mem_vars]):
-      rtor[v.nm] = temp_floats[i] if dtypes.is_float(v[1]) else temp_ints[i]
-      # ARM64 addressing constraints https://devblogs.microsoft.com/oldnewthing/20220728-00/?p=106912
-      ins.append(f"mov x15, {mem_vars[v.nm]}")
-      ins.append(f"ldr {rtor[v.nm]}, [sp, x15]")
-
-    if uop == UOps.SPECIAL:
-      if arg.startswith('data'):
-        # data 8 to n into the stack
-        if int(arg[4:]) >= 8:
-          ins.append(f"ldr x15, [x17, #{(int(arg[4:]) - 8) * 8}]")
-          ins.append(f"mov {rtor[out.nm]}, x15")
-      else:
-        ins.append(f"mov {rtor[out.nm]}, #0")
-        ins.append(f"loop_{arg}:")
-    elif uop == UOps.CAST:
-      if arg == BinaryOps.CMPLT:
-        mov_imm(0.0, 's0')
-        mov_imm(1.0, 's1')
-        ins.append(f"fcsel {rtor[out.nm]}, s1, s0, lt")
-      else:
-        ins.append(f"sxtw {rtor[out.nm]}, w{rtor[vin[0].nm][1:]}")
-    elif uop == UOps.ALU:
-      if len(vin)==2 and vin[1].__class__ is int: mov_imm(vin[1], 'x15')
-      if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
-        ins.append(f"ands {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
-      elif arg == TernaryOps.WHERE:
-        ins.append(f"fcmp {rtor[vin[0].nm]}, #0.0")
-        ins.append(f"{alu[arg]} {rtor[out.nm]}, {rtor[vin[1].nm]}, {rtor[vin[2].nm]}, ne")
-      elif arg in [UnaryOps.LOG2, UnaryOps.SIN, UnaryOps.EXP2, UnaryOps.SQRT]:
-        #NOTE: Not a real instruction, use to emulate a ext call in unicorn
-        if CI: ins.append(f"{alu[arg]} {rtor[out.nm]} {rtor[vin[0].nm]}")
-        else:
-          save_regs = [k for k in rtor.keys() if k != out.nm and k not in mem_vars]
-          ins.append(f"sub sp, sp, #{(len(save_regs))*16}")
-          # Save the registers before they are cleared by func call
-          for i,k in enumerate(save_regs,1):
-            ins.append(f"str {rtor[k]}, [sp, #{16*i}]")
-          ins.append("stp x29, x30, [sp, #0]!")
-          ins.append("mov x29, sp")
-          ins.append(f"fmov s0, {rtor[vin[0].nm]}")
-          ins.append(alu[arg])
-          ins.append(f"fmov {rtor[out.nm]}, s0")
-          ins.append("mov sp, x29")
-          ins.append("ldp x29, x30, [sp], #0")
-          for i,k in enumerate(save_regs,1):
-            ins.append(f"ldr {rtor[k]}, [sp, #{16*i}]")
-          ins.append(f"add sp, sp, #{len(save_regs)*16}")
-      elif arg == BinaryOps.CMPLT:
-        ins.append(f"{alu[arg]} {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}" if not dtypes.is_float(vin[0][1]) else f"fcmp {rtor[vin[0].nm]}, {rtor[vin[1].nm]}")
-      elif arg == BinaryOps.MOD:
-        ins.append(f"udiv x14, {rtor[vin[0].nm]}, x15")
-        ins.append(f"msub {rtor[out.nm]}, x14, x15, {rtor[vin[0].nm]}")
-      else:
-        ins.append(f"{'f' if dtypes.is_float(vin[0][1]) else 's' if arg == BinaryOps.DIV else ''}{alu[arg]} {', '.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
-    elif uop == UOps.LOAD:
-      if arg.__class__ in (int, float):
-        mov_imm(arg, rtor[out.nm])
-      else:
-        #NOTE: if need casting load var in s/h0 or x/w12 temp regs
-        reg_in = type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[out.nm]
-        mov_imm(arg[0], "x15")
-        ins.append(f"add x15, {rtor[vin[0].nm]}, x15")
-        ins.append(f"ldr{'sb' if arg[2] is not None and arg[2] in (dtypes.int8, dtypes.uint8, dtypes.bool) else ''} {reg_in}, [x15]")
-        if arg[2] is not None: ins.append(f"{'fcvt' if arg[2] in [dtypes.half, dtypes.double] else 'scvtf'} {rtor[out.nm]}, {reg_in}")
-    elif uop == UOps.STORE:
-      #NOTE: if need casting load var in s/h0 or x/w12 temp regs
-      reg_out = (type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[vin[1].nm])
-      if arg[2] is not None: ins.append(f"fcvt{'zs' if arg[2] not in [dtypes.half, dtypes.double] else '' } {reg_out}, {rtor[vin[1].nm]}")
-      ins.append(f"mov x15, #{arg[0]}")
-      ins.append(f"str {reg_out}, [{rtor[vin[0].nm]}, x15, lsl #0]")
-    elif uop == UOps.COND_BRANCH:
-      #TODO: this is a hack it shouldn't always be a cmp before a cond branch?
-      if prev_uop == UOps.LOAD:
-        ins.append(f"cmp {rtor[vin[0].nm]}, #0")
-      ins.append(f"b.{'lt' if arg[1] else 'ge'} {arg[0][1:]}")
-    elif uop == UOps.LABEL:
-      ins.append(f"{arg[1:]}:")
-    elif uop == UOps.ENDLOOP:
-      mov_imm(arg[0], "x15")
-      ins.append(f"add {rtor[vin[0].nm]}, {rtor[vin[0].nm]}, #1")
-      ins.append(f"cmp {rtor[vin[0].nm]}, x15")
-      ins.append(f"b.lt loop_{arg[1]}")
-    prev_uop = uop
-    # store regs into memory if needed
-    if out is not None and out.nm in mem_vars:
-      ins.append(f"mov x15, {mem_vars[out.nm]}")
-      ins.append(f"str {rtor[out.nm]}, [sp, x15]")
-  return "\n".join([f"//varsize {var_size}",".arch armv8-a",".text", f".global {get_name(fn_nm)}",".p2align 2", f"{get_name(fn_nm)}:", "mov x17, sp"] + [f"sub sp, sp, #{offset}" for offset in compute_offsets(var_size)]+ ins + [f"add sp, sp, #{offset}" for offset in compute_offsets(var_size)] +["ret", "\n"])
-
-def uops_to_arm64_asm(fn_nm:str, uops:List[UOp]) -> Tuple[str, List[int], List[int], bool]:
-  lang = ARM64Language()
-  global_size, local_size = uops_to_asmstyle(lang, fn_nm, uops)
-  return specialize_to_arm64(fn_nm, lang.ins), global_size[::-1], local_size[::-1], True

tinygrad/renderer/assembly_ptx.py

@@ -1,98 +0,0 @@
-from typing import List
-import struct
-from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.linearizer import UOps, UOp
-from tinygrad.helpers import dtypes
-from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_cuda import arch
-
-dtype_to_nvtype = {dtypes.float32: "f32", dtypes.float16: "f16", dtypes.int64: "s64", dtypes.int32: "s32", dtypes.int8: "s8", dtypes.bool: "pred", dtypes.uint64: "u64", dtypes.uint32: "u32", dtypes.uint16: "u16", dtypes.uint8: "u8", "bits16": "b16", dtypes.float64: "f64"}
-def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
-
-def ptx_needs_cast(dest_dtype, src_dtype): return dtypes.is_float(dest_dtype) and dtypes.is_int(src_dtype) or dtypes.is_int(dest_dtype) and dtypes.is_float(src_dtype) or (dtypes.is_float(src_dtype) and dtypes.is_float(dest_dtype) and dest_dtype.itemsize != src_dtype.itemsize)
-
-def render_cast(ins, inp, out):
-  if inp.dtype == dtypes.bool and (dtypes.is_float(out.dtype) or dtypes.is_int(out.dtype)):
-    ins.append(f"selp.{dtype_to_nvtype[out.dtype]} {out}, {'0f3F800000, 0f00000000' if dtypes.is_float(out.dtype) else '1, 0'}, {inp};")
-  elif out.dtype == dtypes.bool:
-    ins.append(f"setp.ne.{dtype_to_nvtype[inp.dtype]} {out}, {'0f00000000' if dtypes.is_float(inp.dtype) else '0'}, {inp};")
-  else:
-    round_mod = ".rzi" if dtypes.is_int(out.dtype) and dtypes.is_float(inp.dtype) else '.rz' if dtypes.is_float(out.dtype) and (dtypes.is_int(inp.dtype) or dtypes.is_float(inp.dtype) and inp.dtype.itemsize > out.dtype.itemsize) else ''
-    ins.append(f"cvt{round_mod}.{dtype_to_nvtype[out.dtype]}.{dtype_to_nvtype[inp.dtype]} {out}, {inp};")
-
-# https://docs.nvidia.com/cuda/parallel-thread-execution/#
-
-class PTXLanguage(AssemblyLanguage):
-  supports_constant_folding: bool = True
-
-def specialize_to_ptx(lang, function_name):
-  param_cnt = 0
-  ins = []
-  alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
-         BinaryOps.MOD: "rem", BinaryOps.CMPLT: "setp.lt", UnaryOps.SQRT: "sqrt.approx",
-         UnaryOps.NOOP: "mov", UnaryOps.SIN: "sin.approx", UnaryOps.LOG2: "lg2.approx", UnaryOps.EXP2: "ex2.approx.ftz",
-         TernaryOps.MULACC: "fma.rn", TernaryOps.WHERE: "selp"}
-  for uop, out, vin, arg in lang.ins:
-    if uop == UOps.ENDLOOP:
-      ins.append("bar.sync 0;")
-    elif uop == UOps.DEFINE_LOCAL:
-      ins.append(f".shared .align 4 .b8 {arg[0]}[{arg[1]*4}];")
-    elif uop == UOps.SPECIAL:
-      if arg.startswith('data'):
-        param_cnt += 1
-        ins.append(f"ld.param.u64 {out}, [{arg}];")
-        # TODO: we sometimes want this to be local, nvcc converts to global most of the time, not sure when we would need to?
-        # ins.append(f"cvta.to.global.u64 {out}, {out};")
-      elif arg.startswith('gid'):
-        ins.append(f"mov.u32 {out}, %ctaid.{'xyz'[int(arg[3:])]};")
-      elif arg.startswith('lid'):
-        ins.append(f"mov.u32 {out}, %tid.{'xyz'[int(arg[3:])]};")
-    elif uop == UOps.ALU:
-      if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
-        ins.append(f"and.pred {out}, {', '.join(str(x) for x in vin)};")
-      else:
-        otype = vin[0].dtype if arg in [BinaryOps.CMPLT] else out.dtype
-        if arg == TernaryOps.WHERE:
-          reg = lang.newreg((vin[0], 'bool'), dtypes.bool)
-          ins.append(f"setp.ne.{dtype_to_nvtype[vin[0].dtype]} {reg}, {'0f00000000' if dtypes.is_float(vin[0].dtype) else '0'}, {vin[0]};")
-          vin = vin[1:] + [reg]
-        ins.append(f"{alu[arg]}{'.lo' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 else ''}{'.rn' if arg == BinaryOps.DIV and out.dtype == dtypes.float32 else ''}.{dtype_to_nvtype[otype]} {out}, {', '.join(str(x) for x in vin)};")
-    elif uop == UOps.LOAD:
-      if arg.__class__ in (int, float):
-        ins.append(f"mov.{dtype_to_nvtype[out.dtype]} {out}, {'0f'+float_to_hex(arg) if dtypes.is_float(out.dtype) else int(arg)};")
-      elif arg[2] is not None and (arg[2] == dtypes.bool or arg[2] != out.dtype):
-        dt = ('u16', dtypes.uint16) if arg[2] == dtypes.bool == out.dtype else ('u8', dtypes.uint8) if arg[2] == dtypes.bool else ('b16', dtypes.float16) if arg[2] == dtypes.half else (dtype_to_nvtype[arg[2]], arg[2])
-        reg = lang.newreg((out, dt[0]), dtype=dt[1])
-        ins.append(f"ld.{arg[1]}.{dt[0]} {reg}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];")
-        render_cast(ins, reg, out)
-      else:
-        ins.append(f"ld.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} {out}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];")
-    elif uop == UOps.STORE:
-      if ptx_needs_cast(dtypes.float if arg[2] is None else arg[2], vin[1].dtype) or arg[2] == dtypes.bool:
-        if arg[2] == dtypes.bool != vin[1].dtype:
-          prereg = lang.newreg((vin[1],'bool'), dtype=dtypes.bool)
-          render_cast(ins, vin[1], prereg)
-        else: prereg = vin[1]
-        reg = lang.newreg((prereg, dtypes.uint16 if arg[2] == dtypes.bool else arg[2]), dtype=dtypes.uint16 if arg[2] == dtypes.bool else dtypes.float if arg[2] is None else arg[2])
-        render_cast(ins, prereg, reg)
-        ins.append(f"st.{arg[1]}.{dtype_to_nvtype['bits16' if arg[2] == dtypes.float16 else dtypes.uint8 if arg[2] == dtypes.bool else dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {reg};")
-      else:
-        ins.append(f"st.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {vin[1]};")
-    elif uop == UOps.CAST:
-      render_cast(ins, vin[0], out)
-    elif uop == UOps.LABEL:
-      ins.append(f"{arg}:")
-    elif uop == UOps.COND_BRANCH:
-      ins.append(f"@{'!' if not arg[1] else ''}{vin[0]} bra {arg[0]};")
-
-  ins_prefix = [".version 7.8", ".target " + arch(), ".address_size 64",
-                f".visible .entry {function_name}({', '.join(f'.param .u64 data{i}' for i in range(param_cnt))}) {{"]
-  for arg in [(dtype, lang.type_to_letter(dtype), c) for dtype,c in lang.cnts.items()]: ins_prefix.append(f".reg .{dtype_to_nvtype[arg[0][0]]} %{arg[1]}<{arg[2]}>;",)
-  ins = ins_prefix + ins
-  ins += ["ret;", "}"]
-  return '\n'.join(ins)
-
-def uops_to_ptx_asm(function_name:str, uops:List[UOp]):
-  lang = PTXLanguage()
-  global_size, local_size = uops_to_asmstyle(lang, function_name, uops)
-  return specialize_to_ptx(lang, function_name), global_size[::-1], local_size[::-1], True

tinygrad/renderer/wgsl.py

@@ -1,53 +0,0 @@
-from tinygrad.renderer.cstyle import render_cl
-from tinygrad.helpers import dtypes, DType
-from tinygrad.renderer.cstyle import CStyleLanguage
-from typing import List, Union
-from tinygrad.ops import UnaryOps, BinaryOps, TernaryOps
-import math
-from typing import Tuple
-
-type_map = {dtypes.float: "f32", dtypes.half: "f16", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "bool"}
-class WGSLLanguage(CStyleLanguage):
-  gid = [f"i32(gindex.{'xyz'[x]})" for x in range(3)]
-  lid = [f"i32(lindex.{'xyz'[x]})" for x in range(3)]
-  size_prefix = "let"
-  barrier="workgroupBarrier();"
-  generic_var_prefix = "var "
-  external_local_bufs = True
-  code_for_op = {
-    UnaryOps.EXP2: lambda x: f"exp2({x})", UnaryOps.LOG2: lambda x: f"log2({x})", UnaryOps.SIN: lambda x: f"sin({x})", UnaryOps.SQRT: lambda x: f"sqrt({x})",
-    BinaryOps.ADD: lambda x,y: f"({x}+{y})", BinaryOps.SUB: lambda x,y: f"({x}-{y})", BinaryOps.MUL: lambda x,y: f"({x}*{y})", BinaryOps.DIV: lambda x,y: f"({x}/{y})",
-    BinaryOps.MAX: lambda x,y: f"max({x},{y})", BinaryOps.CMPLT: lambda x,y: f"f32({x}<{y})",
-    TernaryOps.MULACC: lambda x,y,z: f"fma({x},{y},{z})", TernaryOps.WHERE: lambda a,b,c: f"select({c},{b},{a}!=0.)"
-  }
-
-  def render_local(self, name: str, size: int):
-    return f"var<workgroup> {name}: array<f32,{size}>;"
-
-  def render_const(self, x:Union[float,int], var_dtype) -> str:
-    if math.isnan(x): val = "nan()"
-    elif math.isinf(x): val = ("-" if x < 0 else "") + "0x1.fffffep+127f"
-    else: val = f"{x}" + ("" if dtypes.is_int(var_dtype) else "f")
-    return self.render_cast([val]*var_dtype.sz, var_dtype) if var_dtype.sz > 1 else val
-
-  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], global_size:List[int], local_size:List[int], prekernel:List[str]) -> Tuple[str, List[int], List[int]]:
-    local_size = local_size[::-1] if local_size else [1]
-    bind_it = iter(range(len(bufs)))
-    prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\n"
-    prg += "\n".join(prekernel+[f"@group(0) @binding({next(bind_it)}) var<storage,read_write> {name}: array<{type_map[dtype]}>;" for name,dtype in bufs])
-    prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>, @builtin(local_invocation_id) lindex: vec3<u32>) {{\n" + "\n".join(kernel) + "\n}"
-    return prg, global_size[::-1] if global_size else [1], local_size
-
-  def render_for(self, expr:str, _min:int, _max:Union[int,str]) -> str:
-    return f"for(var {expr} = {_min}; {expr} <= {_max}; {expr}++) {{"
-
-  def render_conditional(self, cond:str, x:str, y:str) -> str:
-    return f"select(f32({y}), {x}, bool({cond}))"
-
-  def render_load(self, output_dtype, buf_name, buf_dtype, idx, local=False) -> str:
-    return f"f32({super().render_load(output_dtype, buf_name, buf_dtype, idx, local)})"
-
-  def render_store(self, buf_name:str, buf_dtype:DType, var_name:str, var_dtype:DType, idx, local=False) -> str:
-    if buf_dtype != var_dtype:
-      var_name = f"{type_map[buf_dtype]}({var_name})"
-    return f"{buf_name}[{idx.render(render_cl)}] = {var_name};"

tinygrad/runtime/lib.py

@@ -1,113 +0,0 @@
-import ctypes
-import numpy as np
-from collections import defaultdict, deque
-from typing import TypeVar, Type, Any, Dict, Deque, Tuple
-from tinygrad.helpers import DType, dtypes, prod, GlobalCounters, ImageDType
-
-_T = TypeVar("_T")
-class RawBuffer:  # pylint: disable=abstract-method
-  def __init__(self, size:int, dtype:DType, buf:Any=None, allocator:Any=None, **kwargs):
-    self.size: int = size
-    self.dtype: DType = dtype
-    self._buf = buf if buf is not None else (allocator.alloc(size, dtype, **kwargs) if allocator else None) # If buf is provided, use it. Otherwise try to allocate from the allocator.
-    self._memsz: int = size*dtype.itemsize
-    self._allocator = allocator
-    GlobalCounters.mem_used += self._memsz
-  def __del__(self):  # NOTE: if it fails on init (bad dtype), it won't have a _memsz
-    if hasattr(self, '_memsz'): GlobalCounters.mem_used -= self._memsz
-    if hasattr(self, '_allocator') and self._allocator: self._allocator.free(self._buf)
-  def __repr__(self): return f"buffer<{self.size}, {self.dtype}>"
-  @property
-  def key(self): return (self.size, self.dtype)
-
-  # NOTE: this interface allows for 0 copy
-  @classmethod
-  def fromCPU(cls:Type[_T], x:np.ndarray) -> _T: raise NotImplementedError("must be implemented")
-  def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
-
-class RawConst(RawBuffer): # pylint: disable=abstract-method
-  def __repr__(self): return f"const<{self._buf}, {self.dtype}>"
-  @property
-  def key(self): return (str(self._buf), self.dtype)
-
-def buf_is_kernel_arg(x) -> bool:
-  return x.realized is not None and x.realized.__class__ is not RawConst
-
-# --teenygrad--
-
-class RawBufferCopyIn(RawBuffer):
-  def _copyin(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented")
-
-  @classmethod
-  def fromCPU(cls, x:np.ndarray, **kwargs):
-    ret = cls(prod(x.shape), dtypes.from_np(x.dtype), **kwargs)
-    ret._copyin(x)
-    return ret
-
-class RawBufferMapped(RawBufferCopyIn):
-  def _buffer(self) -> memoryview: raise NotImplementedError("must be implemented")
-  # NOTE: this metadata prevents the backing buffer from being freed. hack can be removed with PEP688
-  def toCPU(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=np.dtype(self.dtype.np, metadata={"backing": self}))  # type: ignore
-  def _copyin(self, x:np.ndarray) -> None: np.copyto(self.toCPU(), x.reshape(-1))
-
-# this one is simple enough that i moved it out of the runtimes
-class RawMallocBuffer(RawBufferMapped):
-  def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float64:ctypes.c_double, dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.bfloat16: ctypes.c_int16, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int32: ctypes.c_int32, dtypes.uint32: ctypes.c_uint32, dtypes.int64: ctypes.c_int64, dtypes.uint64: ctypes.c_uint64}[dtype] * size)())
-  def _buffer(self): return memoryview(self._buf)
-
-class RawBufferCopyInOut(RawBufferCopyIn):
-  def _copyout(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented")
-
-  def toCPU(self) -> np.ndarray:
-    x: np.ndarray = np.empty(self.size, dtype=self.dtype.np)
-    self._copyout(x)
-    return x
-
-class RawBufferTransfer(RawBuffer):
-  def _transfer(self, x) -> None: raise NotImplementedError("must be implemented")
-
-  @classmethod
-  def transfer(cls, x, shape, dtype, **kwargs):
-    ret = cls(prod(shape), dtype, **kwargs)
-    ret._transfer(x)
-    return ret
-
-class LRUAllocator:
-  def __init__(self, dev_memsz=(4<<30)):
-    self.epoch = 0
-    self.free_space: Dict[Any, int] = defaultdict(lambda: dev_memsz)
-    self.buffer_info: Dict[Any, Tuple[int, DType, str]] = dict()
-    self.cached_buffers: Dict[Tuple[int, ...], Deque[Tuple[Any, int]]] = defaultdict(deque) # Cached buffer storage, splitted by type and size, newest first.
-    self.aging_order: Dict[Any, Deque[Tuple[Tuple[int, ...], int]]] = defaultdict(deque) # Keys of cached_buffers, ordered from oldest to newest updates.
-  def __del__(self):
-    for v in self.cached_buffers.values():
-      for buf, _ in v: self._free_buffer(buf)
-  def _cache_reuse_buffer(self, rawbufs: Deque[Tuple[Any, int]]): # The newest cached buffer is reused.
-    GlobalCounters.mem_cached -= self._underlying_buf_memsz(rawbufs[0][0])
-    return rawbufs.popleft()[0]
-  def _alloc_buffer(self, size, dtype, device, **kwargs):
-    self.free_space[device] -= size*dtype.itemsize
-    while len(self.aging_order[device]) and self.free_space[device] < 0: # When OOM removing lru buffers.
-      bucket, epoch = self.aging_order[device].popleft()
-      if self.cached_buffers[bucket] and self.cached_buffers[bucket][-1][1] == epoch: self._free_buffer(self.cached_buffers[bucket].pop()[0]) # Free cached buffer if it is still in cache.
-    newbuf = self._do_alloc(size, dtype, device, **kwargs)
-    self.buffer_info[newbuf] = (size, dtype, device)
-    return newbuf
-  def _free_buffer(self, buf_to_free):
-    self.free_space[self.buffer_info[buf_to_free][2]] += self._underlying_buf_memsz(buf_to_free)
-    GlobalCounters.mem_cached -= self._underlying_buf_memsz(buf_to_free)
-    self.buffer_info.pop(buf_to_free)
-    self._do_free(buf_to_free)
-  def alloc(self, size, dtype, device='0', **kwargs):
-    rawbufs = self.cached_buffers.get(self._cached_bufkey(size, dtype, device), None)
-    return self._cache_reuse_buffer(rawbufs) if rawbufs else self._alloc_buffer(size, dtype, device, **kwargs)
-  def free(self, buf): # free() just caches buffer. It might be freed later when OOM during allocation.
-    self.epoch += 1
-    size, dtype, device = self.buffer_info[buf]
-    self.cached_buffers[self._cached_bufkey(size, dtype, device)].appendleft((buf, self.epoch))
-    self.aging_order[device].append((self._cached_bufkey(size, dtype, device), self.epoch))
-    GlobalCounters.mem_cached += self._underlying_buf_memsz(buf)
-  def _underlying_buf_memsz(self, buf): return self.buffer_info[buf][0] * self.buffer_info[buf][1].itemsize
-  def _cached_bufkey(self, size, dtype, device) -> Tuple[int, ...]: return (device, size, dtype, dtype.shape) if isinstance(dtype, ImageDType) else (device, size, dtype) # Provides a key for reusing device buffers with identical keys.
-  def _do_alloc(self, size, dtype, device, **kwargs): raise NotImplementedError("must be implemented")
-  def _do_free(self, buf): pass

tinygrad/runtime/ops_cpu.py

@@ -1,51 +0,0 @@
-import numpy as np
-import operator
-from typing import Callable, Dict, Tuple, Optional
-from tinygrad.helpers import dtypes, DType
-from tinygrad.ops import UnaryOps, BinaryOps, MovementOps, ReduceOps, TernaryOps, Op, Interpreted
-from tinygrad.runtime.lib import RawBuffer
-
-def shape_to_axis(old_shape:Tuple[int, ...], new_shape:Tuple[int, ...]) -> Tuple[int, ...]:
-  assert len(old_shape) == len(new_shape), "reduce shapes must have same dimensions"
-  return tuple(i for i,(a,b) in enumerate(zip(old_shape, new_shape)) if a != b)
-
-base_fxn_for_op: Dict[Op, Callable] = {
-  BinaryOps.ADD: operator.add, BinaryOps.SUB: operator.sub, BinaryOps.MUL: operator.mul, BinaryOps.DIV: operator.truediv,
-  ReduceOps.SUM: lambda x, new_shape: x.sum(shape_to_axis(x.shape, new_shape), keepdims=True) if tuple(x.shape) != tuple(new_shape) else x[:],
-  ReduceOps.MAX: lambda x, new_shape: (x.amax if hasattr(x, 'amax') else x.max)(shape_to_axis(x.shape, new_shape), keepdims=True) if tuple(x.shape) != tuple(new_shape) else x[:],
-  MovementOps.RESHAPE: lambda x, arg: x.reshape(arg), MovementOps.SHRINK: lambda x, arg: x[tuple(slice(p[0], p[1], None) for p in arg)],
-}
-
-def promote_types(x, y): return ret if (ret := np.promote_types(x.dtype, y.dtype)) != np.float64 else np.float32
-def match_types(x, y):
-  up = x.dtype if dtypes.from_np(x.dtype).priority > dtypes.from_np(y.dtype).priority else y.dtype
-  return x.astype(up, copy=False), y.astype(up, copy=False)
-
-def einsum_mulacc(einsum, get_strides, expand):
-  def einscripts(x): return ''.join(["abcdefghijklmnopqrstuvwxyz"[i] for i in x])
-  def axes_slice(strides): return [i for i in range(len(strides)) if strides[i] != 0], tuple(slice(None) if strides[i] != 0 else 0 for i in range(len(strides)))
-  def mulacc(a, b, new_shape):
-    (a_axes, a_slices), (b_axes, b_slices) = axes_slice(get_strides(a)), axes_slice(get_strides(b))
-    out = [i for i in range(len(new_shape)) if a.shape[i] == new_shape[i] and (i in a_axes or i in b_axes)]
-    ret = einsum(f"{einscripts(a_axes)}, {einscripts(b_axes)} -> {einscripts(out)}", a[a_slices], b[b_slices])
-    return expand(ret.reshape([(1 if i not in a_axes and i not in b_axes else s) for i,s in enumerate(new_shape)]), new_shape)
-  return mulacc
-
-numpy_fxn_for_op: Dict[Op, Callable] = {**base_fxn_for_op, **{
-  UnaryOps.NOOP: lambda x: np.require(x, requirements='C'), UnaryOps.EXP2: np.exp2, UnaryOps.LOG2: np.log2, UnaryOps.SIN: np.sin,
-  UnaryOps.CAST: lambda x,y: x.view(y[0].np) if y[1] else x.astype(y[0].np, copy=False),
-  BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: lambda x,y: (x<y).astype(promote_types(x,y)), BinaryOps.ADD: lambda x, y: np.add(*match_types(x, y)),
-  BinaryOps.SUB: lambda x, y: np.subtract(*match_types(x, y)), BinaryOps.MUL: lambda x, y: np.multiply(*match_types(x, y)),
-  BinaryOps.DIV: lambda x, y: np.divide(*match_types(x, y)), UnaryOps.SQRT: np.sqrt,
-  MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to,
-  MovementOps.STRIDE: lambda x, arg: x[tuple(slice(None, None, i) for i in arg)],
-  TernaryOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, *match_types(a.copy(), b.copy()), optimize=True), lambda x: x.strides, np.broadcast_to),
-  TernaryOps.WHERE: np.where,
-}}
-
-class RawNumpyBuffer(RawBuffer):
-  def __init__(self, size:int, dtype:DType, buf:Optional[np.ndarray]=None): super().__init__(size, dtype, buf if buf is not None else np.empty([size], dtype.np))
-  @classmethod
-  def fromCPU(cls, x): return cls(x.size, dtypes.from_np(x.dtype), x)
-  def toCPU(self): return self._buf
-CPUBuffer = Interpreted(RawNumpyBuffer, numpy_fxn_for_op, from_underlying=RawNumpyBuffer.fromCPU)