From 153cce0f7e20aba047c7c42fc0b768cc9a2978de Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 12 Mar 2023 22:31:46 -0700 Subject: [PATCH] tutorial --- docs/abstractions.py | 290 ++++++++++++++++++++++++++++----- tinygrad/codegen/ast.py | 6 +- tinygrad/shape/shapetracker.py | 16 +- 3 files changed, 258 insertions(+), 54 deletions(-) diff --git a/docs/abstractions.py b/docs/abstractions.py index 9aaebcb541..2051e5f2e0 100644 --- a/docs/abstractions.py +++ b/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 + + +# %% diff --git a/tinygrad/codegen/ast.py b/tinygrad/codegen/ast.py index 0aad89921b..8902b28ab0 100644 --- a/tinygrad/codegen/ast.py +++ b/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 diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 67623817d5..daa369ca8b 100644 --- a/tinygrad/shape/shapetracker.py +++ b/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]]]: