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

tinygrad/jit.py

@@ -1,152 +0,0 @@
-from __future__ import annotations
-from typing import Callable, List, Tuple, Dict, cast, Union, Optional, TypeVar, Generic
-import functools, itertools, operator
-from tinygrad.dtype import DType
-from tinygrad.helpers import DEBUG, merge_dicts, getenv, all_int, Context, GRAPH
-from tinygrad.device import Device, JITRunner, CompiledASTRunner, Buffer
-from tinygrad.tensor import Tensor
-from tinygrad.lazy import LazyBuffer
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.symbolic import Variable, NumNode, Node
-from weakref import ref, WeakKeyDictionary
-from dataclasses import dataclass
-
-@dataclass(frozen=True)
-class JitItem:
-  prg: JITRunner  # or a graph executor like MetalGraph
-  rawbufs: List[Optional[Buffer]]
-
-def get_jit_stats(jit_cache: List[JitItem]) -> Tuple[Node, Node]:
-  return functools.reduce(operator.__add__, [ji.prg.op_estimate for ji in jit_cache], NumNode(0)), functools.reduce(operator.__add__, [ji.prg.mem_estimate for ji in jit_cache], NumNode(0))  # noqa: E501
-def get_input_replace(jit_cache: List[JitItem], input_rawbuffers:List[Buffer]) -> Dict[Tuple[int, int], int]:
-  input_replace: Dict[Tuple[int, int], int] = {}
-  for j,ji in enumerate(jit_cache):
-    for i,a in enumerate(ji.rawbufs):
-      if a in input_rawbuffers:
-        input_replace[(j,i)] = input_rawbuffers.index(a)
-  return input_replace
-def get_jc_idxs_with_updatable_launch_dims(jit_cache: List[JitItem]) -> List[int]:
-  return [j for j,ji in enumerate(jit_cache) if isinstance(ji.prg, CompiledASTRunner) and ((ji.prg.global_size and not all_int(tuple(ji.prg.global_size))) or (ji.prg.local_size and not all_int(tuple(ji.prg.local_size))))]  # noqa: E501
-def get_jc_idxs_with_updatable_var_vals(jit_cache: List[JitItem]) -> List[int]:
-  return [j for j,ji in enumerate(jit_cache) if isinstance(ji.prg, CompiledASTRunner) and ji.prg.vars]
-
-class GraphException(Exception): pass
-
-ReturnType = TypeVar('ReturnType')
-class TinyJit(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType]):
-    self.fxn = fxn
-    self.reset()
-
-  def reset(self):
-    self.jit_cache: List[JitItem] = []
-    self.input_replace: Dict[Tuple[int, int], int] = {}
-    self.cnt: int = 0
-    self.ret: Optional[ReturnType] = None
-    self.expected_vals: Optional[Tuple[Variable, ...]] = None
-    self.expected_name_sts_dtype: Optional[Tuple[Tuple[Union[int, str], ShapeTracker, DType], ...]] = None
-
-  # add support for instance methods
-  def __get__(self, obj, objtype): return functools.partial(self.__call__, obj)
-
-  def __call__(self, *args, **kwargs) -> ReturnType:
-    # all inputs (except const) are realized
-    input_tensors: Dict[Union[int, str], LazyBuffer] = {cast(Union[int, str], k):v.realize().lazydata for k,v in itertools.chain(enumerate(args), kwargs.items()) if v.__class__ is Tensor}  # noqa: E501
-    assert all(isinstance(x, LazyBuffer) for x in input_tensors.values()), "multilazybuffer JIT isn't supported"
-    expected_name_sts_dtype = tuple([(k, v.st.unbind(), v.dtype) for k,v in input_tensors.items()])
-
-    # get rawbuffers
-    # TODO: why can .realized have Any type?
-    input_rawbuffers: List[Buffer] = [v.base.realized for v in input_tensors.values() if v.base.realized is not None]
-    assert len(set(input_rawbuffers)) == len(input_rawbuffers), "duplicate inputs to JIT"
-
-    # get variables: they can either be in Tensors or passed in as arguments, and all must be bound. these are all global
-    var_vals: Dict[Variable, int] = merge_dicts([arg.st.var_vals for arg in input_tensors.values()] + [dict(x.unbind() for x in itertools.chain(args, kwargs.values()) if isinstance(x, Variable))])  # noqa: E501
-    expected_vals = tuple(var_vals.keys())
-
-    if self.cnt >= 2:
-      # jit exec
-      assert self.expected_vals == expected_vals, "mismatch of var_vals"
-      assert self.expected_name_sts_dtype == expected_name_sts_dtype, f"mismatch of sts, expected {self.expected_name_sts_dtype} got {expected_name_sts_dtype}"  # noqa: E501
-      for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].rawbufs[i] = input_rawbuffers[input_idx]
-      for ji in self.jit_cache: ji.prg(cast(List[Buffer], ji.rawbufs), var_vals, wait=DEBUG>=2, jit=True)
-    elif self.cnt == 1:
-      # jit capture
-      self.expected_vals, self.expected_name_sts_dtype = expected_vals, expected_name_sts_dtype
-      CacheCollector.start(var_vals)
-      with Context(GRAPH=getenv("JITGRAPH", GRAPH.value)):
-        self.ret = self.fxn(*args, **kwargs)
-      self.jit_cache = CacheCollector.finish()
-      assert len(self.jit_cache) != 0, "didn't JIT anything!"
-      if DEBUG >= 1 and len(set(get_input_replace(self.jit_cache, input_rawbuffers).values())) != len(input_rawbuffers):
-        print("WARNING: some input tensors not found")
-      if DEBUG >= 1: print(f"JIT captured {len(self.jit_cache)} kernels with {len(input_rawbuffers)} inputs")
-
-      # if your Device supports it, condense the items into a graph executor.
-      if (make_graph := Device[Device.DEFAULT].graph) and getenv("JIT") != 2 and len(self.jit_cache) > 1:
-        # Split JIT cache into batches for faster graph execution.
-        # This allows the accelerator to run some batches while subsequent graphs are still being updated.
-        graphed_jit_cache, current_batch = [], []
-        for i,ji in enumerate(self.jit_cache):
-          # If the jit item can potentially be graphed, put it in a batch.
-          if isinstance(ji.prg, CompiledASTRunner): current_batch.append(ji)
-
-          # The flush is done when (1) ji is the last one, (2) the size of batch exceeds the maximum batch size or
-          # (3) the current jit item cannot be graphed, so the current batch is flushed before such a jit item is added.
-          if len(current_batch) > 0 and (i==len(self.jit_cache)-1 or len(current_batch) >= getenv("JIT_BATCH_SIZE", 64) or not isinstance(ji.prg, CompiledASTRunner)):  # noqa: E501
-            try:
-              graphed_jit_cache.append(JitItem(make_graph(current_batch, input_rawbuffers, var_vals), cast(List[Optional[Buffer]], input_rawbuffers)))
-              if DEBUG >= 2: print(f"\tJIT GRAPHing batch with {len(current_batch)} kernels")
-            except GraphException as e:
-              graphed_jit_cache.extend(current_batch)
-              if DEBUG >= 2: print(f"\tJIT GRAPHing failed batch with {len(current_batch)} kernels: {e}")
-            current_batch = []
-
-          # If the jit item cannot be graphed, put it right into the final cache after the flush.
-          if not isinstance(ji.prg, CompiledASTRunner): graphed_jit_cache.append(ji)
-
-        self.jit_cache = graphed_jit_cache
-
-      self.input_replace = get_input_replace(self.jit_cache, input_rawbuffers)
-    elif self.cnt == 0:
-      # jit ignore
-      self.ret = self.fxn(*args, **kwargs)
-
-    # clear jit inputs
-    for (j,i) in self.input_replace.keys(): self.jit_cache[j].rawbufs[i] = None
-
-    self.cnt += 1
-    return cast(ReturnType, self.ret)
-
-class PlaceHolder:
-  def __init__(self, buf:Buffer): self.size, self.dtype, self.device, self.ref, self.bufid = buf.size, buf.dtype, buf.device, ref(buf), id(buf._buf)
-  def to_tuple(self): return (self.size, self.dtype, self.device, self.bufid)
-  def __hash__(self): return hash(self.to_tuple())
-  def __eq__(self, x): return isinstance(x, PlaceHolder) and self.to_tuple() == x.to_tuple()
-  def alloc_if_needed(self, buffer_cache: Dict[PlaceHolder, Buffer]) -> Buffer:
-    ret = self.ref()
-    if ret: return ret
-    if self not in buffer_cache: buffer_cache[self] = Buffer(self.device, self.size, self.dtype)
-    return buffer_cache[self]
-
-class _CacheCollector:
-  def __init__(self):
-    self.cache: Optional[List[Tuple[JITRunner, List[Union[Buffer, PlaceHolder]]]]] = None
-
-  def start(self, var_vals:Optional[Dict[Variable, int]]=None):
-    self.cache = []
-    self.placeholders: WeakKeyDictionary[Buffer, PlaceHolder] = WeakKeyDictionary()
-    self.var_vals = var_vals if var_vals is not None else {}
-
-  def add(self, prg, rawbufs, var_vals):
-    if self.cache is None: return
-    for k,v in var_vals.items(): assert k in self.var_vals and self.var_vals[k] == v, f"var_vals {k} mismatch {v} != {self.var_vals.get(k)}"
-    self.placeholders[rawbufs[0]] = PlaceHolder(rawbufs[0])    # NOTE: this is making an assumption that 0 is special
-    self.cache.append((prg, [self.placeholders.get(x, x) if isinstance(x, Buffer) else x for x in rawbufs]))
-
-  def finish(self) -> List[JitItem]:
-    if self.cache is None: return []
-    buffer_cache: Dict[PlaceHolder, Buffer] = {}
-    saved_cache, self.cache = self.cache, None
-    return [JitItem(prg, [x.alloc_if_needed(buffer_cache) if isinstance(x, PlaceHolder) else x for x in pl]) for prg, pl in saved_cache]
-CacheCollector = _CacheCollector()

tinygrad/realize.py

@@ -1,50 +0,0 @@
-from typing import List, Dict, Optional, cast
-from tinygrad.ops import LoadOps, ScheduleItem, BufferOps, GlobalCounters
-from tinygrad.device import Device, Buffer, BufferCopy, JITRunner, update_stats, InterpretedASTRunner
-from tinygrad.graph import print_tree, realized_lazybuffer
-from tinygrad.helpers import colored, getenv
-from tinygrad.shape.symbolic import Variable
-
-# *** schedule running ***
-
-class CustomOp(JITRunner):
-  def __init__(self, fxn):
-    self.fxn = fxn
-    super().__init__()
-  def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False, jit=False): self.fxn(*rawbufs)
-
-def lower_schedule_item(si:ScheduleItem) -> Optional[JITRunner]:
-  assert all(si.out.device == x.device for x in si.inputs) or si.ast.op is LoadOps.COPY, f"all devices must be the same, {si.out.device} != {[x.device for x in si.inputs]} {print_tree(si.ast) or ''}"  # noqa: E501
-  if si.ast.op is LoadOps.EMPTY: return None
-  if si.ast.op is LoadOps.COPY: return BufferCopy
-  if si.ast.op is LoadOps.CUSTOM: return CustomOp(si.ast.arg)
-  return Device[si.out.device].get_runner(si.ast)
-
-logops = open(getenv("LOGOPS", ""), "a") if getenv("LOGOPS", "") else None
-def run_schedule(schedule:List[ScheduleItem]):
-  while len(schedule):
-    si = schedule.pop(0)
-    if logops and si.ast.op not in LoadOps: logops.write(str(si.ast)+"\n")
-
-    # get the program
-    prg = lower_schedule_item(si)
-
-    # invalidate the output buffer if there's a non contig usage of it in inputs
-    if si.out.output_buffer is not None:
-      for i,a in enumerate(si.inputs):
-        if a.realized == si.out.output_buffer:
-          if any(not x.arg.st.contiguous for x in si.ast.lazyops if x.op == BufferOps.LOAD and x.arg.idx == i+1):
-            si.out.output_buffer = None
-            break
-
-    # we don't have an output buffer, we have to create it, and create to max size if it has symbolic shape
-    si.out.realized = si.out.output_buffer if si.out.output_buffer is not None else \
-      Buffer(si.out.device, si.out.size, si.out.dtype,
-             "PLACEHOLDER" if isinstance(prg, InterpretedASTRunner) else None)
-    del si.out.srcs
-
-    # run the function (put it in JIT)
-    assert all(x.realized is not None for x in si.inputs), f"can't run, some inputs aren't realized {[x for x in si.inputs if x.realized is None]}"
-    if prg: prg.exec([si.out.realized] + [cast(Buffer, x.realized) for x in si.inputs], si.var_vals)
-    else: update_stats(colored(f"empty {si.out.st.size:10d} {si.out.dtype}", "yellow"), 0, 0, {}, None, 1, device=si.out.device)
-    realized_lazybuffer(si.out, GlobalCounters.kernel_count)

tinygrad/runtime/graph/hip.py

@@ -1,24 +0,0 @@
-import ctypes
-from typing import Tuple
-import gpuctypes.hip as hip
-from tinygrad.helpers import init_c_var
-from tinygrad.runtime.ops_hip import check, hip_time_execution
-from tinygrad.runtime.graph.cuda import CUDAGraph
-
-class HIPGraph(CUDAGraph):
-  def __del__(self):
-    check(hip.hipGraphDestroy(self.graph))
-    check(hip.hipGraphExecDestroy(self.instance))
-
-  def encode_args_info(self): return (hip.hipDeviceptr_t, (1,2,3))
-  def graph_create(self): return init_c_var(hip.hipGraph_t(), lambda x: check(hip.hipGraphCreate(ctypes.byref(x), 0)))
-  def graph_instantiate(self, graph):
-    return init_c_var(hip.hipGraphExec_t(), lambda x: check(hip.hipGraphInstantiate(ctypes.byref(x), graph, None, None, 0)))
-  def graph_add_kernel_node(self, graph, c_deps, c_params):
-    return init_c_var(hip.hipGraphNode_t(), lambda x: check(hip.hipGraphAddKernelNode(ctypes.byref(x), graph, c_deps, ctypes.sizeof(c_deps)//8 if c_deps else 0, ctypes.byref(c_params))))  # noqa: E501
-  def graph_launch(self, *args, wait=False): return hip_time_execution(lambda: check(hip.hipGraphLaunch(*args)), enable=wait)
-  def graph_exec_kernel_node_set_params(self, *args): return check(hip.hipGraphExecKernelNodeSetParams(*args))
-  def build_kernel_node_params(self, prg, global_size, local_size, c_config):
-    return hip.hipKernelNodeParams(hip.dim3(*local_size), c_config, ctypes.cast(prg.clprg.prg, ctypes.c_void_p), hip.dim3(*global_size), None, 0)
-  def set_kernel_node_launch_dims(self, node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int]):
-    node.blockDim.x, node.blockDim.y, node.blockDim.z, node.gridDim.x, node.gridDim.y, node.gridDim.z = *local_size, *global_size

tinygrad/runtime/ops_cpu.py

@@ -1,45 +0,0 @@
-import numpy as np
-from typing import Callable, Dict, Tuple
-from tinygrad.helpers import flat_mv
-from tinygrad.ops import BufferOps, UnaryOps, BinaryOps, TernaryOps,  ReduceOps, MovementOps, Op
-from tinygrad.device import Interpreted, Allocator
-
-def reduce_axis(in_shape:Tuple[int, ...], out_shape:Tuple[int, ...]) -> Tuple[int, ...]:
-  assert len(in_shape) == len(out_shape), "reduce shapes must have same dimensions"
-  return tuple(i for i,(a,b) in enumerate(zip(in_shape, out_shape)) if a != b)
-
-def einsum_mulacc(einsum, get_strides, expand):
-  def einscripts(x): return ''.join(["abcdefghijklmnopqrstuvwxyz"[i] for i in x])
-  def axes_slice(strides): return tuple(i for i,s in enumerate(strides) if s != 0), tuple(slice(None) if s != 0 else 0 for s in 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(tuple(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
-
-def as_strided(x, arg):
-  shape, stride, offset = arg
-  return np.ndarray(shape, x.dtype, buffer=np.require(x, requirements='C'), offset=offset*x.dtype.itemsize,
-                    strides=tuple(y*x.dtype.itemsize for y in stride))
-
-numpy_fxn_for_op: Dict[Op, Callable] = {
-  BufferOps.CONST: lambda val, dtype: np.array(val, dtype=dtype.np),
-  UnaryOps.EXP2: np.exp2, UnaryOps.LOG2: np.log2, UnaryOps.SIN: np.sin, UnaryOps.SQRT: np.sqrt,
-  UnaryOps.CAST: lambda x,y: x.view(y[0].np) if y[1] else x.astype(y[0].np, copy=False),
-  UnaryOps.NEG: lambda x: np.logical_not(x) if x.dtype == np.bool_ else np.negative(x),
-  BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: np.less, BinaryOps.CMPEQ: np.equal, BinaryOps.ADD: np.add, BinaryOps.SUB: np.subtract,
-  BinaryOps.MUL: np.multiply, BinaryOps.DIV: lambda x, y: np.divide(x, y).astype(x.dtype, copy=False), BinaryOps.XOR: np.bitwise_xor,
-  ReduceOps.SUM: lambda x, new_shape: x.sum(reduce_axis(x.shape, new_shape), dtype=x.dtype, keepdims=True) if x.shape != new_shape else x,
-  ReduceOps.MAX: lambda x, new_shape: x.max(reduce_axis(x.shape, new_shape), keepdims=True) if x.shape != new_shape else x,
-  TernaryOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy(), optimize=True), lambda x: x.strides, np.broadcast_to),
-  TernaryOps.WHERE: np.where, MovementOps.AS_STRIDED: as_strided, MovementOps.EXPAND: np.broadcast_to, MovementOps.PAD: np.pad
-}
-
-class NumpyAllocator(Allocator):
-  def _alloc(self, size:int): return np.empty(size, dtype=np.uint8)
-  def as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data)
-  def copyin(self, dest:np.ndarray, src:memoryview): np.copyto(dest, np.frombuffer(src, dest.dtype).reshape(dest.shape))
-  def copyout(self, dest:memoryview, src:np.ndarray): np.copyto(np.frombuffer(dest, src.dtype).reshape(src.shape), src)
-
-CPUDevice = Interpreted(NumpyAllocator(), numpy_fxn_for_op)

tinygrad/runtime/ops_hip.py

@@ -1,97 +0,0 @@
-from __future__ import annotations
-import ctypes, functools, subprocess, io
-from typing import Tuple, TypeVar, List
-import gpuctypes.hip as hip
-from tinygrad.helpers import DEBUG, getenv, init_c_var, compile_cuda_style, encode_args_cuda_style, time_execution_cuda_style
-from tinygrad.helpers import from_mv, round_up, to_mv
-from tinygrad.device import Compiled, LRUAllocator, MallocAllocator
-from tinygrad.renderer.cstyle import HIPRenderer
-from tinygrad.codegen.kernel import LinearizerOptions
-
-# The default HIP stream is used for everything.
-MOCKHIP = getenv("MOCKHIP") # for CI. don't run kernels, only check if they compile
-
-def check(status):
-  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
-
-# TODO: remove these helpers, they increase complexity
-def hip_time_execution(cb, enable=False): return time_execution_cuda_style(cb, hip.hipEvent_t, hip.hipEventCreate, hip.hipEventRecord, hip.hipEventSynchronize, hip.hipEventDestroy, hip.hipEventElapsedTime, enable=enable)  # noqa: E501
-
-def compile_hip(prg) -> bytes: return compile_cuda_style(prg, [f'--offload-arch={HIPDevice.default_arch_name}', '-I/opt/rocm/include'], hip.hiprtcProgram, hip.hiprtcCreateProgram, hip.hiprtcCompileProgram, hip.hiprtcGetCode, hip.hiprtcGetCodeSize, hip.hiprtcGetProgramLog, hip.hiprtcGetProgramLogSize, check)  # noqa: E501
-
-class HIPProgram:
-  def __init__(self, device:int, name:str, lib:bytes):
-    self.device, self.name, self.lib = device, name, lib
-
-    if DEBUG >= 6:
-      asm = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib)
-      print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
-
-    if MOCKHIP: return
-    check(hip.hipSetDevice(self.device))
-    self.module = init_c_var(hip.hipModule_t(), lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), lib)))
-    self.prg = init_c_var(hip.hipFunction_t(), lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, name.encode("utf-8"))))
-
-  def __del__(self):
-    if not MOCKHIP: check(hip.hipModuleUnload(self.module))
-
-  def __call__(self, *args, global_size:Tuple[int,int,int], local_size:Tuple[int,int,int], vals:Tuple[int, ...]=(), wait=False):
-    if MOCKHIP: return float("inf")
-    check(hip.hipSetDevice(self.device))
-    return hip_time_execution(lambda: check(hip.hipModuleLaunchKernel(self.prg, *global_size, *local_size, 0, None, None, encode_args_cuda_style(args, vals, hip.hipDeviceptr_t, marks=(1,2,3))[0])), enable=wait)  # noqa: E501
-
-T = TypeVar("T")
-CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
-class HIPAllocator(LRUAllocator):
-  def __init__(self, device:HIPDevice):
-    self.device = device
-    super().__init__()
-  def _alloc(self, size:int):
-    check(hip.hipSetDevice(self.device.device))
-    return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipMalloc(ctypes.byref(x), size)))
-  def _free(self, opaque:T): check(hip.hipFree(opaque))
-  def _hostalloc(self, size:int): return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipHostMalloc(ctypes.byref(x), size, 0)))
-  def copy_from_fd(self, dest, fd, offset, size):
-    check(hip.hipSetDevice(self.device.device))
-    if not hasattr(self, 'hb'): self.hb = [self._hostalloc(CHUNK_SIZE) for _ in range(2)]
-    fo = io.FileIO(fd, "a+b", closefd=False)
-    fo.seek(offset - (minor_offset:=offset % PAGE_SIZE))
-    copied_in = 0
-    for local_offset in range(0, size+minor_offset, CHUNK_SIZE):
-      local_size = min(round_up(size+minor_offset, PAGE_SIZE)-local_offset, CHUNK_SIZE)
-      fo.readinto(to_mv(self.hb[0], local_size))
-      check(hip.hipDeviceSynchronize())
-      check(hip.hipMemcpyAsync(ctypes.c_void_p(dest.value + copied_in), ctypes.c_void_p(self.hb[0].value + minor_offset),
-                               copy_size:=min(local_size-minor_offset, size-copied_in), hip.hipMemcpyHostToDevice, None))
-      copied_in += copy_size
-      self.hb = self.hb[1:] + [self.hb[0]]
-      minor_offset = 0 # only on the first
-  def copyin(self, dest:T, src: memoryview):
-    check(hip.hipSetDevice(self.device.device))
-    host_mem = self._hostalloc(len(src))
-    self.device.pending_copyin.append(host_mem)
-    ctypes.memmove(host_mem, from_mv(src), len(src))
-    check(hip.hipMemcpyAsync(dest, host_mem, len(src), hip.hipMemcpyHostToDevice, None))
-  def copyout(self, dest:memoryview, src:T):
-    check(hip.hipSetDevice(self.device.device))
-    check(hip.hipMemcpy(from_mv(dest), src, len(dest), hip.hipMemcpyDeviceToHost))
-  def transfer(self, dest:T, src:T, sz:int):
-    check(hip.hipSetDevice(self.device.device))
-    # TODO: hipMemcpyAsync, but you have to track the "src" buffer to not free it
-    check(hip.hipMemcpy(dest, src, sz, hip.hipMemcpyDeviceToDevice))
-
-class HIPDevice(Compiled):
-  default_arch_name = "gfx1100"
-  def __init__(self, device:str=""):
-    self.device = int(device.split(":")[1]) if ":" in device else 0
-    self.pending_copyin: List[hip.hipDeviceptr_t] = []
-    if self.device == 0 and not MOCKHIP: HIPDevice.default_arch_name = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device))).gcnArchName.decode()  # noqa: E501
-
-    from tinygrad.runtime.graph.hip import HIPGraph
-    super().__init__(MallocAllocator if MOCKHIP else HIPAllocator(self), LinearizerOptions(device="HIP"), HIPRenderer,
-                     compile_hip, functools.partial(HIPProgram, self.device), HIPGraph)
-  def synchronize(self):
-    check(hip.hipSetDevice(self.device))
-    check(hip.hipDeviceSynchronize())
-    for opaque in self.pending_copyin: check(hip.hipFree(opaque))
-    self.pending_copyin.clear()

