tutorial

docs/abstractions.py

@@ -1,3 +1,17 @@
+"""
+Welcome to the tinygrad documentation
+=================
+
+this file will take you on a whirlwind journey from a Tensor to a Byte
+tinygrad has been aggressively refactored in the 2.5 years it's been worked on.
+what you see here is a refined library (with more refining to go still!)
+
+the whole tinygrad is ~2300 lines, so while it's readable in an evening or two,
+this documentation will help with entry points and understanding the abstraction stack
+"""
+
+# %%
+# == Boilerplate imports (typing mostly) ==
 from __future__ import annotations
 from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List
 from enum import Enum, auto
@@ -5,40 +19,65 @@ from abc import ABC
 import numpy as np
 import torch
 
-# tinygrad has grown a lot since any docs were last written.
-# It's now a terribly large 2300 lines!!
-# Let's trace an addition down through the layers of abstraction:
-# (note: this is a Python file for syntax highlighting)
+# %%
+# == Example: Tensor 2+3 ==
+# Let's trace an addition down through the layers of abstraction.
+# We will be using the clang backend
 
-# Some of this is documentation from the future, meaning this is how things will be refactored to be
-# Though this is pretty close, and is cleaner than what's actually there.
+from tinygrad.lazy import Device
+Device.DEFAULT = "CLANG"
+
+# first, 2+3 as a Tensor
+from tinygrad.tensor import Tensor
+a = Tensor([2])
+b = Tensor([3])
+result = a + b
+print(f"{a.numpy()} + {b.numpy()} = {result.numpy()}")
+assert result.numpy()[0] == 5.
+
+# %%
+# == Tensor (in tinygrad/tensor.py, code 8/10) ==
+# it's worth just reading tinygrad/tensor.py. it's pretty beautiful
+import tinygrad.mlops as mlops
 
-# Tensor (in tinygrad/tensor.py, code 8/10)
 # this is the good old familiar Tensor class
-class Function: pass
 class Tensor:
-  # these two are pretty striaghtforward
+  # these two are pretty straightforward
   grad: Optional[Tensor]
   requires_grad: Optional[bool]
+
   # this is the graph for the autograd engine
   _ctx: Optional[Function]
 
   # this is where the data (and other tensor properties) actually live
   lazydata: LazyBuffer
 
-# all the definitions of the derivatives are superclasses of Function
-# (in tinygrad/mlops.py, code 9/10)
-# they have forward and backward, and operate on LazyBuffers
-# Function.apply is responsible for lowering the Tensors to LazyBuffers and running forward
+  # high level ops (hlops) are defined here. example: relu
+  def relu(self): return self.maximum(0)
 
-# LazyBuffer (in tinygrad/lazy.py, code 5/10)
-# this is where the properties live that you thought were a part of Tensor
+  # log is an mlop, this is the wrapper function in Tensor
+  def log(self): return mlops.Log.apply(self)
+
+# all the definitions of the derivatives are subclasses of Function (like mlops.Log)
+# there's only 18 mlops for derivatives for everything (in tinygrad/mlops.py, code 9/10)
+# if you read one file, read mlops.py. if you read two files, also read tinygrad/tensor.py
+# you can differentiate the world using the chain rule
+class Function:
+  # example types of forward and backward
+  def forward(self, x:LazyBuffer) -> LazyBuffer: pass
+  def backward(self, x:LazyBuffer) -> LazyBuffer: pass
+
+# %%
+# == LazyBuffer (in tinygrad/lazy.py, code 5/10) ==
 from tinygrad.helpers import DType
+
+# this is where the properties live that you thought were a part of Tensor
+# LazyBuffer is like a Tensor without derivatives, at the mlop layer
 class LazyBuffer:
-  # these three define the "type" of the buffer, and they are proxied through Tensor
-  device:str
-  shape:Tuple[int, ...]
-  dtype:DType
+  # these three define the "type" of the buffer, and they are returned as Tensor properties
+  device: str
+  shape: Tuple[int, ...]
+  dtype: DType
 
   # if the lazybuffer is unrealized, it has a LazyOp
   # this LazyOp describes the computation needed to realize this LazyBuffer
@@ -49,12 +88,12 @@ class LazyBuffer:
   realized: Optional[DeviceBuffer]
 
 # LazyOp (in tinygrad/ops.py, code 4/10)
-# it's an AST node, that defines the type of the compute, the sources, and an optional static argument
-# they form an Abstract Syntax Tree
+# they form an Abstract Syntax Tree for a single GPU kernel
+# LazyOp is an AST node that defines:
 class LazyOp:
-  op: Op
-  src: Tuple[Union[LazyOp, LazyBuffer], ...]
-  arg: Optional[Any] = None
+  op: Op                                       # the type of the compute
+  src: Tuple[Union[LazyOp, LazyBuffer], ...]   # the sources
+  arg: Optional[Any] = None                    # and an optional static argument
 
 # there's currently 20 Ops you have to implement for an accelerator.
 class UnaryOps(Enum):    NOOP = auto(); EXP = auto(); LOG = auto(); NEG = auto(); NOT = auto()
@@ -67,10 +106,49 @@ class LoadOps(Enum):     FROMCPU = auto()
 #       as they are handled by the ShapeTracker(in tinygrad/shape/shapetracker.py, code 7/10)
 Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps]
 
-# which reminds me, we should get back to DeviceBuffer
-# it's an abstract class to be implemented for each backend
+# most of tinygrad/lazy.py is concerned with fusing Ops into LazyOps ASTs that map to GPUKernels
+# it's beyond the scope of this tutorial, but you can read the file if interested
+
+# %%
+# == Example: LazyBuffer for 2+3 ==
+
+from tinygrad.tensor import Tensor
+from tinygrad.ops import LazyOp, BinaryOps, LoadOps
+
+# the 2+3 from before
+result = Tensor([2]) + Tensor([3])
+print(type(result.lazydata), result.lazydata)  # let's look at the lazydata of result
+
+# you'll see it has a LazyOp
+# the op type is BinaryOps.ADD
+# and it has two sources, the 2 and the 3
+lazyop: LazyOp = result.lazydata.op
+assert lazyop.op == BinaryOps.ADD
+assert len(lazyop.src) == 2
+
+# the first source is the 2, it comes from the CPU
+# the source is a LazyBuffer, since FROMCPU cannot be folded into LazyOp ASTs
+# again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first
+print(lazyop.src[0].op)
+assert lazyop.src[0].op.op == LoadOps.FROMCPU
+assert lazyop.src[0].op.arg[0] == [2], "the arg of the FROMCPU LazyOP is the [2.]"
+assert result.lazydata.realized is None, "the LazyBuffer is not realized yet"
+
+# now we realize the LazyBuffer
+result.lazydata.realize()
+assert result.lazydata.realized is not None, "the LazyBuffer is realized!"
+# this brings us nicely to DeviceBuffer, of which the realized ClangBuffer is a subclass
+assert 'ClangBuffer' in str(type(result.lazydata.realized))
+# getting ahead of ourselves, but we can copy the DeviceBuffer toCPU
+assert result.lazydata.realized.toCPU()[0] == 5, "when put in numpy with toCPU, it's 5"
+
+# %%
+# == DeviceBuffer (in tinygrad/ops.py, code 4/10) ==
+
+# DeviceBuffer is an abstract class to be implemented for each Device backend
 class DeviceBuffer(ABC):
-  # these two are straightforward. no need for device, since that's contained in the concrete type
+  # these two are straightforward.
+  # unlike LazyBuffer, there's no need for device, since that's contained in the concrete type
   shape: Tuple[int, ...]
   dtype: DType
 
@@ -83,7 +161,9 @@ class DeviceBuffer(ABC):
   def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
 
 # DeviceBuffers come in two flavors, InterpretedBuffer and CompiledBuffer
-# InterpretedBuffers are a lot simpler than CompiledBuffers, and are used to implement the CPU(numpy) and TORCH backends
+# InterpretedBuffers are a lot simpler than CompiledBuffers
+# they are used to implement the CPU(numpy) and TORCH(torch) backends
+# it's worth reading CPUBuffer (in tinygrad/runtime/ops_cpu.py, code 8/10)
 class InterpretedBuffer(DeviceBuffer):
   # this is where the data actually lives
   # finally some classes you recognize!
@@ -94,10 +174,14 @@ class InterpretedBuffer(DeviceBuffer):
   fxn_for_op: Dict[Op, Callable] = {UnaryOps.EXP: lambda x: np.exp(x), BinaryOps.ADD: lambda x,y: x+y}
 
   # NOTE: exec_ast should not need to be overridden!
-  # The actual method lives in tinygrad/ops.py, and it walks the LazyOp tree and calls fxn_for_op as appropriate
+  # The actual method lives in tinygrad/ops.py
+  # it walks the LazyOp tree and calls fxn_for_op as appropriate
 
 # ********** NOTE: for the CPU and TORCH backends, we are done and you can stop reading here **********
 
+# %%
+# == CompiledBuffer (in tinygrad/ops.py, code 4/10) ==
+
 # however, all the magic of tinygrad will come from CompiledBuffer
 # this is used for the GPU(opencl), CUDA, METAL, CLANG, and LLVM backends
 class CompiledBuffer(DeviceBuffer):
@@ -128,10 +212,6 @@ class RawBuffer(ABC):
   # toCPU converts the RawBuffer to a numpy array with shape (size,). many backends are 0 copy here
   def toCPU(self) -> np.ndarray: raise NotImplementedError("must be implemented")
 
-# RawMallocBuffer is the simplest concrete version of this (in tinygrad/ops.py). it's used for the CLANG and LLVM backends
-# it's just malloc(size * dtype.itemsize)
-from tinygrad.ops import RawMallocBuffer
-
 # Runtime is what actually runs the kernels
 class Runtime(ABC):
   # `name` is the name of the function, and `prg` is the code
@@ -140,31 +220,44 @@ class Runtime(ABC):
   # call runs the code on the bufs. NOTE: the output is always bufs[0], but this is just a convention
   def __call__(self, global_size:Optional[List[int]], local_size:Optional[List[int]], *bufs:List[RawBuffer]): pass
 
-# ClangProgram is the simplest version (in tinygrad/runtime/ops_clang.py)
+# %%
+# == Example: 2+3 in raw clang ==
+
+# RawMallocBuffer is the simplest concrete version of RawBuffer (in tinygrad/ops.py)
+# it's used for the CLANG and LLVM backends
+# it's just malloc(size * dtype.itemsize)
+from tinygrad.ops import RawMallocBuffer
+
+# ClangProgram is the simplest runtime (in tinygrad/runtime/ops_clang.py, code 7/10)
 # __init__ calls clang, and __call__ calls the function in the *.so outputted by clang
 # in CLANG, global_size and local_size are ignored
 from tinygrad.runtime.ops_clang import ClangProgram
 
-# a concrete example looks like this, that adds two size 1 RawBuffer
-from tinygrad.helpers import dtypes
-
+# a concrete example looks like this, this adds two size 1 RawBuffer
 # first we create two numpy buffers containing 2 and 3
 # then we copy the numpy in to RawMallocBuffers
 # last, we create an empty output buffer
+from tinygrad.helpers import dtypes
 numpy_a, numpy_b = np.array([2], dtype=np.float32), np.array([3], dtype=np.float32)
 input_a, input_b = RawMallocBuffer.fromCPU(numpy_a), RawMallocBuffer.fromCPU(numpy_b)
 output = RawMallocBuffer(1, dtypes.float32)
 
 # compile the program, run it, and 2+3 does indeed equal 5
 program = ClangProgram("add", "void add(float *a, float *b, float *c) { *a = *b + *c; }")
-program(None, None, output, input_a, input_b)  # NOTE: the None are for global_size and local_size. clean this up?
+program(None, None, output, input_a, input_b)  # NOTE: the None are for global_size and local_size
+print(output.toCPU())
+assert output.toCPU()[0] == 5, "it's still 5"
 np.testing.assert_allclose(output.toCPU(), numpy_a+numpy_b)
 
+# %%
+# == ASTKernel (in tinygrad/codegen/ast.py, code 2/10) ==
+
 # but we are nowhere near done!
-# we need the LazyOp ASTs to actually be turned into code
+# we wrote the code above by hand
+# we need the LazyOp ASTs to be automatically turned into code
 # the current class looks roughly like this, but this will change and we will update the docs
 # this stuff is in the terrible 528 lines of (tinygrad/codegen/*, code 2/10 aka turd quality)
-class ASTKernel:  # (from tinygrad/codegen/ast.py)
+class ASTKernel:
   # create the kernel with the AST
   # NOTE: the AST contains the CompiledBuffers themselves as the root nodes. this will change
   def __init__(self, ast:LazyOp): pass
@@ -178,7 +271,118 @@ class ASTRunner:  # (from tinygrad/ops.py)
 
 # that hides a lot of complexity that will be refactored, but that's the basic idea of code generation
 
-# last, but not least (in fact one of the nicest things in tinygrad). the ShapeTracker
-class ShapeTracker: pass
+# %%
+# == Example: 2+3 autogenerated clang code ==
 
-# TODO: finish this, the coffee shop is closing
+from tinygrad.tensor import Tensor
+result = Tensor([2]) + Tensor([3])
+
+# we have a global cache used by the JIT
+# from there, we can see the generated clang code
+from tinygrad.ops import GlobalCounters
+GlobalCounters.cache = []    # enables the cache
+result.realize()             # create the program and runs it
+cache_saved = GlobalCounters.cache
+GlobalCounters.cache = None  # disable the cache
+
+# there's one ASTRunner in the cache
+assert len(cache_saved) == 1
+prg, bufs = cache_saved[0]
+
+# print the C Program :)
+print(prg.prg)
+
+# after some formatting (the compiler doesn't care)
+# NOTE: the 2 and 3 are constant folded
+"""
+void E_1(float* data0) {
+  for (int idx0 = 0; idx0 < 1; idx0++) {
+    data0[0] = (2.0f) + (3.0f);
+  }
+}
+"""
+
+# %%
+# == Example: ShapeTracker (in tinygrad/shape/shapetracker.py, code 7/10) ==
+
+# remember how I said you don't have to write the MovementOps for CompiledBuffers?
+# that's all thanks to ShapeTracker!
+# ShapeTracker tracks the indices into the RawBuffer
+from tinygrad.shape.shapetracker import ShapeTracker
+
+# create a virtual (10, 10) Tensor. this is just a shape, there's no actual tensor
+a = ShapeTracker((10, 10))
+
+# you'll see it has one view. the (10, 1 are the strides)
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
+
+# we can permute it, and the strides change
+a.permute((1,0))
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
+
+# we can then reshape it, and the strides change again
+# note how the permute stays applied
+a.reshape((5,2,5,2))
+print(a) # ShapeTracker(shape=(5, 2, 5, 2), views=[View((5, 2, 5, 2), (2, 1, 20, 10), 0)])
+
+# now, if we were to reshape it to a (100,) shape tensor, we have to create a second view
+a.reshape((100,))
+print(a) # ShapeTracker(shape=(100,), views=[
+         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
+         #   View((100,), (1,), 0)])
+
+# Views stack on top of each other, to allow zero copy for any number of MovementOps
+# we can render a Python expression for the index at any time
+idx, _ = a.expr_idxs()
+print(idx.render())  # (((idx0%10)*10)+(idx0//10))
+
+# of course, if we reshape it back, the indexes get simple again
+a.reshape((10,10))
+idx, _ = a.expr_idxs()
+print(idx.render())  # ((idx1*10)+idx0)
+
+# the ShapeTracker still has two views though...
+print(a) # ShapeTracker(shape=(10, 10), views=[
+         #   View((5, 2, 5, 2), (2, 1, 20, 10), 0),
+         #   View((10, 10), (10, 1), 0)])
+
+# ...until we simplify it!
+a.simplify()
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)])
+
+# and now we permute it back
+a.permute((1,0))
+print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (10, 1), 0)])
+
+# and it's even contiguous
+assert a.contiguous == True
+
+# %%
+# == Example: Variable (in tinygrad/shape/symbolic.py, code 7/10) ==
+
+# Under the hood, ShapeTracker is powered by a small symbolic algebra library
+from tinygrad.shape.symbolic import Variable
+
+# Variable is the basic class from symbolic
+# it's created with a name and a min and max (inclusive)
+a = Variable("a", 0, 10)
+b = Variable("b", 0, 10)
+
+# some math examples
+print((a*10).min, (a*10).max)  # you'll see a*10 has a min of 0 and max of 100
+print((a+b).min, (a+b).max)    # 0 20, you get the idea
+
+# but complex expressions are where it gets fun
+expr = (a + b*10) % 10
+print(expr.render())   # (a%10)
+# as you can see, b is gone!
+
+# one more
+expr = (a*40 + b) // 20
+print(expr.render())       # (a*2)
+print(expr.min, expr.max)  # 0 20
+# this is just "(a*2)"
+# since b only has a range from 0-10, it can't affect the output
+
+
+# %%

tinygrad/codegen/ast.py

@@ -170,7 +170,7 @@ class ASTKernel:
         if mergeable: rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i])
         else: rets[j].append((shapes[j][i], strides[j][i]))
 
-    for i,x in enumerate(rets): self.sts[i]._reshape(tuple(y[0] for y in x))
+    for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x))
     self.first_reduce = get_first_reduce([x.shape for x in self.sts])
 
   # this should be aware of the three parts to the shape
@@ -179,8 +179,8 @@ class ASTKernel:
   #  * the size outputted by each kernel
   def reshape_and_permute(self, new_shape_fxn, axis):
     for st in self.sts:
