everything is a UOp - Senthilnathan

i&rsquo;ve used PyTorch for years. a @ b, .backward(), .cuda(). works every time. i never thought about what&rsquo;s underneath.<br>then i came across tinygrad. it&rsquo;s the deep learning stack behind comma.ai&rsquo;s openpilot, the open-source self-driving system. 17K lines of Python. PyTorch is 3 million lines of C++. this thing fits the whole pipeline, including the compiler, in less code than some test suites.<br>i figured i&rsquo;d spend an afternoon reading it.<br>to give an outline of how tinygrad works, we would explore these 4 stages:

the first thing i noticed<br>to start with, we need to know how a tensor operation works. the fundamental tensor operation used in DL libraries is matrix multiplication . so, i opened a terminal and typed this out.<br>from tinygrad import Tensor

a = Tensor.rand(4, 4)<br>b = Tensor.rand(4, 4)<br>c = a.matmul(b)

print(c.shape) # (4, 4)

this creates two tensors of shape 4 x 4 with random values and multiplies it, store its result in c. the shape came back as (4, 4). looks fine!!<br>but&mldr; no multiplication had actually happened. a.matmul(b) didn&rsquo;t crunch any numbers. it just built a tree of lazy operations and returned. the real work only runs when you call c.realize().<br>this was new to me because, PyTorch does a @ b and immediately launches the matmul. but, tinygrad waits.<br>this is called lazy evaluation . you describe what you want to do, and the framework postpones the actual computation until you explicitly ask for the answer. function calls build a graph. the graph just sits there waiting for the computation caller to call, which is what .realize() is, and the moment where everything actually runs.<br>so, i was curious what the graph looked like.<br>REDUCE(sum, axis=-1)<br>PERMUTE<br>MUL<br>RESHAPE(4,1,4)<br>a.uop<br>PERMUTE<br>RESHAPE(1,4,4)<br>b.uop

i stared at it for a minute. the tree reads bottom-up :)<br>take a, view it as a 4-by-1-by-4. take b, view it as 1-by-4-by-4, then transpose the last two axes. multiply them elementwise with broadcasting, which gives you a 4-by-4-by-4. then sum along the last axis. that&rsquo;s a 4-by-4 matrix multiplication.

a @ b got turned into (a.reshape(4,1,4) * b.reshape(1,4,4).transpose()).sum(-1).<br>the matmul isn&rsquo;t built into the framework as a primitive. it&rsquo;s syntactic sugar. a convenient way to write reshape, broadcast multiply, and sum. relu is sugar for max(x, 0). sigmoid is sugar for 1 / (1 + exp(-x)). the framework has maybe six actual operations. everything else is convenience that decomposes before the scheduler ever sees it.<br>so if everything is shorthand, and shorthand doesn&rsquo;t trigger any computation&mldr; the framework can see your whole program at once before deciding what to run.<br>what&rsquo;s inside a tensor<br>next i opened tensor.py. i was expecting a struct with strides and a device pointer and maybe a reference count. something that looked like it owned data.<br>__slots__ = "uop", "is_param", "grad"

that&rsquo;s it. a reference to a graph node. a flag for the optimizer. a slot for the gradient after .backward().<br>shape lives on the graph node. self.uop.shape. dtype, device, everything. the tensor object doesn&rsquo;t carry any of that itself.<br>the graph node is called a UOp . it&rsquo;s the central thing in the whole codebase. tensors are UOps. kernels are UOps. compiled binaries are UOps. a UOp has exactly four fields.<br>op says what kind of node. dtype is float32 or int64 or whatever. src is a tuple of child nodes feeding into this one. arg is extra data the operation needs, like the target shape for a reshape.<br>when i realised everything is a UOp i went looking for where the actual operations are defined. add, mul, matmul. and found something that didn&rsquo;t make sense at first.<br>>>> Tensor.add is UOp.add<br>True

the wrapper and the thing it wraps share the same method. not the same name. the same python function object.<br>it works like this. both Tensor and UOp inherit from shared helper classes called mixins . the operations are written once in files like mixin/elementwise.py. both classes pick them up through inheritance.<br>the mixin methods call a few abstract hooks at the bottom, and each class fills in the hooks differently. Tensor&rsquo;s version wraps results back into a new Tensor. UOp&rsquo;s version creates a new graph node. the code in the middle is identical.<br>what this means in practice is you can write at the Tensor level for normal stuff, or drop into raw UOp graphs when you need control over the kernel. same methods either way.<br>the graph doesn&rsquo;t repeat itself<br>UOp nodes get cached. every combination of (op, dtype, children, arg) is stored once in a dictionary. if you build the same expression twice, you get back the same object.<br>>>> Tensor(3) + 4<br>>>> Tensor(3) + 4<br># identical UOp object, served from cache

this is called, hash-consing . it means graph rewriting is just pointer replacement . keep all the shared structure, swap out the one subtree that changed, that&rsquo;s it.<br>a pattern...