tinygrad/runtime/ops_torch.py

@@ -1,49 +0,0 @@
-import torch
-import numpy as np
-from typing import Dict, Callable
-from tinygrad.ops import BufferOps, UnaryOps, BinaryOps, TernaryOps, ReduceOps, MovementOps, Op
-from tinygrad.device import Interpreted, Allocator
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, flatten
-from tinygrad.runtime.ops_cpu import einsum_mulacc, reduce_axis
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu"))
-type_map = {torch.bool: dtypes.bool,
-            torch.int8: dtypes.int8, torch.uint8: dtypes.uint8, torch.int16: dtypes.int16, torch.int32: dtypes.int32, torch.int64: dtypes.int64,
-            torch.float16: dtypes.float16, torch.bfloat16: dtypes.bfloat16, torch.float32: dtypes.float32, torch.float64: dtypes.float64}
-inverse_type_map = {v: k for k,v in type_map.items()}
-# TODO: should unsupported types fail instead of implicit conversion?
-inverse_type_map.update({dtypes.uint16: torch.int16, dtypes.uint32: torch.int32, dtypes.uint64: torch.int64})
-def np_type_cvt(t): return {np.uint32: np.int32, np.uint64: np.int64}.get(t, t)
-
-def as_strided(x, arg):
-  shape, stride, offset = arg
-  x = x.contiguous()
-  offset += x.storage_offset()    # NOTE: contiguous can still have a storage_offset, so we adjust for it
-  if any(i < 0 for i in stride):
-    return torch.as_strided(x, shape, tuple(abs(i) for i in stride),
-      offset + sum((s-1)*a if a < 0 else 0 for (s,a) in zip(shape, stride))).flip([i for i,a in enumerate(stride) if a < 0])
-  return torch.as_strided(x, shape, stride, offset)
-
-torch_fxn_for_op: Dict[Op, Callable] = {
-  # TODO: torch.tensor should work here. it doesn't due to "overflow" in uint8
-  #BufferOps.CONST: lambda val, dtype: torch.tensor(val, device=device, dtype=inverse_type_map[dtype]),
-  BufferOps.CONST: lambda val, dtype: torch.from_numpy(np.array(val, dtype=np_type_cvt(dtype.np))).to(device),
-  UnaryOps.EXP2: torch.exp2, UnaryOps.LOG2: torch.log2, UnaryOps.SIN: torch.sin, UnaryOps.SQRT: torch.sqrt,
-  UnaryOps.CAST: lambda x,y: (x.view if y[1] else x.type)(inverse_type_map[y[0]]),
-  UnaryOps.NEG: lambda x: torch.logical_not(x) if x.dtype is torch.bool else torch.neg(x),
-  BinaryOps.ADD: torch.add, BinaryOps.SUB: torch.sub, BinaryOps.MUL: torch.mul, BinaryOps.DIV: lambda x,y: torch.div(x, y).type(x.dtype),
-  BinaryOps.XOR: torch.bitwise_xor, BinaryOps.MAX: torch.maximum, BinaryOps.CMPLT: torch.lt, BinaryOps.CMPEQ: torch.eq,
-  ReduceOps.SUM: lambda x, new_shape: x.sum(reduce_axis(x.shape, new_shape), dtype=x.dtype, keepdims=True) if x.shape != new_shape else x,
-  ReduceOps.MAX: lambda x, new_shape: x.amax(reduce_axis(x.shape, new_shape), keepdims=True) if x.shape != new_shape else x,
-  TernaryOps.MULACC: einsum_mulacc(lambda s,a,b: torch.einsum(s, a.float(), b.float()).type(a.dtype), lambda x: x.stride(), lambda x,s: x.expand(s)),
-  TernaryOps.WHERE: torch.where, MovementOps.AS_STRIDED: as_strided, MovementOps.EXPAND: lambda x, arg: x.expand(arg),
-  MovementOps.PAD: lambda x, padding: torch.nn.functional.pad(x, flatten(padding[::-1])),
-}
-
-class TorchAllocator(Allocator):
-  def _alloc(self, size:int): return torch.empty([size], device=device, dtype=torch.uint8)
-  def copyin(self, dest:torch.Tensor, src:memoryview): dest.copy_(torch.frombuffer(src, dtype=dest.dtype))
-  def copyout(self, dest:memoryview, src:torch.Tensor): torch.frombuffer(dest, dtype=src.dtype).copy_(src.flatten())
-
-TorchDevice = Interpreted(TorchAllocator(), torch_fxn_for_op)

tinygrad-0.8.0.dist-info/RECORD

@@ -1,41 +0,0 @@
-tinygrad/__init__.py,sha256=FnazNjFEkM_gHbdFHnawSXPH2yWh4HzVlBjz9K1foEc,353
-tinygrad/device.py,sha256=BC-KRXOMVQC9ALfZCpC82r9PieQv4wgOEX4dwW9j9XU,18438
-tinygrad/dtype.py,sha256=Hw-GBXcJ1DO24aKxO4fSc1ZIbOi03hHFS0gj4ScVl4w,5555
-tinygrad/graph.py,sha256=qIl7pLl_31bhxRIEBVqVPgdSSSf6sGoGT9PODCdqsto,5393
-tinygrad/helpers.py,sha256=s4fCWZfTRzg_Y66esZhzd_a48qaPgbSecXsrIUG9x4k,11771
-tinygrad/jit.py,sha256=9u8bozQce2ZFCFP2m2MhKTNyJ3fesXQuE8xm79DIGVU,8743
-tinygrad/lazy.py,sha256=MPKzyOQhMFA6-P9uAHMvn_9vhydrTVRkCOI60kuS7zk,16383
-tinygrad/mlops.py,sha256=IYUZ5eZ5xuLWworlPVwplabMMIRpCYLNJc54U2fSxZk,8628
-tinygrad/ops.py,sha256=5r4t0AKDwwXXC5sOhltUOtwUIMwJrWdZjmjHa9gjir4,5498
-tinygrad/realize.py,sha256=WcMxjNbettjj2O3Y3Q0FuzwOhGKeg3XZDGB9jaZ3pog,2611
-tinygrad/tensor.py,sha256=mFljNbp57lGmEe7J8gaRwfVQqdI8ie35L5QsWCgFg1E,58106
-tinygrad/codegen/kernel.py,sha256=1EVZ6lRFJxd2B9wmgEjj4l4vFEpksFDAPiLuTmq75lA,35020
-tinygrad/codegen/linearizer.py,sha256=16ZjlYLGaFz0DP2jtRMgMLbkvt8JDMeNWLr9OtJ8Cl8,33166
-tinygrad/features/image.py,sha256=L2nGvL63WCWXf5hQKWMYs362cEU4jhtZtpnC5g5s0kY,4843
-tinygrad/features/multi.py,sha256=B231xBGUCCPuFgLTj65Lkv0H7LEQN3SaSP1A23jHTns,6212
-tinygrad/features/search.py,sha256=Sof7SjoQvWvPlkjKaW5zQGEUiQwbwmgQvxhvvJCr3XE,9031
-tinygrad/nn/__init__.py,sha256=wWpObBZfHRmdLR6AC7F07xFpskXzyvndsKR3yrjCE_M,7645
-tinygrad/nn/optim.py,sha256=4989TCNNFHb5_HJe7fhwaqMmHXRZb83xdeiNm8wztwo,3682
-tinygrad/nn/state.py,sha256=UgQRNfMyJOpOdPE87QXfzfpJHCZL2kfZ7Us6WzW4BsM,8294
-tinygrad/renderer/cstyle.py,sha256=IojfT4-N3Fi8bJFOBPYUhotR8BU-dz6DQMsGK3azHXA,16951
-tinygrad/renderer/llvmir.py,sha256=p9FPUimWPleYj4bE1VsylMzDwREkX7e6Q4tQJMN2QAw,10652
-tinygrad/runtime/ops_clang.py,sha256=9RKXciiOr9m6nJjZafR3G3ygHJoroq-C9DxcxPHmAG0,1595
-tinygrad/runtime/ops_cpu.py,sha256=0E4umRHMuaUw97kEEFN3f1HKgsBBMisJU9X6BB-8Bjg,3188
-tinygrad/runtime/ops_cuda.py,sha256=f46g7sNq0VqqL-XPqyJSAiZgXsXF2K3aThCKEYRsnO4,6466
-tinygrad/runtime/ops_disk.py,sha256=rBiUqPF1zDrh9jufkqkHk8bzbhrv4zefqlu49XHWs7w,2835
-tinygrad/runtime/ops_gpu.py,sha256=7dH_aR44Z-9UOtZnjwMaLP0n9B6mV7Xf-J6F2xSivK8,7714
-tinygrad/runtime/ops_hip.py,sha256=WcZ81n_0yqgs-A8FVtmX6pBrdVAI5UhQ3xQzY5VOu9M,5645
-tinygrad/runtime/ops_llvm.py,sha256=X1nn9laLZseHhbKSrjSXXn_FDboknuTe3njaPJ6ueks,2787
-tinygrad/runtime/ops_metal.py,sha256=5JGkOrkjGcvGil-eznA0uAt9HR8d9uDx57OlBscdCXQ,5106
-tinygrad/runtime/ops_torch.py,sha256=fEDBRwT-jUicaa9tAABCvZ7meLGQrh7nlsVhV3IZV5c,3470
-tinygrad/runtime/graph/cuda.py,sha256=UTK1d6Q2Cog83KRNDchMy4Qh6QoocDi3WmsnKNnmsss,5148
-tinygrad/runtime/graph/hip.py,sha256=a47xegB9Li6y0IZeAm_B_q357RTnYC5gAYhNbB7r2qY,1680
-tinygrad/runtime/graph/metal.py,sha256=nKKPKjMRckKfdUAQRMFOo9DaW-GZ0CdFLwvJdR3bl3o,5123
-tinygrad/shape/shapetracker.py,sha256=EzU6zouSs9A6-E8KdjTxG3CWlFbpWruVlyxqxhuj1yE,9843
-tinygrad/shape/symbolic.py,sha256=n90f4njHOJlmkZ87sVJn8MauIjsaAJF0DExCIfXVtcg,16565
-tinygrad/shape/view.py,sha256=1ubLD4q2tnzq3_WBqv40ckwajfXLIZsVMfDgv7PRQBk,12060
-tinygrad-0.8.0.dist-info/LICENSE,sha256=6cp1Hqk0v7NMg1j6OXty_1vAZ4EIwZdCySIoHrCS6RI,1055
-tinygrad-0.8.0.dist-info/METADATA,sha256=2MEwTsFJMDgjuxUImazkcp-vwGKNeLtx4RHhHZdovhg,9529
-tinygrad-0.8.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-tinygrad-0.8.0.dist-info/top_level.txt,sha256=vDABMCWBFQnx2kn9Azueu88FP-1klQdePoHikQhHymc,9
-tinygrad-0.8.0.dist-info/RECORD,,