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

tinygrad/runtime/support/llvm.py

@@ -0,0 +1,32 @@
+import ctypes, ctypes.util, os, sys, subprocess
+from tinygrad.helpers import DEBUG, OSX, getenv
+
+if sys.platform == 'win32':
+  # Windows llvm distribution doesn't seem to add itself to PATH or anywhere else where it can be easily retrieved from.
+  # winget also doesn't have something like `brew --prefix llvm` so just hardcode default installation path with an option to override
+  LLVM_PATH = getenv('LLVM_PATH', 'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll')
+  if not os.path.exists(LLVM_PATH):
+    raise RuntimeError('LLVM not found, you can install it with `winget install LLVM.LLVM` or point at a custom dll with LLVM_PATH')
+elif OSX and 'tinygrad.runtime.ops_metal' in sys.modules:
+  # Opening METAL after LLVM doesn't fail because ctypes.CDLL opens with RTLD_LOCAL but MTLCompiler opens it's own llvm with RTLD_GLOBAL
+  # This means that MTLCompiler's llvm will create it's own instances of global state because RTLD_LOCAL doesn't export symbols, but if RTLD_GLOBAL
+  # library is loaded first then RTLD_LOCAL library will just use it's symbols. On linux there is RTLD_DEEPBIND to prevent that, but on macos there
+  # doesn't seem to be anything we can do.
+  LLVM_PATH = ctypes.util.find_library('tinyllvm')
+  if LLVM_PATH is None:
+    raise RuntimeError("LLVM can't be opened in the same process with metal. You can install llvm distribution which supports that via `brew install uuuvn/tinygrad/tinyllvm`") # noqa: E501
+elif OSX:
+  brew_prefix = subprocess.check_output(['brew', '--prefix', 'llvm']).decode().strip()
+  # `brew --prefix` will return even if formula is not installed
+  if not os.path.exists(brew_prefix):
+    raise RuntimeError('LLVM not found, you can install it with `brew install llvm`')
+  LLVM_PATH = os.path.join(brew_prefix, 'lib', 'libLLVM.dylib')
+else:
+  LLVM_PATH = ctypes.util.find_library('LLVM')
+  for ver in range(14, 19+1):
+    if LLVM_PATH is not None: break
+    LLVM_PATH = ctypes.util.find_library(f'LLVM-{ver}')
+  if LLVM_PATH is None:
+    raise RuntimeError("No LLVM library found on the system. Install it via your distro's package manager and ensure it's findable as 'LLVM'")
+
+if DEBUG>=2: print(f'Using LLVM at {repr(LLVM_PATH)}')

tinygrad/shape/shapetracker.py

@@ -1,30 +1,79 @@
 # ShapeTracker allows movement operations to a buffer that don't require a copy to be made.
 from __future__ import annotations
 from dataclasses import dataclass
-from typing import Tuple, List, Optional, Dict, Set
+import functools
+from typing import Optional, Callable
 from tinygrad.helpers import merge_dicts, getenv
-from tinygrad.shape.view import View, strides_for_shape
+from tinygrad.shape.view import View, strides_for_shape, unravel
 from tinygrad.dtype import dtypes
-from tinygrad.ops import UOp, Ops, graph_rewrite, split_uop, symbolic_flat, Variable, sint, uop_given_valid, simplify_valid
+from tinygrad.ops import UOp, Ops, graph_rewrite, split_uop, symbolic_flat, Variable, sint, uop_given_valid, simplify_valid, sint_to_uop, Context
+from tinygrad.codegen.rewriter import sym
+
+def overflow(u: UOp): return u.vmax > dtypes.max(dtypes.int) or u.vmin < dtypes.min(dtypes.int)
+
+# If a node overflow, its srcs need to be checked to see if this overflow is the result of an ALU operation,
+# or that the node simply inherits the dtype from srcs. Upcast is either `Ops.CAST`+`replace` or just `replace`.
+def upcast(u: UOp):
+  srcs = tuple(upcast(_src) for _src in u.src)
+  if u.dtype.scalar() is dtypes.int:
+    dtype = dtypes.int64.vec(u.dtype.count) if u.dtype.count > 1 else dtypes.int64
+    upcasted = u.replace(dtype=dtype, src=tuple([_src.cast(dtype) for _src in srcs]))
+    if overflow(u): return upcasted
+    # Check the original src, new srcs has Ops.CAST whose vmin, vmax change the real bounds
+    # Cast back is required because if the node is in range, siblings would never be upcasted
+    if any((overflow(src) for src in u.src)): return upcasted.cast(u.dtype)
+  return u.replace(src=tuple(srcs))
+
+# pooling op may overflow before folding causing unnecessary upcast
+def folded_upcast(u: UOp):
+  with Context(TRACK_MATCH_STATS=0):
+    return upcast(graph_rewrite(u, sym, {}))
+
+@functools.lru_cache(None)
+def views_to_indexed_uops(views: tuple[View, ...], _idxs:Optional[tuple[UOp, ...]]=None) -> tuple[UOp, UOp]:
+  idx, valid = views[-1].to_indexed_uops(_idxs)
+  for view in reversed(views[0:-1]):
+    view = view.minify()
+    idx, valid = view.to_indexed_uops([sint_to_uop(i) for i in unravel(view.shape, idx)], valid)
+  return idx, valid
+
+@functools.lru_cache(None)
+def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[Optional[sint], ...]:
+  # NOTE: if a stride is not always valid, it will be None
+  if len(views) == 1 and views[-1].mask is None: return views[-1].strides
+  ret: list[Optional[sint]] = [None] * len(views[-1].shape)
+  idx, valid = (graph_rewrite(u, symbolic_flat) for u in views_to_indexed_uops(views))
+  # TODO: always apply these in to_indexed_uops?
+  if (newvalid:=simplify_valid(valid)) is not None: valid = newvalid
+  if (newidx:=uop_given_valid(valid, idx)) is not None: idx = graph_rewrite(newidx, symbolic_flat)
+  for c in split_uop(idx, Ops.ADD):
+    if c.op is Ops.RANGE: ret[c.arg] = 1
+    if c.op is Ops.MUL and c.src[0].op is Ops.RANGE and c.src[1].op is Ops.CONST: ret[c.src[0].arg] = c.src[1].arg
+    if c.op is Ops.MUL and c.src[1].op is Ops.RANGE and c.src[0].op is Ops.CONST: ret[c.src[1].arg] = c.src[0].arg
+  used_ranges = [x.arg for x in idx.toposort if x.op is Ops.RANGE]
+  ret = [x if i in used_ranges else 0 for i,x in enumerate(ret)]
+  if not ignore_valid:
+    for masked_axis in [x.arg for x in valid.toposort if x.op is Ops.RANGE]: ret[masked_axis] = None
+  return tuple(ret)
 
 @dataclass(frozen=True, order=True)
 class ShapeTracker:
-  views: Tuple[View, ...]
+  views: tuple[View, ...]
 
   def __add__(self, st:ShapeTracker) -> ShapeTracker:
     ret = self
     for v in st.views: ret = ShapeTracker(ret.views + (v,)).simplify() # one view at a time = better simplification
     return ret
 
-  def invert(self, out_shape:Tuple[sint, ...]) -> Optional[ShapeTracker]:
-    inverted_views:List[View] = []
+  def invert(self, out_shape:tuple[sint, ...]) -> Optional[ShapeTracker]:
+    inverted_views:list[View] = []
     for v,s in zip(self.views[::-1], [x.shape for x in self.views[::-1][1:]]+[out_shape]):
       if (inverted:= v.invert(s)) is None: return None
       inverted_views.append(inverted)
     return ShapeTracker(tuple(inverted_views)).reshape(out_shape)
 
   @staticmethod
-  def from_shape(shape:Tuple[sint, ...]) -> ShapeTracker: return ShapeTracker((View.create(shape),))
+  def from_shape(shape:tuple[sint, ...]) -> ShapeTracker: return ShapeTracker((View.create(shape),))
 
   @property
   def contiguous(self) -> bool: return len(self.views) == 1 and self.views[0].contiguous
@@ -33,65 +82,42 @@ class ShapeTracker:
   def consecutive(self) -> bool: return len(self.views) == 1 and (v:=self.views[0]).mask is None and v.strides == strides_for_shape(v.shape)
 
   @property
-  def shape(self) -> Tuple[sint, ...]: return self.views[-1].shape
+  def shape(self) -> tuple[sint, ...]: return self.views[-1].shape
 
   @property
   def size(self) -> int: return self.views[-1].size()
 
-  def reduce(self, axis:Tuple[int, ...]) -> Tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape))
+  def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape))
 
   def to_uop(self) -> UOp: return UOp(Ops.VIEW, dtypes.void, (), self)
+  def to_indexed_uops(self, _idxs:Optional[list[UOp]|tuple[UOp, ...]]=None) -> tuple[UOp, UOp]:
+    idx, valid = views_to_indexed_uops(self.views, tuple(_idxs) if _idxs is not None else None)
+    return folded_upcast(idx), folded_upcast(valid)
 
-  def to_indexed_uops(self, _idxs:Optional[List[UOp]]=None) -> Tuple[UOp, UOp]:
-    idx, valid = self.views[-1].to_indexed_uops(_idxs)
-    for view in reversed(self.views[0:-1]):
-      view = view.minify()
-      acc, idxs = 1, []
-      for d in reversed(view.shape):
-        idxs.append((idx//acc)%d)
-        acc *= d
-      idx, valid = view.to_indexed_uops(idxs[::-1], valid)
-    return idx, valid
-
+  # upper bound on buffer size required to fit this shapetracker
   def real_size(self) -> int:
     if 0 in self.shape: return 0
-    idx, valid = self.to_indexed_uops()
-    if not valid.vmax: return 0
+    view = (v.shrink(v.mask) if (v:=self.views[0]).mask else v)
+    idx, _ = views_to_indexed_uops((view,))
     assert idx.vmax < 1e12, f"real_size broken for {self}"
-    return int(idx.vmax+1)
+    return int(idx.vmax + 1)
 
-  def vars(self) -> Set[Variable]: return set().union(*[v.vars() for v in self.views])
+  def vars(self) -> set[Variable]: return set().union(*[v.vars() for v in self.views])
 
   @property
-  def var_vals(self) -> Dict[Variable, int]: return merge_dicts([dict([v.unbind()]) for v in self.vars()])
+  def var_vals(self) -> dict[Variable, int]: return merge_dicts([dict([v.unbind()]) for v in self.vars()])
 
-  def unbind(self) -> Tuple[ShapeTracker, Dict[Variable, int]]:
+  def unbind(self) -> tuple[ShapeTracker, dict[Variable, int]]:
     unbound_views, var_vals = zip(*[v.unbind() for v in self.views])
     return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals)
 
-  # NOTE: if a stride is not always valid, it will be None
-  def real_strides(self, ignore_valid=False) -> Tuple[Optional[sint], ...]:
-    if len(self.views) == 1 and self.views[-1].mask is None: return self.views[-1].strides
-    ret: List[Optional[sint]] = [None] * len(self.shape)
-    idx, valid = (graph_rewrite(u, symbolic_flat) for u in self.to_indexed_uops())
-    # TODO: always apply these in to_indexed_uops?
-    if (newvalid:=simplify_valid(valid)) is not None: valid = newvalid
-    if (newidx:=uop_given_valid(valid, idx)) is not None: idx = graph_rewrite(newidx, symbolic_flat)
-    for c in split_uop(idx, Ops.ADD):
-      if c.op is Ops.RANGE: ret[c.arg[0]] = 1
-      if c.op is Ops.MUL and c.src[0].op is Ops.RANGE and c.src[1].op is Ops.CONST: ret[c.src[0].arg[0]] = c.src[1].arg
-      if c.op is Ops.MUL and c.src[1].op is Ops.RANGE and c.src[0].op is Ops.CONST: ret[c.src[1].arg[0]] = c.src[0].arg
-    used_ranges = [x.arg[0] for x in idx.sparents if x.op is Ops.RANGE]
-    ret = [x if i in used_ranges else 0 for i,x in enumerate(ret)]
-    if not ignore_valid:
-      for masked_axis in [x.arg[0] for x in valid.sparents if x.op is Ops.RANGE]: ret[masked_axis] = None
-    return tuple(ret)
-
-  def unit_stride_axes(self, ignore_valid=False) -> List[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1]
+  def real_strides(self, ignore_valid=False) -> tuple[Optional[sint], ...]: return views_to_real_strides(self.views, ignore_valid)
+  def unit_stride_axes(self, ignore_valid=False) -> list[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1]
 
   def axis_is_masked(self, axis:int) -> bool:
-    _, valid = self.to_indexed_uops()
-    return axis in [x.arg[0] for x in graph_rewrite(valid, symbolic_flat).sparents if x.op is Ops.RANGE]
+    with Context(TRACK_MATCH_STATS=0):
+      _, valid = self.to_indexed_uops()
+      return axis in [x.arg for x in graph_rewrite(valid, symbolic_flat).toposort if x.op is Ops.RANGE]
 
   def simplify(self) -> ShapeTracker:
     if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:
@@ -100,12 +126,17 @@ class ShapeTracker:
 
   # *** under this line are the movement ops ***
 
-  def pad(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ))
-  def shrink(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ))
-  def expand(self, new_shape: Tuple[sint, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ))
-  def permute(self, axis: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ))
-  def stride(self, mul: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].stride(mul), ))
+  def pad(self, arg: tuple[tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ))
+  def shrink(self, arg: tuple[tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ))
+  def expand(self, new_shape: tuple[sint, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ))
+  def permute(self, axis: tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ))
+  def flip(self, mul: tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].flip(mul), ))
 
-  def reshape(self, new_shape: Tuple[sint, ...]) -> ShapeTracker:
+  def reshape(self, new_shape: tuple[sint, ...]) -> ShapeTracker:
     if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None: return ShapeTracker(self.views[0:-1] + (new_view,))
     return ShapeTracker(self.views + (View.create(new_shape), ))
+
+  def mop(self, op, arg): return mops[op](self, arg)
+
+mops: dict[Ops, Callable] = {Ops.RESHAPE: ShapeTracker.reshape, Ops.PERMUTE: ShapeTracker.permute, Ops.EXPAND: ShapeTracker.expand,
+                             Ops.SHRINK: ShapeTracker.shrink, Ops.FLIP: ShapeTracker.flip, Ops.PAD: ShapeTracker.pad}