{"slug": "everything-is-a-uop", "title": "everything is a UOp", "summary": "Tinygrad, the 17,000-line Python deep learning framework behind comma.ai's openpilot self-driving system, represents all tensor operations as UOps (universal operations) and uses lazy evaluation, deferring computation until .realize() is called, in contrast to PyTorch's 3 million lines of C++ and immediate execution. The framework decomposes high-level operations like matmul into a small set of primitives (reshape, broadcast multiply, sum), and its Tensor object contains only a UOp reference, a parameter flag, and a gradient slot, with all metadata stored on the graph node.", "body_md": "i’ve used PyTorch for years. `a @ b`\n\n, `.backward()`\n\n, `.cuda()`\n\n. works every time. i never thought about what’s underneath.\n\nthen i came across tinygrad. it’s the deep learning stack behind [comma.ai](https://comma.ai/)’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.\n\ni figured i’d spend an afternoon reading it.\n\nto give an outline of how tinygrad works, we would explore these 4 stages:\n\n## the first thing i noticed\n\nto 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.\n\n``` python\nfrom tinygrad import Tensor\n\na = Tensor.rand(4, 4)\nb = Tensor.rand(4, 4)\nc = a.matmul(b)\n\nprint(c.shape)   # (4, 4)\n```\n\nthis 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)`\n\n. looks fine!!\n\nbut… no multiplication had actually happened. `a.matmul(b)`\n\ndidn’t crunch any numbers. it just built a tree of lazy operations and returned. the real work only runs when you call `c.realize()`\n\n.\nthis was new to me because, PyTorch does `a @ b`\n\nand immediately launches the matmul. but, tinygrad waits.\n\nthis 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()`\n\nis, and the moment where everything actually runs.\n\nso, i was curious what the graph looked like.\n\n```\nREDUCE(sum, axis=-1)\n  PERMUTE\n    MUL\n      RESHAPE(4,1,4)\n        a.uop\n      PERMUTE\n        RESHAPE(1,4,4)\n          b.uop\n```\n\ni stared at it for a minute. the tree reads bottom-up :)\n\ntake `a`\n\n, view it as a 4-by-1-by-4. take `b`\n\n, 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’s a 4-by-4 matrix multiplication.\n\n`a @ b`\n\ngot turned into `(a.reshape(4,1,4) * b.reshape(1,4,4).transpose()).sum(-1)`\n\n.\n\nthe matmul isn’t built into the framework as a primitive. it’s syntactic sugar. a convenient way to write reshape, broadcast multiply, and sum. `relu`\n\nis sugar for `max(x, 0)`\n\n. `sigmoid`\n\nis sugar for `1 / (1 + exp(-x))`\n\n. the framework has maybe six actual operations. everything else is convenience that decomposes before the scheduler ever sees it.\n\nso if everything is shorthand, and shorthand doesn’t trigger any computation… the framework can see your whole program at once before deciding what to run.\n\n## what’s inside a tensor\n\nnext i opened `tensor.py`\n\n. i was expecting a struct with strides and a device pointer and maybe a reference count. something that looked like it owned data.\n\n```\n__slots__ = \"uop\", \"is_param\", \"grad\"\n```\n\nthat’s it. a reference to a graph node. a flag for the optimizer. a slot for the gradient after `.backward()`\n\n.\n\nshape lives on the graph node. `self.uop.shape`\n\n. dtype, device, everything. the tensor object doesn’t carry any of that itself.\n\nthe graph node is called a **UOp**. it’s the central thing in the whole codebase. tensors are UOps. kernels are UOps. compiled binaries are UOps. a UOp has exactly four fields.\n\n`op`\n\nsays what kind of node. `dtype`\n\nis float32 or int64 or whatever. `src`\n\nis a tuple of child nodes feeding into this one. `arg`\n\nis extra data the operation needs, like the target shape for a reshape.\n\nwhen i realised everything is a UOp i went looking for where the actual operations are defined. `add`\n\n, `mul`\n\n, `matmul`\n\n. and found something that didn’t make sense at first.\n\n```\n>>> Tensor.add is UOp.add\nTrue\n```\n\nthe wrapper and the thing it wraps share the same method. not the same name. the same python function object.\n\nit works like this. both `Tensor`\n\nand `UOp`\n\ninherit from shared helper classes called **mixins**. the operations are written once in files like `mixin/elementwise.py`\n\n. both classes pick them up through inheritance.\n\nthe mixin methods call a few abstract hooks at the bottom, and each class fills in the hooks differently. `Tensor`\n\n’s version wraps results back into a new `Tensor`\n\n. `UOp`\n\n’s version creates a new graph node. the code in the middle is identical.\n\nwhat this means in practice is you can write at the `Tensor`\n\nlevel for normal stuff, or drop into raw `UOp`\n\ngraphs when you need control over the kernel. same methods either way.\n\n## the graph doesn’t repeat itself\n\nUOp nodes get cached. every combination of `(op, dtype, children, arg)`\n\nis stored once in a dictionary. if you build the same expression twice, you get back the same object.\n\n```\n>>> Tensor(3) + 4\n>>> Tensor(3) + 4\n# identical UOp object, served from cache\n```\n\nthis 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’s it.\n\na pattern matcher is a rewrite rule. “if you see this arrangement of nodes, replace it with this other arrangement.” the code has comments like `pm_move_where_on_load`\n\n. the `pm`\n\nstands for pattern matcher. it moves WHERE nodes below LOAD nodes. the whole compilation pipeline is these chained together.\n\nhere’s a real one from the codebase, `pm_cast_float_alu`\n\nin `codegen/__init__.py`\n\n:\n\n```\npm_cast_float_alu = PatternMatcher([\n  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL),\n        src=(UPat(name=\"x\"),), name=\"u\"),\n   lambda u, x: u.replace(src=(x.cast(u.dtype),))\n                if x.dtype != u.dtype else None),\n])\n```\n\nit says: find any node whose operation is `SIN`\n\n, `LOG2`\n\n, `EXP2`\n\n, `SQRT`\n\n, or `RECIPROCAL`\n\n. grab its single input `x`\n\n. if `x`\n\nis a different dtype than the result, insert a cast node between them. returning `None`\n\nmeans “no match, leave it alone.” returning a new UOp replaces the old one.\n\nthat’s the template. a `UPat`\n\ndescribes the subtree to find. a lambda takes the matched pieces and returns either a replacement or `None`\n\n. the `PatternMatcher`\n\nclass stores these pairs in a dictionary keyed by the root operation, so it only checks patterns against nodes whose op matches. no scanning irrelevant nodes.\n\nthe power of this comes from chaining. `graph_rewrite`\n\nwalks the UOp tree and applies every matching rule. when a rule fires, the tree changes, and the walker rechecks the new node against the rules. this repeats until nothing matches or a depth limit is hit. one pass can trigger cascading rewrites.\n\na few more examples from the pipeline so you can see the pattern:\n\n`pm_simplify_ranges`\n\nmatches a loop from 0 to 1 and inlines it: “this isn’t a loop, it’s one iteration, just substitute the index with 0.”\n\n`pm_load_collapse`\n\nmatches two adjacent loads of `buf[i]`\n\nand `buf[i+1]`\n\nand merges them into one wider load. fewer memory instructions, better bandwidth.\n\n`pm_move_gates_from_index`\n\ntakes a WHERE node that guards an invalid index and moves the guard to the LOAD level: “instead of computing a masked index and then loading, load a safe value and use the gate to pick between the real value and zero.” this avoids out-of-bounds reads.\n\nthe compilation pipeline is exactly this, thirty times in a row. each pass is one `PatternMatcher`\n\nwith a few patterns. the passes compose because each one creates shapes the next pass knows how to match.\n\n## what happens when you finally compute\n\nwhen i called `c.realize()`\n\n, i expected the magic to happen in one step. it doesn’t. two things happen. rangeify, then the scheduler.\n\n### views are not copies\n\nreshape and permute don’t copy data. reshape changes the shape of a tensor without moving anything in memory,\n\nfor example: `Tensor.arange(6).reshape(2, 3)`\n\nturns six numbers into a 2-by-3 grid, but the numbers stay in the same order underneath.\n\nand… permute swaps axes:`tensor.permute(1, 0)`\n\ntransposes a matrix by flipping which dimension is rows and which is columns. again, no data moves. these are just alternate ways to index into the same bytes.\n\nif you have 16 floats in a flat buffer and call `.reshape(4, 4)`\n\n, all that happens is the reshape records that element `(row, col)`\n\nmaps to byte offset `(row * 4 + col) * 4`\n\n.\n\nrangeify walks the graph and replaces every reshape, permute, and expand node with explicit loop variables called RANGEs. a reshape becomes a RANGE that loops over the new dimensions and computes the index into the flat buffer. but rangeify does something else too. it splits the full compute graph into kernels.\n\ntake the matmul kernel `r_4_4_4`\n\n. here’s what it computes:\n\n```\nfor i in range(4):\n    for j in range(4):\n        out[i][j] = sum(a[i][k] * b[k][j] for k in range(4))\n```\n\nrangeify turns this into a graph of low-level ops:\n\neach `out[i][j] = ...`\n\nis one **kernel invocation**. two loads from the input buffers, a multiply and accumulation chain, one store to the output buffer. that’s the graph structure inside every kernel.\n\nkernels are split at buffer writes. As long as the graph only reshapes, multiplies, and reduces values, those operations stay in the same kernel. Once a result is written to a buffer and later read back, that write marks the boundary between two kernels.\n\nfor my 4×4 matmul, rangeify produced eight kernels. Seven handled RNG bookkeeping: `Tensor.rand`\n\nuses a counter-based PRNG, and advancing its state is a separate computation. The matrix multiplication itself became a single kernel: `r_4_4_4`\n\n, named for its three loops, each of size 4.\n\nthis is where laziness pays off. Because reshape, multiply, and sum were all still in the graph, rangeify fused them into one kernel. An eager framework would have executed them separately, writing intermediates to memory and reading them back between each step.\n\n### ordering the kernels\n\nrangeify gives you a pile of kernels. the scheduler figures out what order to run them in.\n\neach kernel reads from some buffers and writes to others. if kernel A writes buffer X and kernel B reads buffer X, then A must run before B. if kernel B wants to write buffer X and kernel A is still reading it, B has to wait. the scheduler builds a dependency graph from these constraints and topologically sorts it into a flat list. the output is called a LINEAR.\n\nthere’s also a memory planner. if kernel A writes buffer X and kernel B is the last reader, buffer X’s memory can be handed to kernel C later. no need to allocate fresh space.\n\n## the compilation pipeline\n\neach kernel now enters the codegen. a chain of rewrite passes walks the UOp tree and transforms it. i’ll show three passes where the change is visible.\n\n`symbolic simplify`\n\nresolves shape math at compile time:\n\n```\n// before\nint offset = (Lidx1 * 4) + 0;\n// after\nint offset = Lidx1 * 4;\n```\n\n`apply_opts`\n\npicks an execution strategy. for my matmul it unrolled the inner reduction:\n\n```\n// before: a loop\nfor (int k = 0; k < 4; k++)\n    sum += a[k] * b[k];\n\n// after: straight-line multiply-adds\nsum = a0*b0 + a1*b1 + a2*b2 + a3*b3;\n```\n\n`devectorize`\n\nsplits vector ops into scalars, then `memory coalescing`\n\nre-merges adjacent ones. a `float4`\n\nadd becomes four `float`\n\nadds, and if coalescing spots four adjacent loads, it re-packs them into one `float4`\n\nload:\n\n``` php\n// devectorize: float4 -> four scalars\nfloat4 tmp = buf[0];    // before\nfloat a = buf[0];       // after\nfloat b = buf[1];\nfloat c = buf[2];\nfloat d = buf[3];\n\n// coalesce: adjacent scalars -> float4\n// four individual loads become one vector load\nfloat4 val = *((float4*)(buf));  // after\n```\n\nthe order matters. devectorize first gives coalescing full visibility into memory adjacency.\n\nafter the chain, a renderer turns the flat UOp instructions into source code. on CPU, C with GCC vector extensions. on CUDA, CUDA. on AMD GPUs, machine code bytes. then a compiler produces the binary, cached by hash.\n\nhere’s the full kernel after all sixteen passes:\n\n```\ntypedef float float4 __attribute__((aligned(16),ext_vector_type(4)));\nvoid r_4_4_4(float* restrict data0, float* restrict data1, float* restrict data2) {\n  for (int Lidx1 = 0; Lidx1 < 4; Lidx1++) {\n    int alu0 = (Lidx1<<2);\n    float4 val0 = (*((float4*)((data1+alu0))));\n    for (int Lidx2 = 0; Lidx2 < 4; Lidx2++) {\n      float val1 = (*(data2+(Lidx2+4)));\n      float val2 = (*(data2+(Lidx2+8)));\n      float val3 = (*(data2+(Lidx2+12)));\n      float val4 = (*(data2+Lidx2));\n      *(data0+(alu0+Lidx2)) = ((val0[0]*val4)+(val0[1]*val1)+(val0[2]*val2)+(val0[3]*val3));\n    }\n  }\n}\n```\n\nthe `float4`\n\nis from coalescing. the straight-line multiply-add is from unroll. `Lidx1<<2`\n\nis from symbolic simplify. every line comes from a pass you can find.\n\n## running it\n\nafter compilation, each kernel is a binary blob with metadata. the runtime’s job is to get these binaries onto the hardware.\n\na function called `run_linear`\n\nwalks the LINEAR list one entry at a time. each entry is a `CALL`\n\nnode. the CALL says what kind of work this is and which buffers are involved.\n\nfor a compute kernel, the runtime looks up the cached binary by the kernel’s hash. it computes the launch dimensions from the kernel metadata. on a GPU, this means deciding how many thread blocks and threads per block. on CPU, it’s the loop bounds. then it calls into the device driver to launch.\n\nfor a copy between devices, the runtime picks the fastest available path. if the two devices share a PCIe link, it uses direct memory access - the GPU can read from CPU memory without involving the CPU. if one end is a disk tensor backed by a file on an NVMe drive, it might use memory-mapped I/O to bypass the kernel’s page cache. the fallback is always a CPU-hosted copy, but it’s the slowest option and the runtime avoids it when it can.\n\neach device backend lives in a file under `tinygrad/runtime/`\n\n. every backend implements four things. an allocator that gives you device memory and frees it. a renderer that turns UOp instructions into source code. a compiler that turns source into a binary the device can run. a runtime that loads and launches that binary.\n\nthe CPU backend uses gcc to compile C code and launches it as a function call. the CUDA backend uses nvrtc to compile CUDA source and launches via the CUDA driver API. the Metal backend compiles MSL and dispatches through Metal’s command queue. the AMD and NVIDIA backends go further. they skip the vendor runtime entirely. they format the commands directly in the GPU’s hardware command queue format and write them to memory-mapped registers. same approach Apple’s Metal driver uses internally. no ROCm. no CUDA runtime. just Python talking to hardware.\n\nyou can watch the whole thing by turning up a number. `DEBUG=4`\n\nshows the generated source code for every kernel before it gets compiled.\n\n``` python\nDEBUG=4 python -c \"from tinygrad import Tensor; (Tensor.rand(16,16)@Tensor.rand(16,16)).realize()\"\n```\n\na matmul that took seconds to write becomes 30 lines of unrolled, vectorized C on your terminal. and you can trace which pass produced each line.\n\ni opened this codebase expecting a Python frontend with a black-box compiler underneath. instead i found 17K lines where a matmul is visible at every step. the graph you build, the kernel it becomes, the loops it unrolls, the C it generates. nothing is tucked away.\n\nit’s wild how much design went into this.\n\n- the\n**mixin** system that lets Tensor and UOp share code. - the\n**hash-consed graph** where rewriting is free. - the\n**scheduler** that fuses ops because it sees the whole program at once. - the pipeline of pattern matchers where each pass knows exactly what shape to look for.\n\nthe whole thing is at [github.com/tinygrad/tinygrad](https://github.com/tinygrad/tinygrad). the `DEBUG`\n\nflags are your way in. `1`\n\nfor kernels, `2`\n\nfor timing, `3`\n\nfor the AST, `4`\n\nfor the source. start with a matmul and follow it down.", "url": "https://wpnews.pro/news/everything-is-a-uop", "canonical_source": "https://ssenthilnathan3.github.io/blog/tinygrad/", "published_at": "2026-08-06 00:00:00+00:00", "updated_at": "2026-08-10 09:16:48.453053+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools", "ai-infrastructure"], "entities": ["tinygrad", "comma.ai", "openpilot", "PyTorch", "Tensor", "UOp"], "alternates": {"html": "https://wpnews.pro/news/everything-is-a-uop", "markdown": "https://wpnews.pro/news/everything-is-a-uop.md", "text": "https://wpnews.pro/news/everything-is-a-uop.txt", "jsonld": "https://wpnews.pro/news/everything-is-a-uop.jsonld"}}