""" Welcome to the tinygrad documentation ================= this file will take you on a whirlwind journey from a Tensor all the way down 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 for typing == from __future__ import annotations from typing import Optional, Tuple, Union, Any, Dict, Callable, Type, List, ClassVar from enum import Enum, auto from abc import ABC # %% # == Example: Tensor 2+3 == # let's trace an addition down through the layers of abstraction. # we will be using the clang backend from tinygrad import Device Device.DEFAULT = "CLANG" # first, 2+3 as a Tensor, the highest level 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.0 # %% # == Tensor (in tinygrad/tensor.py, code 8/10) == # it's worth reading tinygrad/tensor.py. it's pretty beautiful import tinygrad.mlops as mlops # this is the good old familiar Tensor class class Tensor: # 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 # high level ops (hlops) are defined on this class. example: relu def relu(self): return self.maximum(0) # 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 returned as Tensor properties device: str shape: Tuple[int, ...] dtype: DType # a ShapeTracker is used to track things like reshapes and permutes # all MovementOps are zero copy in tinygrad! # the ShapeTracker specifies how the data in the RawBuffer matches to the shape # we'll come back to this later st: ShapeTracker # if the LazyBuffer is realized, it has a Buffer # we will come back to Buffer later realized: Optional[Buffer] # if the lazybuffer is unrealized, it has a LazyOp # this LazyOp describes the computation needed to realize this LazyBuffer op: Optional[LazyOp] # LazyOp (in tinygrad/ops.py, code 4/10) # in a tree they form an Abstract Syntax Tree for a single GPU kernel class LazyOp: 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 26 Ops you have to implement for an accelerator. class UnaryOps(Enum): EXP2 = auto() LOG2 = auto() CAST = auto() SIN = auto() SQRT = auto() class BinaryOps(Enum): ADD = auto() SUB = auto() MUL = auto() DIV = auto() CMPLT = auto() MAX = auto() class ReduceOps(Enum): SUM = auto() MAX = auto() class MovementOps(Enum): RESHAPE = auto() PERMUTE = auto() EXPAND = auto() PAD = auto() SHRINK = auto() STRIDE = auto() class TernaryOps(Enum): MULACC = auto() WHERE = auto() class LoadOps(Enum): EMPTY = auto() CONST = auto() FROM = auto() CONTIGUOUS = auto() CUSTOM = auto() # NOTE: if you have a CompiledBuffer(DeviceBuffer) # you do not need to implement the MovementOps # as they are handled by the ShapeTracker(in tinygrad/shape/shapetracker.py, code 7/10) Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, TernaryOps, LoadOps] # 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 that is a "CPU" Tensor # again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first assert lazyop.src[0].op.op == LoadOps.FROM assert lazyop.src[0].op.src[0].device == "CPU" assert ( lazyop.src[0].op.src[0].op.src[0].realized._buf[0] == 2 ), "the src of the FROM LazyOP is a LazyBuffer on the CPU holding [2.]" assert result.lazydata.realized is None, "the LazyBuffer is not realized yet" # now we realize the LazyBuffer result.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 'RawMallocBuffer' 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" # %% # == Union[Interpreted, Compiled] (in tinygrad/ops.py, code 5/10) == # Now you have a choice, you can either write a "Interpreted" backend or "Compiled" backend # Interpreted backends are very simple (example: CPU and TORCH) class Interpreted: # and they have a lookup table to functions for the Ops fxn_for_op: Dict[Op, Callable] = { UnaryOps.EXP2: lambda x: np.exp2(x), BinaryOps.ADD: lambda x, y: x + y, } # Compiled backends take a little more (example: GPU and LLVM) class Compiled: # a code generator, which compiles the AST codegen: Type[Linearizer] # and a runtime, which runs the generated code runtime: Type[Runtime] # Runtime is what actually runs the kernels for a compiled backend class Runtime(ABC): # `name` is the name of the function, and `prg` is the code # the constructor compiles the code def __init__(self, name: str, prg: str): pass # call runs the code on the bufs. NOTE: the output is always bufs[0], but this is just a convention def __call__( self, *bufs: List[Buffer], global_size: Optional[List[int]], local_size: Optional[List[int]], ): pass # %% # == Buffer (in tinygrad/device.py, code 6/10) == import numpy as np # Buffer is where the data is actually held. it's pretty close to just memory class Buffer(ABC): # create an empty rawbuffer that holds `size` elements of type `dtype` # `opaque` is an opaque container class def __init__(self, device: str, size: int, dtype: DType, opaque: Any = None): pass # toCPU converts the RawBuffer to a numpy array with shape (size,) def toCPU(self) -> np.ndarray: pass # %% # == Example: 2+3 in raw clang == # MallocAllocator is the simplest concrete version of Allocator (in tinygrad/device.py) # it's used for the CLANG and LLVM backends # it's just malloc(size * dtype.itemsize) from tinygrad.device import MallocAllocator # 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, compile_clang # 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 input_a, input_b = MallocAllocator.alloc(4), MallocAllocator.alloc(4) output = MallocAllocator.alloc(4) # now we copy in the values numpy_a, numpy_b = np.array([2], dtype=np.float32), np.array([3], dtype=np.float32) MallocAllocator.copyin(input_a, numpy_a.data.cast("B")) MallocAllocator.copyin(input_b, numpy_b.data.cast("B")) # compile the program, run it, and 2+3 does indeed equal 5 program = ClangProgram( "add", compile_clang(f"void add(float *a, float *b, float *c) {{ *a = *b + *c; }}") ) program(output, input_a, input_b) numpy_out = np.empty(1, dtype=np.float32) MallocAllocator.copyout(numpy_out.data.cast("B"), output) assert numpy_out[0] == 5, "it's still 5" np.testing.assert_allclose(numpy_out, numpy_a + numpy_b) # %% # == Linearizer (in tinygrad/codegen/linearizer.py, code 4/10) == # in the above example, we wrote the code by hand # normally while using tinygrad you don't do that # the first step of transforming an AST into code is to "linearize" it, think like toposort on the AST # for that, we use the Linearizer, which turns an AST into a list of (linear) UOps class UOps(Enum): LOOP = auto() DEFINE_LOCAL = auto() LOAD = auto() ALU = auto() CONST = auto() ENDLOOP = auto() STORE = auto() class UOp: uop: UOps dtype: Optional[DType] vin: Tuple[UOp, ...] arg: Any num: int # UOps are unique class Linearizer: # 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 def linearize(self): pass # when linearize is run, it fills in this list uops: List[UOp] from tinygrad.tensor import Tensor result = Tensor(2).realize() + Tensor(3).realize() # use the real Linearizer to linearize 2+3 from tinygrad.codegen.linearizer import Linearizer sched = result.lazydata.schedule() linearizer = Linearizer(sched[-1].ast) linearizer.linearize() # print the uops for uop in linearizer.uops: print(uop) # output: """ 0 UOps.DEFINE_GLOBAL : ptr.dtypes.float [] ('data0', dtypes.float) 1 UOps.CONST : dtypes.float [] 2.0 2 UOps.CONST : dtypes.float [] 3.0 3 UOps.ALU : dtypes.float [1, 2] BinaryOps.ADD 4 UOps.CONST : dtypes.int [] 0 5 UOps.STORE : [0, 4, 3] None """ # %% # == Example: 2+3 autogenerated clang code == # to generate clang code, the Linearizer is wrapped with CStyleCodegen # here, we have an example where we fetch the generated code from the JIT 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.jit import CacheCollector CacheCollector.start() # enables the cache result.realize() # create the program and runs it cache_saved = CacheCollector.finish() # disable the cache # there's one ASTRunner in the cache assert len(cache_saved) == 1 # print the C Program :) print(cache_saved[0].prg.prg) # after some formatting (the compiler doesn't care) # NOTE: the 2 and 3 are constant folded """ void E_n2(float* restrict data0) { 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.from_shape((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 = 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 = 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 = 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 = 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 = a.simplify() print(a) # ShapeTracker(shape=(10, 10), views=[View((10, 10), (1, 10), 0)]) # and now we permute it back a = 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 6/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 # %%