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

tinygrad/engine/lazy.py

@@ -1,228 +0,0 @@
-from __future__ import annotations
-from typing import Optional, Any, Tuple, List, get_args
-from tinygrad.dtype import dtypes, DType, ConstType, to_dtype, ImageDType
-from tinygrad.helpers import prod, getenv, all_int, all_same, DEBUG, _METADATA, Metadata, SPLIT_REDUCEOP, LAZYCACHE
-from tinygrad.ops import exec_alu, python_alu
-from tinygrad.ops import identity_element, MathTrait, resolve, UOp, sint, GroupOp, Ops
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.device import Buffer
-from weakref import ref, ReferenceType, WeakValueDictionary
-
-lazycache: WeakValueDictionary[Any, LazyBuffer] = WeakValueDictionary()
-def create_lazybuffer(device:str, st:ShapeTracker, dtype:DType, op:Optional[Ops]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),
-                      base:Optional[LazyBuffer]=None, enable_cache=bool(LAZYCACHE)):
-  if st.size == 0: op, arg, srcs, base = Ops.CONST, 0, (), None
-  dtype = to_dtype(dtype)
-  if op is Ops.CONST: arg, enable_cache = dtypes.as_const(arg, dtype) if not isinstance(arg, UOp) else arg, True
-
-  cache_key = (device, st, dtype, op, arg, tuple(ref(x) for x in srcs)) if base is None else (st, ref(base))
-  if enable_cache and (rret := lazycache.get(cache_key, None)) is not None: return rret
-
-  ret = LazyBuffer(device, st, dtype, op, arg, srcs, base=base, metadata=_METADATA.get())
-  if enable_cache: lazycache[cache_key] = ret
-  return ret
-
-view_supported_devices = {"LLVM", "CLANG", "CUDA", "NV", "AMD", "METAL", "QCOM", "DSP", "DISK"}
-class LazyBuffer(MathTrait):
-  def __init__(self, device:str, st:ShapeTracker, dtype:DType,
-               op:Optional[Ops]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),
-               base:Optional[LazyBuffer]=None, metadata:Optional[Metadata]=None):
-    self.device, self.st, self.dtype, self.shape, self.size, self.metadata = device, st, to_dtype(dtype), st.shape, st.size, metadata
-    self._base: Optional[LazyBuffer] = None
-    if base is None:
-      # properties on base
-      self.op, self.arg, self.srcs = op, arg, srcs  # this is a UOp, except the src is LazyBuffers and not UOps
-      assert self.op is not Ops.ASSIGN or srcs[0].base.realized is not None, "assign target must be realized"
-
-      if self.op is Ops.BUFFER_VIEW:
-        # some LazyBuffers can be processed with only a view, no AST required
-        self.buffer: Buffer = srcs[0].base.buffer.view(st.size, self.dtype, srcs[0].st.views[0].offset * srcs[0].dtype.itemsize)
-      else:
-        self.buffer = srcs[0].base.buffer if self.op is Ops.ASSIGN else Buffer(device, self.size, self.dtype)
-      self.buffer.ref(1)
-      self.contiguous_child: Optional[Tuple[ReferenceType[LazyBuffer], ShapeTracker]] = None
-      self.forced_realize = False
-    else:
-      # properties on view
-      assert base.base == base, "base must be a base itself"
-      self._base = base
-
-  def __del__(self):
-    if hasattr(self, 'buffer'): self.buffer.ref(-1)
-
-  def __repr__(self) -> str:
-    return f"<LB {self.device} {self.shape} {str(self.dtype)[7:]} {self.st if self.base is not self else (self.op, self.realized)}>"
-
-  @property
-  def realized(self) -> Optional[Buffer]:
-    # NOTE: we check for a lack of srcs instead of an allocated buffer to make unrealized assigns return None here
-    return self.buffer if self._base is None and not hasattr(self, 'srcs') else None
-
-  # NOTE: this has to be a function to prevent self reference
-  @property
-  def base(self) -> LazyBuffer: return self._base if self._base is not None else self
-
-  # same API as multi
-  @property
-  def lbs(self) -> List[LazyBuffer]: return [self]
-
-  @staticmethod
-  def metaop(op, shape:Tuple[sint,...], dtype:DType, device:str, arg=None, src:Tuple[LazyBuffer, ...]=(), enable_cache=False) -> LazyBuffer:
-    assert isinstance(src, tuple)
-    return create_lazybuffer(device, ShapeTracker.from_shape(shape), dtype, op, arg, src, enable_cache=enable_cache)
-
-  def const_like(self, b): return self.const_with_shape(b, self.shape)
-  def const_with_shape(self, val:ConstType, shape:Tuple[sint,...]) -> LazyBuffer:
-    assert isinstance(val, get_args(ConstType)), f"{val=} has {type(val)=}, not a ConstType"
-    return LazyBuffer.metaop(Ops.CONST, tuple(), self.dtype, self.device, arg=val).reshape((1,)*len(shape)).expand(shape)
-
-  @property
-  def is_realized(self) -> bool: return self.base.realized is not None
-
-  def assign(self, x:LazyBuffer) -> LazyBuffer:
-    assert x.size == self.size, f"assign target must have same size {self.size=} != {x.size=}"
-    assert self.is_realized, f"assign target must be realized {self}"
-    return LazyBuffer.metaop(Ops.ASSIGN, self.shape, self.dtype, self.device, arg=() if self.st.contiguous else (self.st,),
-                             src=(self.base, x), enable_cache=True)
-
-  def can_view(self):
-    return (self.st.consecutive and not self.is_unrealized_const() and not isinstance(self.dtype, ImageDType) and
-            self.device.split(":")[0] in view_supported_devices)
-
-  def contiguous(self, allow_buffer_view=True):
-    if not self.st.contiguous or self.size != self.base.size or self.is_unrealized_const():
-      ret = self.alu(Ops.BUFFER_VIEW) if allow_buffer_view and self.can_view() else self.alu(Ops.CONTIGUOUS)
-      if (sti := self.st.invert(self.base.shape)) is not None: self.base.contiguous_child = ref(ret), sti
-      return ret
-    self.base.forced_realize = True
-    return self
-
-  def bitcast(self, dtype:DType) -> LazyBuffer: return self.cast(dtype, bitcast=True)
-  def cast(self, dtype:DType, bitcast:bool=False, allow_buffer_view=True) -> LazyBuffer:
-    if self.dtype == dtype: return self
-    if self.device.startswith("DISK") and not bitcast: raise RuntimeError("attempted to cast disk buffer (bitcast only)")
-    if self.is_unrealized_unmasked_const() and not bitcast:
-      return create_lazybuffer(self.device, self.st, dtype, Ops.CONST, dtypes.as_const(self.base.arg, dtype))
-    new_shape = self.shape
-    if bitcast and self.dtype.itemsize != dtype.itemsize:
-      if not self.device.startswith("DISK"): raise RuntimeError("shape changing bitcast only supported on DISK right now")
-      if not all_int(new_shape): raise RuntimeError("shape changing bitcast with symbolic shape isn't supported yet")
-      # https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
-      if not (new_shape[-1]*self.dtype.itemsize) % dtype.itemsize == 0: raise RuntimeError("unsupported size in bitcast")
-      new_shape = new_shape[:-1] + ((new_shape[-1]*self.dtype.itemsize) // dtype.itemsize,)
-    elif getenv("CAST_BEFORE_VIEW", 1) and dtype.itemsize <= self.dtype.itemsize and self is not self.base:
-      # TODO: applying this makes gpt2 slower
-      return self.base.cast(dtype, bitcast)._view(self.st)
-    cast_op: Ops = (Ops.BUFFER_VIEW if self.can_view() and allow_buffer_view else Ops.BITCAST) if bitcast else Ops.CAST
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(new_shape), dtype, cast_op, dtype, (self,))
-
-  def is_unrealized_const(self): return self.base.realized is None and self.base.op is Ops.CONST and not isinstance(self.base.arg, UOp)
-  def is_unrealized_unmasked_const(self): return self.is_unrealized_const() and all(v.mask is None for v in self.st.views)
-
-  def _copy(self, device:str) -> LazyBuffer:
-    assert self.st.contiguous and self.size == self.base.size, f"can only copy contig {self} {self.base}"
-    return create_lazybuffer(device, ShapeTracker.from_shape(self.shape), self.dtype, Ops.COPY, self.buffer.nbytes, (self,), enable_cache=False)
-
-  def copy_to_device(self, device:str, force:bool=False, clone:bool=False) -> LazyBuffer:
-    # no COPY
-    if self.device == device and not clone: return self
-
-    # double COPY = one COPY
-    if not force and self.st.contiguous and self.size == self.base.size and not self.base.realized and self.base.op is Ops.COPY:
-      return self.base.srcs[0].copy_to_device(device).reshape(self.st.shape)
-
-    # const doesn't have to be copied (issues with disk tensor)
-    if self.is_unrealized_const():
-      return LazyBuffer.metaop(Ops.CONST, tuple(), self.dtype, device, arg=self.base.arg)._view(self.st)
-
-    # if it's a shrink, do the shrink before the copy with CONTIGUOUS
-    if prod(self.st.shape) < prod(self.base.st.shape): return self.contiguous()._copy(device)
-
-    # copy the base and apply the shapetracker on the new device
-    return self.base._copy(device)._view(self.st)
-
-  def clone(self) -> LazyBuffer: return self.copy_to_device(self.device, clone=True)
-
-  def alu(self, op:Ops, *in_srcs:LazyBuffer) -> LazyBuffer:
-    srcs: List[LazyBuffer] = []
-    for s in (self,)+in_srcs:
-      if s == s.base and s.base.contiguous_child and (root:=s.base.contiguous_child[0]()) is not None:
-        srcs.append(root._view(s.base.contiguous_child[1]))
-      else:
-        srcs.append(s)
-    if not all_same(dts:=[x.dtype.base for x in (srcs[1:] if op is Ops.WHERE else srcs)]):
-      raise AssertionError(f"all dtypes must match {dts} on {op}")
-    assert all_same([x.shape for x in srcs]), f"all shapes must be the same {[x.shape for x in srcs]}"
-    if op is Ops.WHERE: assert srcs[0].dtype == dtypes.bool, "Ops.WHERE must have the first arg be bool"
-
-    out_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else srcs[-1].dtype
-
-    # const folding
-    if op in python_alu and all(s.is_unrealized_unmasked_const() for s in srcs):
-      return self.cast(out_dtype).const_like(exec_alu(op, out_dtype, [s.base.arg for s in srcs]))
-    if op in GroupOp.Binary:
-      x, y = self, in_srcs[0]
-      if op is Ops.ADD:
-        if y.is_unrealized_unmasked_const() and y.base.arg == 0: return x
-        if x.is_unrealized_unmasked_const() and x.base.arg == 0: return y
-      if op is Ops.MUL:
-        if x.is_unrealized_unmasked_const() and (val := x.base.arg) in (1, 0): return y if val == 1 else y.const_like(0)
-        if y.is_unrealized_unmasked_const() and (val := y.base.arg) in (1, 0): return x if val == 1 else x.const_like(0)
-      if op is Ops.IDIV and y.is_unrealized_unmasked_const() and y.base.arg == 1: return x
-
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(self.shape), out_dtype, op, None, tuple(srcs))
-
-  # *** reduce ops ***
-
-  def _reduce_op(self, op:Ops, axis:Tuple[int, ...]) -> LazyBuffer:
-    assert all(0 <= x < len(self.shape) for x in axis), f"axis args {axis} out of range for shape {self.shape}"
-    axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
-    if len(axis) == 0: return self
-    return create_lazybuffer(self.device, ShapeTracker.from_shape(self.st.reduce(axis)), self.dtype, Ops.REDUCE_AXIS, (op, axis), (self,))
-
-  def r(self, op:Ops, axis:Tuple[int, ...]) -> LazyBuffer:
-    new_shape = self.st.reduce(axis)
-    # TODO: this logic should move to the scheduler
-    if 0 in self.shape and 0 not in new_shape: return self.const_with_shape(identity_element(op, self.dtype), new_shape)
-
-    # const folding
-    # TODO: fold this for symbolic?
-    if self.is_unrealized_unmasked_const() and all_int(self.shape):
-      if op is Ops.ADD: return self.const_with_shape(self.base.arg * prod(self.shape[i] for i in axis), new_shape)
-      if op is Ops.MUL: return self.const_with_shape(self.base.arg ** prod(self.shape[i] for i in axis), new_shape)
-      if op is Ops.MAX: return self.const_with_shape(self.base.arg, new_shape)
-
-    # TODO: can we split symbolic shape if the reduce axis is not symbolic?
-    if not SPLIT_REDUCEOP or not all_int(self.shape) or (0 in self.shape) or \
-      prod(self.shape) // prod(new_shape) < getenv("REDUCEOP_SPLIT_THRESHOLD", 32768):
-      return self._reduce_op(op, axis)
-
-    # if there are few globals, make some reduces into globals by splitting into two kernels
-    # cap output buffer to 2**22: heuristic number of global outputs to achieve max occupancy with enough locals+upcasts for gemm
-    #   ~2**10 should be enough if GROUP is used
-    # 256 split maximum should be "negligible reduce" for low prod(new_shape), 8 split minimum.
-    # split is moved to the end to provide maximum locality for the second phase reduce.
-    self_real_strides = self.st.real_strides(ignore_valid=True)
-    split_candidates = [(i, x) for i in axis for x in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(new_shape)),8-1,-1)
-                        if self.shape[i] % x == 0 and self_real_strides[i] != 0]
-    if not split_candidates: return self._reduce_op(op, axis)
-    dim_to_split, divisor = split_candidates[0]
-    splitted_shape = self.shape[:dim_to_split] + (divisor,) + (self.shape[dim_to_split]//divisor,) + self.shape[dim_to_split+1:]
-    splitted = self.reshape(splitted_shape).permute(tuple([x for x in range(len(splitted_shape)) if x != dim_to_split]+[dim_to_split]))
-    if DEBUG >= 3: print(f"split {divisor}: {self.shape} -> {splitted.shape} -> {new_shape}")
-    return splitted._reduce_op(op, axis)._reduce_op(op, (len(new_shape),)).reshape(new_shape)  # reduce original axes, then split
-
-  # *** movement ops ***
-
-  def _view(self, new_st:ShapeTracker) -> LazyBuffer:
-    if self.st.size == 0 or (new_st.views[-1].mask is not None and any((x[1]-x[0]) == 0 for x in new_st.views[-1].mask)):
-      return self.const_with_shape(0, new_st.shape)
-    if new_st.contiguous and self.base.shape == new_st.shape: return self.base
-    return create_lazybuffer(self.device, new_st, self.dtype, base=self.base)
-
-  def reshape(self, arg:Tuple[sint, ...]): return self._view(self.st.reshape(arg))
-  def pad(self, arg:Tuple[Tuple[sint, sint], ...]): return self._view(self.st.pad(arg))
-  def expand(self, arg:Tuple[sint, ...]): return self._view(self.st.expand(arg))
-  def permute(self, arg:Tuple[int, ...]): return self._view(self.st.permute(arg))
-  def shrink(self, arg:Tuple[Tuple[sint, sint], ...]): return self._view(self.st.shrink(arg))
-  def stride(self, arg:Tuple[int, ...]): return self._view(self.st.stride(arg))

tinygrad/function.py

@@ -1,212 +0,0 @@
-"""This is where the forwards and backwards passes live."""
-import math
-from typing import Tuple, Optional
-from tinygrad.helpers import argsort
-from tinygrad.dtype import dtypes, DType, sum_acc_dtype
-from tinygrad.ops import Ops, resolve, sint
-from tinygrad.tensor import Function
-from tinygrad.engine.lazy import LazyBuffer
-
-class Contiguous(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x.contiguous()
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output
-
-class ContiguousBackward(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.contiguous()
-
-class Cast(Function):
-  def forward(self, x:LazyBuffer, dtype:DType, bitcast:bool=False) -> LazyBuffer:
-    self.input_dtype, self.bitcast = x.dtype, bitcast
-    return x.bitcast(dtype) if self.bitcast else x.cast(dtype)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    if self.bitcast: raise RuntimeError("bitcast cannot backward")
-    return grad_output.cast(self.input_dtype)
-
-# ************* unary ops *************
-
-class Reciprocal(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.reciprocal()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return -grad_output * self.ret * self.ret
-
-class Sin(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return x.sin()
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return (math.pi/2 - self.x).sin() * grad_output
-
-class Relu(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.maximum(0)
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return self.ret.gt(0).cast(grad_output.dtype) * grad_output
-
-class Log(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return x.log2() * math.log(2)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output / self.x
-
-class Exp(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = (x * (1/math.log(2))).exp2()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return self.ret * grad_output
-
-class Sqrt(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = x.sqrt()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output / (self.ret*2)
-
-# NOTE: the implicit derivative of sigmoid is not stable
-# https://towardsdatascience.com/derivative-of-the-sigmoid-function-536880cf918e
-# TODO: have the backend automatically find this
-class Sigmoid(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer:
-    self.ret = (1 + (x * (-1/math.log(2))).exp2()).reciprocal()
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return (self.ret * (1 - self.ret)) * grad_output
-
-class Sign(Function):
-  def forward(self, x:LazyBuffer) -> LazyBuffer: return x.ne(0).where(x.lt(0).where(x.const_like(-1), x.const_like(1)), x.const_like(0))
-  # backward always return 0 to match torch
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.const_like(0)
-
-# ************* binary ops *************
-
-class Less(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.lt(y)
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None
-
-class Neq(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x.ne(y)
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]: return None, None
-
-class Xor(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x^y
-
-class BitwiseAnd(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x&y
-
-class BitwiseOr(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x|y
-
-class Threefry(Function):
-  def forward(self, x:LazyBuffer, seed:LazyBuffer) -> LazyBuffer: return x.threefry(seed)
-
-class Add(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x+y
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return grad_output if self.needs_input_grad[0] else None, \
-           grad_output if self.needs_input_grad[1] else None
-
-class Mul(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
-    self.x, self.y = x, y
-    return x * y
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return (self.y * grad_output) if self.needs_input_grad[0] else None, \
-           (self.x * grad_output) if self.needs_input_grad[1] else None
-
-class IDiv(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: return x // y
-
-# ************* ternary ops *************
-
-class Where(Function):
-  def forward(self, x:LazyBuffer, y:LazyBuffer, z:LazyBuffer) -> LazyBuffer:
-    self.x = x
-    return self.x.where(y, z)
-
-  def backward(self, grad_output:LazyBuffer) -> Tuple[None, Optional[LazyBuffer], Optional[LazyBuffer]]:
-    return None, \
-      self.x.where(grad_output, grad_output.const_like(0)) if self.needs_input_grad[1] else None, \
-      self.x.where(grad_output.const_like(0), grad_output) if self.needs_input_grad[2] else None
-
-# ************* reduce ops *************
-
-class Sum(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.input_shape = x.shape
-    return x.r(Ops.ADD, axis)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.expand(self.input_shape)
-
-class Prod(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.x, self.ret = x, x.r(Ops.MUL, axis)
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return (grad_output * self.ret).expand(self.x.shape) / self.x
-
-class Max(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.x, self.ret, self.axis = x, x.r(Ops.MAX, axis), axis
-    return self.ret
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    # 1s in locations where the max was chosen (can be two locations)
-    max_is_1s = self.x.ne(self.ret.expand(self.x.shape)).ne(self.x.const_like(1).cast(dtypes.bool)).cast(grad_output.dtype)
-    div = max_is_1s.r(Ops.ADD, self.axis).expand(self.x.shape)
-    return (max_is_1s/div) * grad_output.expand(self.x.shape)
-
-# ************* movement ops *************
-
-# NOTE: this is sum in reverse
-class Expand(Function):
-  def forward(self, x:LazyBuffer, shape:Tuple[int, ...]) -> LazyBuffer:
-    self.expanded_axis = tuple(i for i, (si, so) in enumerate(zip(x.shape, shape)) if resolve(si != so))
-    return x.expand(shape)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer:
-    return grad_output.cast(sum_acc_dtype(grad_output.dtype)).r(Ops.ADD, self.expanded_axis).cast(grad_output.dtype)
-
-class Reshape(Function):
-  def forward(self, x:LazyBuffer, shape:Tuple[int, ...]) -> LazyBuffer:
-    self.input_shape = x.shape
-    return x.reshape(shape)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.reshape(self.input_shape)
-
-class Permute(Function):
-  def forward(self, x:LazyBuffer, order:Tuple[int, ...]) -> LazyBuffer:
-    self.input_order = order
-    return x.permute(order)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.permute(argsort(self.input_order))
-
-class Pad(Function):
-  def forward(self, x:LazyBuffer, arg:Tuple[Tuple[int, int], ...]) -> LazyBuffer:
-    self.narg = tuple([(p[0], s+p[0]) for s,p in zip(x.shape, arg)])
-    return x.pad(arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.shrink(self.narg)
-
-class Shrink(Function):
-  def forward(self, x:LazyBuffer, arg:Tuple[Tuple[sint, sint], ...]) -> LazyBuffer:
-    self.narg = tuple([(p[0], s-p[1]) for s,p in zip(x.shape, arg)])
-    return x.shrink(arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.pad(self.narg)
-
-class Flip(Function):
-  def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
-    self.arg = tuple([-1 if i in axis else 1 for i in range(len(x.shape))])
-    return x.stride(self.arg)
-
-  def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.stride(self.arg)

tinygrad/multi.py

@@ -1,177 +0,0 @@
-from __future__ import annotations
-from typing import Optional, Tuple, List, Dict
-import functools, itertools, operator
-from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv
-from tinygrad.dtype import DType
-from tinygrad.ops import Ops, MathTrait
-from tinygrad.engine.lazy import LazyBuffer
-from tinygrad.shape.shapetracker import sint
-
-def all_reduce(bop: Ops, lbs: List[LazyBuffer]) -> List[LazyBuffer]:
-  assert all_int(lbs[0].shape), f"does not support symbolic shape {lbs[0].shape}"
-  assert all_same([lb.shape[0] for lb in lbs]), "allreduce with uneven shards is undefined"
-  n_lbs, shape, numel = len(lbs), lbs[0].shape, prod(lbs[0].shape)
-  # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
-  # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
-  use_ring = (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
-  if DEBUG >= 2: print(f"{'RING ALLREDUCE' if use_ring else 'NAIVE ALLREDUCE'} {n_lbs}x{numel} | {lbs[0].dtype}")
-  if not use_ring: return [functools.reduce(lambda x,y: x.alu(bop, y), [x.copy_to_device(lb.device) for x in lbs]) for lb in lbs]
-
-  factor = next(f for f in [32, 16, 8, 4, 2, 1] if numel % f == 0)
-  base, left = (numel // factor) // n_lbs, (numel // factor) % n_lbs
-  chunk_sizes = [(base + 1) * factor] * left + [base * factor] * (n_lbs - left)
-  acc = 0
-  chunks = [(acc, (acc := acc + i)) for i in chunk_sizes if i > 0]
-  chunked = [[lb.reshape((numel,)).shrink(((s,e),)) for s,e in chunks] for lb in lbs]
-
-  # scatter-reduce
-  for step in range(n_lbs-1):
-    for i in range(len(chunks)):
-      src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs
-      chunked[dest][i] = chunked[dest][i].alu(bop, chunked[src][i].copy_to_device(chunked[dest][i].device, force=True))
-
-  # allgather
-  for step in range(n_lbs-1):
-    for i in range(len(chunks)):
-      src, dest = (i+step-1)%n_lbs, (i+step)%n_lbs
-      chunked[dest][i] = chunked[src][i].copy_to_device(chunked[dest][i].device, force=True)
-
-  # assemble chunks back
-  pads = [((s,numel-e),) for s,e in chunks]
-  return [functools.reduce(operator.add, [c.pad(pad) for pad,c in zip(pads,lb_c)]).reshape(shape) for lb_c in chunked]
-
-def to_sharded(lbs:List[LazyBuffer], axis:int, bounds: Tuple[Tuple[int, int], ...]) -> List[LazyBuffer]:
-  if DEBUG >= 3 and lbs[0].shape[axis] % len(lbs) != 0: print(f"multi axis uneven: {lbs[0].shape=} {axis=} {len(lbs)=}, bounds={bounds}")
-  return [lb.shrink(tuple((0,s) if a != axis else bound for a,s in enumerate(lb.shape))) for i, (bound, lb) in enumerate(zip(bounds, lbs))]
-
-class MultiLazyBuffer(MathTrait):
-  def __init__(self, lbs:List[LazyBuffer], axis:Optional[int], real:Optional[List[bool]]=None):
-    assert all(isinstance(x, LazyBuffer) for x in lbs) and len(lbs), "all lbs must be LazyBuffers, and we need at least one of them"
-    assert all_same([x.dtype for x in lbs]), f"all multilazybuffer needs same dtype, getting {[x.dtype for x in lbs]}"
-    self.lbs, self.axis, self.dtype, self.device, self.real = lbs, axis, lbs[0].dtype, tuple(x.device for x in lbs), real or [True]*len(lbs)
-    if axis is not None:
-      splits = list(itertools.accumulate([lb.shape[axis] for lb in lbs], initial=0))
-      self.bounds = tuple(zip(splits, splits[1:]))
-
-  @property
-  def shape(self): return tuple(sum(y.shape[a] for y in self.real_lbs) if a == self.axis else s for a,s in enumerate(self.real_lbs[0].shape))
-
-  @property
-  def size(self): return sum(x.size for x in self.real_lbs)
-
-  @property
-  def real_lbs(self): return [lb for lb,r in zip(self.lbs, self.real) if r]
-
-  def __repr__(self): return f"<MLB {self.axis=} {self.real=} {chr(10)}{chr(10).join([f'{x.device} {x.st}' for x in self.lbs])}>"
-
-  @staticmethod
-  def from_sharded(lb:LazyBuffer, devices:Tuple[str, ...], axis:Optional[int], bounds:Optional[Tuple[Tuple[int, int], ...]]):
-    assert (axis is None) == (bounds is None), "must specify bounds iff axis is specified"
-    lbs = [lb] * len(devices)
-    sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(to_sharded(lbs, axis, bounds) if axis is not None and bounds is not None else lbs, devices)]
-    return MultiLazyBuffer([lb if lb.is_unrealized_unmasked_const() else lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis)
-
-  def copy_to_device(self, device:str) -> LazyBuffer:
-    if self.axis is None:
-      # if we already have a copy on the device, return that
-      return next((lb for lb in self.real_lbs if lb.device == device), self.real_lbs[0].copy_to_device(device))
-    # copy lbs to device, pad to final shape, and sum
-    llbs:List[LazyBuffer] = []
-    for lb,real,(start,end) in zip(self.lbs, self.real, self.bounds):
-      if not real: continue
-      pad_arg = tuple((0,0) if a != self.axis else (start, self.bounds[-1][1]-end) for a in range(len(lb.shape)))
-      llbs.append(lb.copy_to_device(device).pad(pad_arg))
-    return functools.reduce(operator.add, llbs)
-
-  # passthroughs
-  @property
-  def is_realized(self) -> bool: return all(lb.base.realized is not None for lb in self.real_lbs)
-  def cast(self, dtype:DType, bitcast:bool=False, allow_buffer_view=True):
-    return MultiLazyBuffer([x.cast(dtype, bitcast, allow_buffer_view) for x in self.lbs], self.axis, self.real)
-  def const_like(self, b) -> MultiLazyBuffer: return MultiLazyBuffer([x.const_like(b) for x in self.lbs], self.axis, self.real)
-  def assign(self, x:MultiLazyBuffer): return MultiLazyBuffer([s.assign(d) for s,d in zip(self.lbs, x.lbs)], self.axis, self.real)
-  def contiguous(self): return MultiLazyBuffer([x.contiguous() for x in self.lbs], self.axis, self.real)
-  def clone(self) -> MultiLazyBuffer: return MultiLazyBuffer([lb.clone() for lb in self.lbs], self.axis, self.real)
-
-  # elementwise is simple
-  def alu(self, op:Ops, *in_srcs:MultiLazyBuffer) -> MultiLazyBuffer:
-    msrcs = (self,)+in_srcs
-    assert all(isinstance(x, MultiLazyBuffer) for x in msrcs), f"all buffers must be MultiLazyBuffer {msrcs}"
-    assert all_same([x.device for x in msrcs]), f"all buffers must have the same device {[x.device for x in msrcs]}"
-
-    # NOTE: they all have to share an axis, we always choose [-1]
-    axis, bounds = axes[-1] if len(axes := dedup([(x.axis, x.bounds) for x in msrcs if x.axis is not None])) else (None, None)
-    srcs:List[List[LazyBuffer]] = []
-    not_all_real = not all(all(mlb.real) for mlb in msrcs)
-    new_real = [all(transposed) for transposed in zip(*[mlb.real for mlb in msrcs])] if not_all_real else self.real
-    assert any(new_real), "output contains no real lb"
-    for mlb in msrcs:
-      if (mlb.axis == axis and (mlb.axis is None or mlb.bounds == bounds)) or not_all_real: srcs.append(mlb.lbs)
-      elif mlb.axis is None and axis is not None: srcs.append(to_sharded(mlb.lbs, axis, bounds))
-      else: srcs.append(to_sharded([mlb.copy_to_device(lb.device) for lb in mlb.lbs], axis, bounds))
-    new_real_lbs:Dict[int,LazyBuffer] = {i:lsrcs[0].alu(op, *lsrcs[1:]) for i,(lsrcs,r) in enumerate(zip(zip(*srcs), new_real)) if r}
-    # NOTE: const dtype should match real
-    new_dtype = next(iter(new_real_lbs.values())).dtype
-    return MultiLazyBuffer([new_real_lbs.get(i, lsrcs[0].const_like(0).cast(new_dtype)) for i,lsrcs in enumerate(zip(*srcs))], axis, new_real)
-
-  def r(self, op:Ops, axis:Tuple[int, ...]) -> MultiLazyBuffer:
-    if self.axis is not None and self.axis in axis:
-      # all-reduce on sharded axes
-      reduced_parts = [(x if r else x.const_like(0)).r(op, axis) for x,r in zip(self.lbs, self.real)]
-      # if all partitions are real, do all_reduce
-      if all(self.real): return MultiLazyBuffer(all_reduce(op, reduced_parts), None)
-      # only one partition is real, keep it
-      return MultiLazyBuffer(reduced_parts, None, self.real)
-    # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct
-    return MultiLazyBuffer([x.r(op, axis) for x in self.lbs], self.axis, self.real)
-
-  # *** movement ops ***
-
-  def _shape_to_single_shard(self, shape:Tuple[sint, ...], lb:LazyBuffer) -> Tuple[sint, ...]:
-    return tuple(lb.shape[self.axis] if a == self.axis else s for a,s in enumerate(shape))
-
-  def reshape(self, arg:Tuple[sint, ...]):
-    if self.axis is None: return MultiLazyBuffer([x.reshape(arg) for x in self.lbs], None, self.real)
-    assert prod(self.shape) == prod(arg), "reshape must maintain prod(shape)"
-    arg_acc:List[sint] = list(itertools.accumulate(arg, operator.mul, initial=1))
-    # new_axis is the last one that preserves prod(prior to new_axis) and must not move items between shards
-    # todo: what to do about shrinking to self.shape[self.axis]==1 len(self.real_lbs)==1?
-    new_axis = len(arg_acc) - arg_acc[::-1].index(prod(self.shape[:self.axis])) - 1
-    assert all(prod(lb.shape[self.axis:])%prod(arg[new_axis+1:])==0 for lb in self.lbs), f"reshape cannot move items between shards {self=} {arg=}"
-    lbs = [x.reshape(tuple(s if a!=new_axis else prod(x.shape[self.axis:])//prod(arg[new_axis+1:]) for a,s in enumerate(arg))) for x in self.lbs]
-    return MultiLazyBuffer(lbs, new_axis, self.real)
-
-  def pad(self, arg:Tuple[Tuple[sint, sint], ...]):
-    assert self.axis is None or arg[self.axis] == (0,0) or not all(self.real), f"padding not supported for {arg=}"
-    # pad on shard axis -> fill others with zeros and set real to all True
-    if self.axis is not None and arg[self.axis] != (0,0):
-      # pad back to whole axis, remove real mask
-      assert all(arg[i] == (0, 0) for i in range(len(self.shape)) if i != self.axis), "cannot pad sharded and non-sharded axis at the same time"
-      dim, bound = sum(lb.shape[self.axis] for lb in self.lbs), self.bounds[self.real.index(True)]
-      assert arg[self.axis] == (bound[0], dim-bound[1]), "can only pad to whole axis"
-      return MultiLazyBuffer([x if r else x.const_like(0) for x,r in zip(self.lbs, self.real)], self.axis)
-    return MultiLazyBuffer([x.pad(arg) for x in self.lbs], self.axis, self.real)
-
-  def expand(self, arg:Tuple[sint, ...]):
-    # NOTE: this assert isn't needed, sharded axis can have dim 1
-    assert self.axis is None or arg[self.axis] == self.shape[self.axis], f"expand not supported on sharded axis {arg=}"
-    return MultiLazyBuffer([x.expand(self._shape_to_single_shard(arg, x)) for x in self.lbs], self.axis, self.real)
-
-  def permute(self, arg:Tuple[int, ...]):
-    # all permutes supported!
-    return MultiLazyBuffer([x.permute(arg) for x in self.lbs], arg.index(self.axis) if self.axis is not None else None, self.real)
-
-  def shrink(self, arg:Tuple[Tuple[sint, sint], ...]):
-    assert self.axis is None or arg[self.axis] == (0, self.shape[self.axis]) or arg[self.axis] in self.bounds, f"shrinking not supported for {arg=}"
-    if self.axis is not None and arg[self.axis] in self.bounds and arg[self.axis] != (0, self.shape[self.axis]):
-      assert all(arg[i] == (0, s) or i == self.axis for i,s in enumerate(self.shape)), "cannot shrink sharded and non-sharded axis at the same time"
-      # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real
-      idx = self.bounds.index(arg[self.axis])
-      # zero out other lbs to not create lb reference
-      return MultiLazyBuffer([lb if i==idx else lb.const_like(0) for i,lb in enumerate(self.lbs)], self.axis, [i==idx for i in range(len(self.lbs))])
-    return MultiLazyBuffer([x.shrink(tuple((0, x.shape[self.axis]) if a == self.axis else s for a,s in enumerate(arg))) for x in self.lbs],
-                           self.axis, self.real)
-
-  def stride(self, arg:Tuple[int, ...]):
-    assert self.axis is None or arg[self.axis] == 1, "flipping not supported on sharded axis"
-    return MultiLazyBuffer([x.stride(arg) for x in self.lbs], self.axis, self.real)

tinygrad/runtime/graph/clang.py

@@ -1,39 +0,0 @@
-from typing import List, Dict, cast
-import ctypes
-from tinygrad.helpers import dedup, cpu_time_execution, DEBUG
-from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.device import Buffer, Device
-from tinygrad.engine.realize import ExecItem, CompiledRunner
-from tinygrad.ops import Variable
-from tinygrad.runtime.ops_clang import ClangProgram
-from tinygrad.renderer.cstyle import ClangRenderer
-render_dtype = ClangRenderer().render_dtype
-
-class ClangGraph(GraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
-    super().__init__(jit_cache, input_rawbuffers, var_vals)
-    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
-
-    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
-    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v.expr}" for v in var_vals])
-    code = ["void batched("+','.join(args)+") {"]
-    for ji in jit_cache:
-      args = []
-      for buf in ji.bufs:
-        assert buf is not None
-        if buf in input_rawbuffers:
-          args.append(f"arg{input_rawbuffers.index(buf)}")
-        else:
-          args.append(f"({render_dtype(buf.dtype)}*)0x{ctypes.addressof(buf._buf):X}")
-      args += [x.expr for x in cast(CompiledRunner, ji.prg).p.vars]
-      code.append(f"  {cast(CompiledRunner, ji.prg).p.function_name}({','.join(args)});")
-    code.append("}")
-    if DEBUG >= 4: print("\n".join(code))
-    compiler = Device["CLANG"].compiler
-    assert compiler is not None
-    self.clprg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers
-
-  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
-    return cpu_time_execution(
-    lambda: self.clprg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), enable=wait)