-      if new_shape_fxn is not None: st._reshape(tuple(new_shape_fxn(st.shape)))
-      if axis is not None: st._permute(tuple(axis))
+      if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape)))
+      if axis is not None: st.permute(tuple(axis))
 
   # axis : the axis to pull from
   # amount : the amount to take

tinygrad/shape/shapetracker.py

@@ -149,7 +149,7 @@ class ShapeTracker:
     offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)])
     self.views[-1] = View(tuple(y-x for x,y in arg), self.strides, self.offset+offset)
 
-  def _pad(self, arg: Tuple[Tuple[int, int], ...]):
+  def pad(self, arg: Tuple[Tuple[int, int], ...]):
     assert all((b>=0 and e>=0) for b,e in arg) and len(arg) == len(self.shape)
     if all(b==0 and e==0 for b,e in arg): return self   # ZeroView is expensive if we don't need it
     zvarg = tuple((-b,s+e) for s,(b,e) in zip(self.shape, arg))
@@ -158,15 +158,15 @@ class ShapeTracker:
     # if we add a ZeroView, we add another (stock) view also for modding
     self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))]
 
-  def _shrink(self, arg: Tuple[Tuple[int, int], ...]):
+  def shrink(self, arg: Tuple[Tuple[int, int], ...]):
     assert all((b>=0 and e<=s) for s,(b,e) in zip(self.shape,arg)) and len(arg) == len(self.shape)
     self.__unsafe_resize(arg)
 
-  def _expand(self, new_shape: Tuple[int, ...]):
+  def expand(self, new_shape: Tuple[int, ...]):
     assert all(isinstance(x, int) and (s == x or (s == 1 and st == 0)) for s,x,st in zip(self.shape, new_shape, self.strides)), f"can't expand {self.shape} into {new_shape}"
     self.views[-1] = View(new_shape, self.strides, self.offset)
 
-  def _reshape(self, new_shape: Tuple[int, ...]):
+  def reshape(self, new_shape: Tuple[int, ...]):
     if self.shape == new_shape: return self
     assert all(isinstance(x, int) and x != 0 for x in new_shape), f"shape must be ints and can't contain 0 {new_shape}"
     assert prod(self.shape) == prod(new_shape), f"can't reshape {self.shape} -> {new_shape}"
@@ -186,13 +186,13 @@ class ShapeTracker:
       if (merged_view := merge_views(cast(View, self.views[-1]), view)) is not None: self.views[-1] = merged_view
       else: self.views.append(view)
 
-  def _permute(self, axis: Tuple[int, ...]):
+  def permute(self, axis: Tuple[int, ...]):
     assert all(isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis), f"invalid permute {axis} for {self.shape}"
     assert len(set(axis)) == len(axis) and len(axis) == len(self.shape), f"can't permute {self.shape} with {axis}"
     self.views[-1] = View(tuple(self.shape[a] for a in axis), tuple(self.strides[a] for a in axis), self.offset)
 
   # except for the negative case, you can build this from the others. invertible in the negative case
-  def _stride(self, mul: Tuple[int, ...]):
+  def stride(self, mul: Tuple[int, ...]):
     assert all(isinstance(x, int) and x != 0 for x in mul), f"invalid stride {mul} for {self.shape}"
     strides = tuple(z*m for z,m in zip(self.strides, mul))
     new_shape = tuple((s+(abs(m)-1))//abs(m) for s,m in zip(self.shape, mul))
@@ -206,8 +206,8 @@ class ShapeTracker:
     dispatch[op](self, arg)
     return self
 
-dispatch: Dict[MovementOps, Callable] = {MovementOps.RESHAPE: ShapeTracker._reshape, MovementOps.EXPAND: ShapeTracker._expand, MovementOps.PAD: ShapeTracker._pad,
-                                         MovementOps.SHRINK: ShapeTracker._shrink, MovementOps.PERMUTE: ShapeTracker._permute, MovementOps.STRIDE: ShapeTracker._stride}
+dispatch: Dict[MovementOps, Callable] = {MovementOps.RESHAPE: ShapeTracker.reshape, MovementOps.EXPAND: ShapeTracker.expand, MovementOps.PAD: ShapeTracker.pad,
+                                         MovementOps.SHRINK: ShapeTracker.shrink, MovementOps.PERMUTE: ShapeTracker.permute, MovementOps.STRIDE: ShapeTracker.stride}
 
 # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
 def get_contraction(old_shape:Tuple[int, ...], new_shape:Tuple[int, ...]) -> Optional[List[List[int]]]: