Appendix D: Complete source code
This appendix contains the whole repository that the book builds: the compiler (uopc/), its tests
(tests/), and the example programs (examples/). Each file is included at build time from the same commit
as the text, so it always matches the listings in the chapters. Click a file name to expand it. In the PDF
every file is printed in full.
The files are in the order the course introduces them. Line counts include comments and blank lines.
D.1 The compiler: uopc/
| Week | Modules |
|---|---|
| 1 | uop.py, dtype.py, rewrite.py, symbolic.py, uir.py |
| 2–3 | render_c.py, runtime.py, linearize.py |
| 4 | rangeify.py, composite.py |
| 5 | kernelize.py, autodiff.py, tensor.py, nn.py |
| 6 | opt.py |
| 7–8 | gpu.py, cuda.py |
uopc/init.py · 13 lines
"""uopc: a machine learning compiler built from scratch, one week at a time.
week 1 uop, rewrite, symbolic, uir the IR, rewriting, algebra, the text format
week 2 render_c, runtime, linearize memory, C code, compiling and running it
week 3 linearize loops and scopes
week 4 rangeify, composite shapes -> loops, PAD/FLIP from primitives
week 5 kernelize, autodiff, tensor, nn fusion/fission, gradients, a Tensor library
week 6 opt upcast, unroll, threads, layout
week 7-8 gpu, cuda GPU mapping, CUDA, shared memory, tensor cores
"""
from uopc.uop import UOp, Ops, AxisType, ParamArg, CallArg, RangeArg, sink, range_
from uopc.dtype import dtypes, DType, AddrSpace
from uopc.tensor import Tensor
uopc/dtype.py · 72 lines
"""Week 1: data types and address spaces.
A DType is a tiny frozen record. We keep the set deliberately small:
void, bool, i32, f32 (weeks 1-6) and f16 (week 8, tensor cores).
"""
from __future__ import annotations
import math, struct
from dataclasses import dataclass
from enum import Enum, auto
@dataclass(frozen=True, eq=True)
class DType:
name: str
itemsize: int
ctype: str
fmt: str | None # struct / numpy format character
def __repr__(self): return f"dtypes.{self.name}"
@property
def is_float(self) -> bool: return self.name in ("f32", "f16")
@property
def is_int(self) -> bool: return self.name == "i32"
class dtypes:
void = DType("void", 0, "void", None)
bool = DType("bool", 1, "_Bool", "?")
i32 = DType("i32", 4, "int", "i")
f32 = DType("f32", 4, "float", "f")
f16 = DType("f16", 2, "_Float16", "e")
all = (void, bool, i32, f32, f16)
@staticmethod
def from_name(name: str) -> DType:
for d in dtypes.all:
if d.name == name: return d
raise ValueError(f"unknown dtype {name!r}")
@staticmethod
def from_py(x) -> DType:
if isinstance(x, bool): return dtypes.bool
if isinstance(x, int): return dtypes.i32
if isinstance(x, float): return dtypes.f32
raise TypeError(f"no dtype for python value {x!r}")
@staticmethod
def min(d: DType):
return {"bool": False, "i32": -2**31}.get(d.name, -math.inf)
@staticmethod
def max(d: DType):
return {"bool": True, "i32": 2**31 - 1}.get(d.name, math.inf)
def truncate(d: DType, x):
"""Bring a python value into the exact value set of dtype d (used by const folding)."""
if d == dtypes.bool: return bool(x)
if d == dtypes.i32: return ((int(x) + 2**31) % 2**32) - 2**31 # two's complement wraparound
if d == dtypes.f32: return struct.unpack("f", struct.pack("f", float(x)))[0] if math.isfinite(x) else float(x)
if d == dtypes.f16: return struct.unpack("e", struct.pack("e", float(x)))[0] if math.isfinite(x) and abs(x) < 65520 else float(x)
raise TypeError(f"can't truncate to {d}")
class AddrSpace(Enum):
"""Where a UOp lives. ALU = a plain value in a register. The rest are memory."""
ALU = auto() # a value (week 1)
GLOBAL = auto() # device memory, a function argument (week 2)
LOCAL = auto() # call-scope storage: `float acc[4];` (week 2)
SHARED = auto() # GPU workgroup shared memory (week 7)
def __repr__(self): return f"AddrSpace.{self.name}"
uopc/uop.py · 320 lines
"""Week 1: the UOp.
UOp(op, src: tuple[UOp, ...], arg)
Every program in this course is a DAG of UOps. UOps are *hash consed*:
constructing a UOp that is structurally identical to a live one returns the
very same Python object. That gives us O(1) structural equality (`is`),
free common-subexpression elimination, and cheap memoization keyed on nodes.
dtype, shape and addrspace are *recursive properties*: they are computed
from the op, the arg and the sources, never stored by the user.
"""
from __future__ import annotations
import math, weakref
from dataclasses import dataclass
from enum import Enum, auto
from typing import Any
from uopc.dtype import DType, dtypes, AddrSpace
class Ops(Enum):
# ---- week 1: values ---------------------------------------------------
PARAM = auto() # arg: ParamArg. An input to the enclosing CALL.
CALL = auto() # src[0] = body, src[1:] = arguments. arg: CallArg
CONST = auto() # arg: python value (bool/int/float). shape ()
STACK = auto() # new leading axis; all srcs share one shape
INDEX = auto() # src[0] indexed by src[1:]; pops axes off the left
CAST = auto() # arg: DType
RECIP = auto() # 1/x
ADD = auto(); MUL = auto(); CMPLT = auto()
FLOORDIV = auto(); FLOORMOD = auto()
WHERE = auto() # (cond, a, b)
# extra ALU, introduced when we need them (masks, max-reduce, softmax)
CMPNE = auto(); MAX = auto()
EXP2 = auto(); LOG2 = auto(); SQRT = auto(); SIN = auto()
# ---- week 2: memory ---------------------------------------------------
BUFFER = auto() # arg: ParamArg. global persistent memory (never in a CALL body)
ALLOC = auto() # arg: ParamArg. call-scope storage `float a[10];`
LOAD = auto() # (ptr,)
STORE = auto() # (ptr, value)
AFTER = auto() # (mem, *deps): `mem`, but only once deps have happened
SINK = auto() # (*effects): groups side effects, the root of a program
# ---- week 3: loops ----------------------------------------------------
RANGE = auto() # src[0] = end; arg: RangeArg. value is the loop counter
END = auto() # (passthrough, *ranges): closes the ranges
# ---- week 4: movement + reduction -------------------------------------
RESHAPE = auto(); EXPAND = auto(); PERMUTE = auto(); SHRINK = auto()
REDUCE = auto() # arg: (Ops.ADD | Ops.MAX, n) -- reduces the last n axes
# ---- week 5: kernels ----------------------------------------------------
STAGE = auto() # materialize src[0] into a temp buffer you can index into
# ---- weeks 7-8: GPUs ----------------------------------------------------
BARRIER = auto() # (*deps): __syncthreads()
WMMA = auto() # reserved for tensor cores (week 8 lowers GEMMs to WMMA by pattern instead)
def __repr__(self): return f"Ops.{self.name}"
__hash__ = object.__hash__ # members are singletons: identity hashing is correct and much faster
UNARY = {Ops.RECIP, Ops.CAST, Ops.EXP2, Ops.LOG2, Ops.SQRT, Ops.SIN}
BINARY = {Ops.ADD, Ops.MUL, Ops.CMPLT, Ops.CMPNE, Ops.FLOORDIV, Ops.FLOORMOD, Ops.MAX}
ALU = UNARY | BINARY | {Ops.WHERE}
COMMUTATIVE = {Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE}
MOVEMENT = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.SHRINK}
class AxisType(Enum):
LOOP = auto() # a plain serial for-loop
REDUCE = auto() # a for-loop that a reduction accumulates over
UPCAST = auto() # fully unrolled (week 6). uopc unrolls *during* rangeify (Opts.upcast), so no RANGE of
# this kind is ever built; the name stays for the course's vocabulary
THREAD = auto() # parallel across CPU threads (week 6)
GLOBAL = auto() # GPU blockIdx (week 7)
LOCAL = auto() # GPU threadIdx (week 7)
def __repr__(self): return f"AxisType.{self.name}"
__hash__ = object.__hash__
@dataclass(frozen=True)
class ParamArg:
slot: int
dtype: DType
shape: tuple[int, ...] = ()
addrspace: AddrSpace | None = None # only used by ALLOC (LOCAL or SHARED)
@dataclass(frozen=True)
class CallArg:
name: str = "f"
@dataclass(frozen=True, order=True)
class RangeArg:
idx: int # loops nest in increasing idx order
kind: AxisType = AxisType.LOOP
def prod(xs) -> int: return math.prod(xs)
def _argkey(arg):
# 1 == 1.0 == True and 0.0 == -0.0 in python; the cache must not merge them.
if isinstance(arg, float): return (float, repr(arg))
if isinstance(arg, tuple): return (tuple, tuple(_argkey(a) for a in arg))
return (type(arg), arg)
class UOp:
__slots__ = ("op", "src", "arg", "dtype", "shape", "addrspace", "__weakref__", "__dict__")
_cache: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
def __new__(cls, op: Ops, src: tuple[UOp, ...] = (), arg: Any = None):
key = (op, src, _argkey(arg))
if (ret := UOp._cache.get(key)) is not None: return ret
ret = super().__new__(cls)
ret.op, ret.src, ret.arg = op, src, arg
ret.dtype, ret.shape, ret.addrspace = _infer(op, src, arg)
UOp._cache[key] = ret
return ret
# hash consing means identity *is* structural equality, so we keep object's default
# __eq__/__hash__ (identity, implemented in C). Note __lt__ etc. build UOps instead.
def __reduce__(self): return (UOp, (self.op, self.src, self.arg))
def __repr__(self):
s = f"UOp({self.op}, arg={self.arg!r}" if self.arg is not None else f"UOp({self.op}"
return s + (f", src=({', '.join(repr(x) for x in self.src)},))" if self.src else ")")
def replace(self, **kw) -> UOp:
return UOp(kw.get("op", self.op), kw.get("src", self.src), kw.get("arg", self.arg))
# ---- graph helpers ----------------------------------------------------
def toposort(self) -> list[UOp]:
"""Every node reachable from self, children before parents. Iterative (graphs get deep)."""
out, seen, stack = [], set(), [(self, False)]
while stack:
u, done = stack.pop()
if done: out.append(u); continue
if u in seen: continue
seen.add(u)
stack.append((u, True))
stack.extend((s, False) for s in reversed(u.src) if s not in seen)
return out
def substitute(self, mapping: dict[UOp, UOp]) -> UOp:
"""Replace nodes (by identity) everywhere in the graph. Rebuilds only what changed."""
new: dict[UOp, UOp] = {}
for u in self.toposort():
if u in mapping: new[u] = mapping[u]; continue
src = tuple(new[s] for s in u.src)
new[u] = u if src == u.src else u.replace(src=src)
return new[self]
@staticmethod
def wrap(roots: list[UOp], targets: set[UOp], fn) -> list[UOp]:
"""Rebuild the graph bottom-up, replacing each node t in targets by fn(t'), where t' is t
with *its* sources already rebuilt. (substitute() would wrap the original, unrewritten t.)"""
new: dict[UOp, UOp] = {}
for u in UOp(Ops.SINK, tuple(roots)).toposort():
src = tuple(new[s] for s in u.src)
n = u if src == u.src else u.replace(src=src)
new[u] = fn(n) if u in targets else n
return [new[r] for r in roots]
def ranges(self) -> list[UOp]:
return sorted((u for u in self.toposort() if u.op is Ops.RANGE), key=lambda u: u.arg)
# ---- constructors -----------------------------------------------------
@staticmethod
def const(v, dtype: DType | None = None) -> UOp:
if dtype == dtypes.f16: return UOp(Ops.CONST, (), float(v)).cast(dtypes.f16) # python has no f16 values
if dtype is not None and dtype != dtypes.from_py(v):
v = {"bool": bool, "i32": int}.get(dtype.name, float)(v)
return UOp(Ops.CONST, (), v)
def const_like(self, v) -> UOp:
"""A constant with self's dtype *and shape* (broadcast with RESHAPE + EXPAND)."""
c = UOp.const(v, self.dtype)
return c.reshape((1,) * len(self.shape)).expand(self.shape) if self.shape else c
def _c(self, x) -> UOp: return x if isinstance(x, UOp) else self.const_like(x)
def alu(self, op: Ops, *src) -> UOp: return UOp(op, (self, *(self._c(s) for s in src)))
def __add__(self, x): return self.alu(Ops.ADD, x)
def __radd__(self, x): return self._c(x).alu(Ops.ADD, self)
def __mul__(self, x): return self.alu(Ops.MUL, x)
def __rmul__(self, x): return self._c(x).alu(Ops.MUL, self)
def __neg__(self): return self.alu(Ops.MUL, -1 if self.dtype == dtypes.i32 else -1.0)
def __sub__(self, x): return self + -self._c(x)
def __rsub__(self, x): return self._c(x) + -self
def __floordiv__(self, x): return self.alu(Ops.FLOORDIV, x)
def __mod__(self, x): return self.alu(Ops.FLOORMOD, x)
def __lt__(self, x): return self.alu(Ops.CMPLT, x)
def __gt__(self, x): return self._c(x).alu(Ops.CMPLT, self)
def ne(self, x): return self.alu(Ops.CMPNE, x)
def eq(self, x): return self.ne(x).ne(True)
def maximum(self, x): return self.alu(Ops.MAX, x)
def recip(self): return UOp(Ops.RECIP, (self,))
def cast(self, d: DType): return self if d == self.dtype else UOp(Ops.CAST, (self,), d)
def where(self, a, b):
a = a if isinstance(a, UOp) else (b.const_like(a) if isinstance(b, UOp) else UOp.const(a))
b = b if isinstance(b, UOp) else a.const_like(b)
return UOp(Ops.WHERE, (self, a, b))
def index(self, *idx) -> UOp:
return UOp(Ops.INDEX, (self, *(i if isinstance(i, UOp) else UOp.const(int(i)) for i in idx)))
def load(self) -> UOp: return UOp(Ops.LOAD, (self,))
def store(self, v) -> UOp: return UOp(Ops.STORE, (self, v))
def after(self, *deps) -> UOp: return UOp(Ops.AFTER, (self, *deps)) if deps else self
def end(self, *rngs) -> UOp: return UOp(Ops.END, (self, *rngs)) if rngs else self
def reshape(self, shape): return UOp(Ops.RESHAPE, (self,), tuple(shape))
def expand(self, shape): return UOp(Ops.EXPAND, (self,), tuple(shape))
def permute(self, perm): return UOp(Ops.PERMUTE, (self,), tuple(perm))
def shrink(self, arg): return UOp(Ops.SHRINK, (self,), tuple(tuple(x) for x in arg))
def reduce(self, op: Ops, n: int = 1): return UOp(Ops.REDUCE, (self,), (op, n))
def sink(*srcs: UOp) -> UOp: return UOp(Ops.SINK, tuple(srcs))
def range_(end: int | UOp, idx: int, kind: AxisType = AxisType.LOOP) -> UOp:
return UOp(Ops.RANGE, (end if isinstance(end, UOp) else UOp.const(int(end)),), RangeArg(idx, kind))
# ---------------------------------------------------------------------------
# shape / dtype / addrspace inference. This is the type checker of the IR:
# every malformed UOp dies here, at construction time, with a useful message.
class UOpError(Exception): pass
def _check(cond, msg):
if not cond: raise UOpError(msg)
def _infer(op: Ops, src: tuple[UOp, ...], arg) -> tuple[DType, tuple[int, ...], AddrSpace]:
A = AddrSpace.ALU
if op is Ops.CONST:
_check(not src, "CONST has no sources")
return dtypes.from_py(arg), (), A
if op is Ops.PARAM: return arg.dtype, arg.shape, A
if op is Ops.BUFFER:
_check(len(arg.shape) == 1, "BUFFER is flat")
return arg.dtype, arg.shape, AddrSpace.GLOBAL
if op is Ops.ALLOC:
_check(len(arg.shape) == 1, "ALLOC is flat")
return arg.dtype, arg.shape, arg.addrspace or AddrSpace.LOCAL
if op is Ops.CALL:
params = {u.arg.slot: u for u in call_params(src[0])}
_check(sorted(params) == list(range(len(src) - 1)), f"CALL needs params 0..{len(src)-2}, body has {sorted(params)}")
for i, a in enumerate(src[1:]):
p = params[i]
_check(a.shape == p.shape and a.dtype == p.dtype, f"CALL arg {i}: {a.dtype}{a.shape} != param {p.dtype}{p.shape}")
return src[0].dtype, src[0].shape, src[0].addrspace
if op is Ops.STACK:
_check(len(src) > 0, "STACK of nothing")
_check(all(s.shape == src[0].shape and s.dtype == src[0].dtype for s in src), "STACK srcs must agree in shape and dtype")
return src[0].dtype, (len(src),) + src[0].shape, A
if op is Ops.INDEX:
x, idx = src[0], src[1:]
_check(len(idx) <= len(x.shape), f"INDEX: {len(idx)} indices into shape {x.shape}")
_check(all(i.dtype == dtypes.i32 and i.shape == () for i in idx), "INDEX indices are i32 scalars")
return x.dtype, x.shape[len(idx):], x.addrspace
if op in ALU:
if op is Ops.WHERE:
c, a, b = src
_check(c.dtype == dtypes.bool, "WHERE condition must be bool")
_check(a.dtype == b.dtype, f"WHERE branches disagree: {a.dtype} vs {b.dtype}")
_check(c.shape == a.shape == b.shape, f"WHERE shape mismatch {c.shape} {a.shape} {b.shape}")
return a.dtype, a.shape, A
if op in UNARY:
_check(len(src) == 1, f"{op} is unary")
if op is Ops.CAST: return arg, src[0].shape, A
_check(src[0].dtype.is_float, f"{op} needs a float, got {src[0].dtype}")
return src[0].dtype, src[0].shape, A
a, b = src
_check(a.shape == b.shape, f"{op}: shape mismatch {a.shape} vs {b.shape} (broadcast with EXPAND)")
_check(a.dtype == b.dtype, f"{op}: dtype mismatch {a.dtype} vs {b.dtype} (use CAST)")
if op in (Ops.FLOORDIV, Ops.FLOORMOD): _check(a.dtype == dtypes.i32, f"{op} is integer only")
return (dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else a.dtype), a.shape, A
if op is Ops.LOAD:
_check(src[0].addrspace is not A, "LOAD needs memory")
return src[0].dtype, src[0].shape, A
if op is Ops.STORE:
_check(src[0].addrspace is not A, "STORE needs memory")
_check(src[0].dtype == src[1].dtype, f"STORE of {src[1].dtype} into {src[0].dtype}")
return dtypes.void, (), A
if op is Ops.AFTER: return src[0].dtype, src[0].shape, src[0].addrspace
if op in (Ops.SINK, Ops.END, Ops.BARRIER): return dtypes.void, (), A
if op is Ops.RANGE:
_check(src[0].dtype == dtypes.i32, "RANGE end is i32")
return dtypes.i32, (), A
if op is Ops.RESHAPE:
_check(prod(arg) == prod(src[0].shape), f"RESHAPE {src[0].shape} -> {arg}")
return src[0].dtype, arg, src[0].addrspace
if op is Ops.EXPAND:
s = src[0].shape
_check(len(arg) == len(s) and all(a == b or a == 1 for a, b in zip(s, arg)), f"EXPAND {s} -> {arg}")
return src[0].dtype, arg, src[0].addrspace
if op is Ops.PERMUTE:
_check(sorted(arg) == list(range(len(src[0].shape))), f"PERMUTE {arg} of {src[0].shape}")
return src[0].dtype, tuple(src[0].shape[i] for i in arg), src[0].addrspace
if op is Ops.SHRINK:
s = src[0].shape
_check(len(arg) == len(s) and all(0 <= b <= e <= n for (b, e), n in zip(arg, s)), f"SHRINK {arg} of {s}")
return src[0].dtype, tuple(e - b for b, e in arg), src[0].addrspace
if op is Ops.REDUCE:
rop, n = arg
_check(rop in (Ops.ADD, Ops.MAX) and 0 <= n <= len(src[0].shape), f"REDUCE {arg} of {src[0].shape}")
return src[0].dtype, src[0].shape[:len(src[0].shape) - n], A
if op is Ops.STAGE: return src[0].dtype, src[0].shape, A
if op is Ops.WMMA: return src[2].dtype, src[2].shape, A
raise UOpError(f"no inference rule for {op}")
def call_params(body: UOp) -> list[UOp]:
"""PARAMs that belong to this body. A nested CALL's body is its own scope, so we don't look inside it."""
out, seen, stack = [], set(), [body]
while stack:
u = stack.pop()
if u in seen: continue
seen.add(u)
if u.op is Ops.PARAM: out.append(u)
stack.extend(u.src[1:] if u.op is Ops.CALL else u.src)
return out
uopc/rewrite.py · 115 lines
"""Week 1: pattern matching and graph rewriting.
A *rewrite rule* is (pattern, function). The pattern is a UPat tree that matches
UOps structurally and binds names; the function receives the bindings as keyword
arguments and returns a replacement UOp, or None for "rule does not apply".
graph_rewrite applies a PatternMatcher bottom-up and greedily: children are
rewritten to a fixed point before their parent is looked at, and whenever a rule
fires the result is itself rewritten again. Because UOps are hash consed, the
memo table is just a dict keyed by node.
"""
from __future__ import annotations
import itertools
from typing import Any, Callable
from uopc.uop import UOp, Ops
from uopc.dtype import DType
class UPat:
__slots__ = ("op", "src", "name", "arg", "dtype", "allow_any_len")
def __init__(self, op: Ops | tuple[Ops, ...] | set[Ops] | None = None, src=None, name: str | None = None,
arg: Any = None, dtype: DType | tuple[DType, ...] | None = None, allow_any_len: bool = False):
self.op = None if op is None else (frozenset([op]) if isinstance(op, Ops) else frozenset(op))
self.dtype = None if dtype is None else (frozenset([dtype]) if isinstance(dtype, DType) else frozenset(dtype))
# src: tuple -> match in order. list -> match in any order (for commutative ops).
self.src, self.name, self.arg, self.allow_any_len = src, name, arg, allow_any_len
@staticmethod
def var(name: str | None = None, dtype=None) -> UPat: return UPat(name=name, dtype=dtype)
@staticmethod
def cvar(name: str | None = None, dtype=None) -> UPat: return UPat(Ops.CONST, name=name, dtype=dtype)
def __repr__(self): return f"UPat({self.op and set(self.op)}, name={self.name})"
def match(self, u: UOp, store: dict[str, UOp]) -> list[dict[str, UOp]]:
"""All ways this pattern matches u, each a dict of bindings."""
if self.op is not None and u.op not in self.op: return []
if self.dtype is not None and u.dtype not in self.dtype: return []
if self.arg is not None and u.arg != self.arg: return []
if self.name is not None:
if self.name in store and store[self.name] is not u: return [] # same name must bind same node
store = {**store, self.name: u}
if self.src is None: return [store]
orders = itertools.permutations(self.src) if isinstance(self.src, list) else [self.src]
out = []
for pats in orders:
if len(pats) != len(u.src) and not (self.allow_any_len and len(pats) <= len(u.src)): continue
stores = [store]
for p, s in zip(pats, u.src):
stores = [m for st in stores for m in p.match(s, st)]
if not stores: break
out.extend(stores)
return out
# sugar so rules read like math: UPat.var("x") + UPat.cvar("c")
def _bin(self, op, other, swap=False):
other = other if isinstance(other, UPat) else UPat(Ops.CONST, arg=other)
src = [other, self] if swap else [self, other]
return UPat(op, src if op in (Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE) else tuple(src))
def __add__(self, o): return self._bin(Ops.ADD, o)
def __mul__(self, o): return self._bin(Ops.MUL, o)
def __floordiv__(self, o): return self._bin(Ops.FLOORDIV, o)
def __mod__(self, o): return self._bin(Ops.FLOORMOD, o)
def __lt__(self, o): return self._bin(Ops.CMPLT, o)
Rule = tuple[UPat, Callable[..., UOp | None]]
class PatternMatcher:
def __init__(self, rules: list[Rule]):
self.rules = rules
# index by root op so we only try rules that could possibly match
self.by_op: dict[Ops, list[Rule]] = {}
for pat, fn in rules:
for op in (pat.op if pat.op is not None else Ops):
self.by_op.setdefault(op, []).append((pat, fn))
def __add__(self, other: PatternMatcher) -> PatternMatcher: return PatternMatcher(self.rules + other.rules)
def rewrite(self, u: UOp, ctx=None) -> UOp | None:
for pat, fn in self.by_op.get(u.op, ()):
for binding in pat.match(u, {}):
kw = {**binding, "ctx": ctx} if ctx is not None else binding
if (ret := fn(**kw)) is not None and ret is not u: return ret
return None
def graph_rewrite(root: UOp, pm: PatternMatcher, ctx=None, max_steps: int = 10_000_000) -> UOp:
"""Bottom-up greedy rewrite to a fixed point. Iterative, so deep graphs don't blow the stack.
Stage 0: schedule children. Stage 1: children done -> rebuild, try rules.
Stage 2: a rule fired and we rewrote its result -> record that as our answer.
"""
done: dict[UOp, UOp] = {}
stack: list[tuple[UOp, int, UOp | None]] = [(root, 0, None)]
steps = 0
while stack:
u, stage, aux = stack.pop()
if u in done and stage != 2: continue
if stage == 0:
stack.append((u, 1, None))
stack.extend((s, 0, None) for s in reversed(u.src) if s not in done)
elif stage == 1:
src = tuple(done[s] for s in u.src)
n = u if src == u.src else u.replace(src=src)
steps += 1
if steps > max_steps: raise RuntimeError("graph_rewrite did not converge (a rule loops?)")
if (m := pm.rewrite(n, ctx)) is None: done[u] = n
elif m in done: done[u] = done[m]
else: stack.extend([(u, 2, m), (m, 0, None)])
else:
done[u] = done[aux]
return done[root]
uopc/symbolic.py · 302 lines
"""Week 1: constant folding, symbolic simplification, and "running" programs by rewriting.
Three rule sets:
* fold -- evaluate ALU ops whose inputs are all constants
* symbolic -- algebra: identities, reassociation, div/mod, bounds (vmin/vmax)
* evaluate -- push ALU through STACK, resolve INDEX of STACK, inline CALLs
graph_rewrite(program, fold + symbolic + evaluate) on a closed program (no PARAMs,
no memory) reduces it all the way to CONST/STACK: rewriting *is* the interpreter.
"""
from __future__ import annotations
import math
from uopc.uop import UOp, Ops, ALU, call_params, prod
from uopc.dtype import dtypes, truncate
from uopc.rewrite import UPat, PatternMatcher, graph_rewrite
# ---------------------------------------------------------------------------
# 1. the reference semantics of every ALU op. The C renderer must agree with this.
def _safe(f, *a):
try: return f(*a)
except (OverflowError, ValueError, ZeroDivisionError): return None
def exec_alu(op: Ops, dtype, vals: list):
# inputs first: an f32 CONST may hold a double (16777217.0) that C rounds when it parses the literal
if op is not Ops.CAST and dtype.is_float: vals = [truncate(dtype, v) if isinstance(v, float) else v for v in vals]
if op is Ops.WHERE and dtype.is_float: vals = [vals[0]] + [truncate(dtype, v) for v in vals[1:]]
match op:
case Ops.ADD: r = vals[0] + vals[1]
case Ops.MUL: r = vals[0] * vals[1] if not (dtype.is_float and any(math.isinf(v) for v in vals) and 0 in vals) else math.nan
case Ops.MAX: r = max(vals[0], vals[1])
case Ops.CMPLT: return vals[0] < vals[1]
case Ops.CMPNE: return vals[0] != vals[1]
case Ops.FLOORDIV:
if vals[1] == 0 or (dtype == dtypes.i32 and vals[0] == -2**31 and vals[1] == -1): return None # C traps here
r = vals[0] // vals[1]
case Ops.FLOORMOD: r = vals[0] % vals[1] if vals[1] != 0 else None
case Ops.WHERE: r = vals[1] if vals[0] else vals[2]
case Ops.RECIP: r = 1.0 / vals[0] if vals[0] != 0 else math.copysign(math.inf, vals[0])
case Ops.EXP2: r = _safe(math.pow, 2.0, vals[0]) if vals[0] < 1024 else math.inf
case Ops.LOG2: r = math.log2(vals[0]) if vals[0] > 0 else (-math.inf if vals[0] == 0 else math.nan)
case Ops.SQRT: r = math.sqrt(vals[0]) if vals[0] >= 0 else math.nan
case Ops.SIN: r = math.sin(vals[0]) if math.isfinite(vals[0]) else math.nan
case Ops.CAST:
v = vals[0]
if dtype == dtypes.i32:
if not math.isfinite(v): return None # undefined behavior in C: refuse to fold
r = int(v) # C truncates toward zero
elif dtype == dtypes.bool: return bool(v)
else: r = float(v)
case _: return None
return None if r is None else truncate(dtype, r)
# ---------------------------------------------------------------------------
# 2. bounds analysis. vmin/vmax are conservative: the true value is always inside.
def bounds(u: UOp) -> tuple:
if (b := u.__dict__.get("_bounds")) is not None: return b
b = _bounds(u)
if u.dtype == dtypes.i32 and (b[0] < dtypes.min(u.dtype) or b[1] > dtypes.max(u.dtype)):
b = (dtypes.min(u.dtype), dtypes.max(u.dtype)) # might wrap around: give up
u.__dict__["_bounds"] = b
return b
def vmin(u: UOp): return bounds(u)[0]
def vmax(u: UOp): return bounds(u)[1]
def _bounds(u: UOp) -> tuple:
d = u.dtype
full = (dtypes.min(d), dtypes.max(d))
if u.op is Ops.CONST: return (u.arg, u.arg)
if u.op is Ops.RANGE: return (0, max(0, vmax(u.src[0]) - 1))
if d.is_float or u.shape != (): return full
if u.op is Ops.CAST and u.src[0].dtype == dtypes.bool: return (int(vmin(u.src[0])), int(vmax(u.src[0])))
if u.op is Ops.WHERE:
(c0, c1), (a0, a1), (b0, b1) = bounds(u.src[0]), bounds(u.src[1]), bounds(u.src[2])
if c0: return (a0, a1)
if not c1: return (b0, b1)
return (min(a0, b0), max(a1, b1))
if u.op not in ALU or len(u.src) != 2: return full
(a0, a1), (b0, b1) = bounds(u.src[0]), bounds(u.src[1])
match u.op:
case Ops.ADD: return (a0 + b0, a1 + b1)
case Ops.MUL:
if any(math.isinf(x) for x in (a0, a1, b0, b1)): return full
ps = [a0 * b0, a0 * b1, a1 * b0, a1 * b1]
return (min(ps), max(ps))
case Ops.MAX: return (max(a0, b0), max(a1, b1))
case Ops.FLOORDIV if b0 == b1 and b0 > 0: return (a0 // b0, a1 // b0)
case Ops.FLOORMOD if b0 == b1 and b0 > 0:
return (a0, a1) if 0 <= a0 and a1 < b0 else (0, b0 - 1) # floor mod by c>0 is always in [0, c)
case Ops.CMPLT: return (True, True) if a1 < b0 else (False, False) if a0 >= b1 else (False, True)
case Ops.CMPNE: return (True, True) if a1 < b0 or b1 < a0 else (False, False) if a0 == a1 == b0 == b1 else (False, True)
return full
# ---------------------------------------------------------------------------
# 3. rules
def _fold(x: UOp) -> UOp | None:
if x.shape != () or not all(s.op is Ops.CONST for s in x.src): return None
r = exec_alu(x.op, x.dtype, [s.arg for s in x.src])
return None if r is None else UOp.const(r, x.dtype)
fold = PatternMatcher([(UPat(ALU, name="x"), _fold)])
def _divides(u: UOp, c: int) -> UOp | None:
"""u / c if every term of u is a multiple of c, else None."""
if u.op is Ops.CONST: return u.const_like(u.arg // c) if u.arg % c == 0 else None
if u.op is Ops.MUL and u.src[1].op is Ops.CONST:
return u.src[0] * (u.src[1].arg // c) if u.src[1].arg % c == 0 else None
if u.op is Ops.ADD:
if (a := _divides(u.src[0], c)) is not None and (b := _divides(u.src[1], c)) is not None: return a + b
return None
def _split_terms(u: UOp) -> list[UOp]:
return _split_terms(u.src[0]) + _split_terms(u.src[1]) if u.op is Ops.ADD else [u]
def _div_sum(x: UOp, c: UOp, mod: bool) -> UOp | None:
"""(a*c + b) // c == a + b//c and (a*c + b) % c == b % c. Always true for *floor* division."""
if c.arg <= 0: return None
terms = _split_terms(x)
quot, rest = [], []
for t in terms: (quot if _divides(t, c.arg) is not None else rest).append(t)
if not quot: return None
r = sum(rest[1:], rest[0]) if rest else x.const_like(0)
if mod: return r % c
return sum((_divides(t, c.arg) for t in quot[1:]), _divides(quot[0], c.arg)) + r // c
def _const_bounds(x: UOp) -> UOp | None:
# an int/bool expression whose bounds pin it to one value *is* that value
if x.shape == () and not x.dtype.is_float and x.op in ALU and (b := bounds(x))[0] == b[1]:
return UOp.const(b[0], x.dtype)
return None
def _lt_mul(x: UOp, c1: UOp, c2: UOp) -> UOp | None:
# x*c1 < c2 <=> x < ceil(c2/c1) for c1 > 0
return x < -(-c2.arg // c1.arg) if c1.arg > 0 else None
def _i32(v: int) -> int: return truncate(dtypes.i32, v) # int32 arithmetic wraps, so folded constants must too
x, y, z = UPat.var("x"), UPat.var("y"), UPat.var("z")
c1, c2 = UPat.cvar("c1"), UPat.cvar("c2")
ints = (dtypes.i32, dtypes.bool)
xi = UPat.var("x", dtype=dtypes.i32)
symbolic = fold + PatternMatcher([
# canonical form: constants on the right of commutative ops
(UPat((Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE), src=(UPat(Ops.CONST, name="c"), UPat(name="y")), name="u"),
lambda u, c, y: u.replace(src=(y, c)) if y.op is not Ops.CONST else None),
# identities
(x + 0, lambda x: x), (x + 0.0, lambda x: x), # x+0.0 is not exact for x=-0.0; we accept that
(x * 1, lambda x: x), (x * 1.0, lambda x: x),
(UPat.var("x", dtype=ints) * 0, lambda x: x.const_like(0)),
(x // 1, lambda x: x), (x % 1, lambda x: x.const_like(0)),
(UPat(Ops.MAX, src=(x, x)), lambda x: x),
(UPat(Ops.WHERE, src=(UPat(), x, x)), lambda x: x),
(UPat(Ops.WHERE, src=(UPat.cvar("c"), x, y)), lambda c, x, y: x if c.arg else y),
(UPat(Ops.WHERE, src=(z, UPat(Ops.WHERE, src=(z, x, UPat())), y)), lambda z, x, y: z.where(x, y)),
(UPat(Ops.WHERE, src=(z, x, UPat(Ops.WHERE, src=(z, UPat(), y)))), lambda z, x, y: z.where(x, y)),
(UPat(Ops.CAST, src=(x,), name="c"), lambda c, x: x if c.dtype == x.dtype else None),
(UPat(Ops.CMPNE, src=(UPat.var("x", dtype=dtypes.bool), UPat(Ops.CONST, arg=False))), lambda x: x),
(UPat((Ops.CMPLT, Ops.CMPNE), src=(xi, xi)), lambda x: UOp.const(False)),
(xi + xi, lambda x: x * 2),
(xi + xi * -1, lambda x: x.const_like(0)),
# reassociate constants outward, so that they can meet and fold
((xi + c1) + c2, lambda x, c1, c2: x + _i32(c1.arg + c2.arg)), # ints only: float addition isn't associative
((xi * c1) * c2, lambda x, c1, c2: x * _i32(c1.arg * c2.arg)),
((xi + c1) * c2, lambda x, c1, c2: x * c2 + _i32(c1.arg * c2.arg)),
# distribute constants over sums so div/mod can see every term: (a + b) * c -> a*c + b*c
(UPat(Ops.MUL, src=(UPat(Ops.ADD, src=(xi, UPat.var("y")), name="s"), UPat.cvar("c"))), lambda x, y, s, c: x * c + y * c),
(UPat(Ops.ADD, src=[xi + c1, UPat.var("y")]), lambda x, c1, y: (x + y) + c1 if y.op is not Ops.CONST else None),
(xi * c1 + xi * c2, lambda x, c1, c2: x * _i32(c1.arg + c2.arg)),
(xi * c1 + xi, lambda x, c1: x * _i32(c1.arg + 1)),
# division and modulo (floor semantics make these unconditional)
((xi // c1) // c2, lambda x, c1, c2: x // (c1.arg * c2.arg) if c1.arg > 0 and c2.arg > 0 and c1.arg * c2.arg < 2**31 else None),
((xi % c1) % c2, lambda x, c1, c2: x % c2 if c2.arg > 0 and c1.arg % c2.arg == 0 else None),
((xi % c1) // c2, lambda x, c1, c2: (x // c2) % (c1.arg // c2.arg) if c2.arg > 0 and c1.arg % c2.arg == 0 else None),
(UPat(Ops.FLOORDIV, src=(xi, UPat.cvar("c"))), lambda x, c: _div_sum(x, c, mod=False)),
(UPat(Ops.FLOORMOD, src=(xi, UPat.cvar("c"))), lambda x, c: _div_sum(x, c, mod=True)),
(UPat(Ops.FLOORDIV, src=(xi, UPat.cvar("c"))), lambda x, c: x.const_like(0) if 0 <= vmin(x) and vmax(x) < c.arg else None),
(UPat(Ops.FLOORMOD, src=(xi, UPat.cvar("c"))), lambda x, c: x if 0 <= vmin(x) and vmax(x) < c.arg else None),
# recombine: (x//c)*c + x%c == x
(UPat(Ops.ADD, src=[(xi // c1) * c1, xi % c1]), lambda x, c1: x),
(UPat(Ops.ADD, src=[(xi // c1) * c1 + y, xi % c1]), lambda x, c1, y: x + y),
# comparisons
((xi + c1) < c2, lambda x, c1, c2: x < (c2.arg - c1.arg) if -2**31 <= c2.arg - c1.arg < 2**31 else None),
((xi * c1) < c2, _lt_mul),
# bounds say this int expression is a constant
(UPat(ALU, name="x"), _const_bounds),
])
# ---------------------------------------------------------------------------
# 4. evaluation: STACK, INDEX and CALL
def _elem(s: UOp, i: int) -> UOp: return s.src[i] if s.op is Ops.STACK else s.index(i)
def _push_stack(x: UOp) -> UOp | None:
# f(STACK(a0, a1), STACK(b0, b1)) -> STACK(f(a0, b0), f(a1, b1))
if x.shape == () or not any(s.op is Ops.STACK for s in x.src): return None
return UOp(Ops.STACK, tuple(x.replace(src=tuple(_elem(s, i) for s in x.src)) for i in range(x.shape[0])))
def _index_stack(s: UOp, i: UOp, rest: tuple[UOp, ...]) -> UOp | None:
if not 0 <= i.arg < len(s.src): raise IndexError(f"index {i.arg} out of bounds for STACK of {len(s.src)}")
return s.src[i.arg].index(*rest) if rest else s.src[i.arg]
def inline_call(call: UOp) -> UOp:
"""Beta reduction: replace each PARAM of the body with its argument.
Scoped: a nested CALL's body (its src[0]) has its own PARAMs, so we never rewrite inside it."""
args = {p: call.src[1 + p.arg.slot] for p in call_params(call.src[0])}
order, seen, stack = [], set(), [(call.src[0], False)]
while stack:
u, done = stack.pop()
if done: order.append(u); continue
if u in seen: continue
seen.add(u)
stack.append((u, True))
stack.extend((s, False) for s in (u.src[1:] if u.op is Ops.CALL else u.src))
new: dict[UOp, UOp] = {}
for u in order:
if u in args: new[u] = args[u]; continue
src = (u.src[0], *(new[s] for s in u.src[1:])) if u.op is Ops.CALL else tuple(new[s] for s in u.src)
new[u] = u if src == u.src else u.replace(src=src)
return new[call.src[0]]
def _literal(u: UOp):
"""A CONST/STACK tree as (flat list of python values, shape), or None if anything is symbolic."""
if u.op is Ops.CONST: return [u.arg], ()
if u.op is not Ops.STACK: return None
parts = [_literal(s) for s in u.src]
if any(p is None for p in parts): return None
return [v for p in parts for v in p[0]], u.shape
def _from_flat(flat: list, shape: tuple, dtype) -> UOp:
if not shape: return UOp.const(flat[0], dtype)
n = len(flat) // shape[0]
return UOp(Ops.STACK, tuple(_from_flat(flat[i*n:(i+1)*n], shape[1:], dtype) for i in range(shape[0])))
def _unravel(i: int, shape) -> list[int]:
out = []
for s in reversed(shape): out.append(i % s); i //= s
return out[::-1]
def _ravel(idx, shape) -> int:
i = 0
for x, s in zip(idx, shape): i = i * s + x
return i
def _eval_movement(u: UOp) -> UOp | None:
"""Movement ops and REDUCE on fully known values: just compute them."""
if (lit := _literal(u.src[0])) is None: return None
flat, ishape = lit
oshape, n_out = u.shape, prod(u.shape)
match u.op:
case Ops.RESHAPE: out = flat
case Ops.EXPAND: out = [flat[_ravel([0 if s == 1 else x for x, s in zip(_unravel(i, oshape), ishape)], ishape)] for i in range(n_out)]
case Ops.PERMUTE:
inv = [u.arg.index(k) for k in range(len(ishape))]
out = [flat[_ravel([_unravel(i, oshape)[inv[k]] for k in range(len(ishape))], ishape)] for i in range(n_out)]
case Ops.SHRINK: out = [flat[_ravel([x + b for x, (b, _) in zip(_unravel(i, oshape), u.arg)], ishape)] for i in range(n_out)]
case Ops.REDUCE:
rop, k = u.arg
inner = prod(ishape[len(ishape) - k:])
out = []
for i in range(n_out):
acc = flat[i * inner]
for v in flat[i * inner + 1:(i + 1) * inner]: acc = exec_alu(rop, u.dtype, [acc, v])
out.append(acc)
return _from_flat(out, oshape, u.dtype)
evaluate = PatternMatcher([
(UPat({Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.SHRINK, Ops.REDUCE}, name="u"), _eval_movement),
(UPat(Ops.INDEX, src=(x,)), lambda x: x),
(UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="a"),), allow_any_len=True, name="b"),
lambda a, b: a.src[0].index(*a.src[1:], *b.src[1:])),
(UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="s"), UPat.cvar("i")), allow_any_len=True, name="u"),
lambda s, i, u: _index_stack(s, i, u.src[2:])),
(UPat(ALU, name="x"), _push_stack),
(UPat(Ops.CALL, name="call"), inline_call),
])
_simplified: dict[UOp, UOp] = {}
def simplify(u: UOp) -> UOp:
"""graph_rewrite(u, symbolic), memoized: rewriting is a pure function of a (hash-consed) node."""
if (r := _simplified.get(u)) is not None: return r
if len(_simplified) > 200_000: _simplified.clear()
r = _simplified[u] = graph_rewrite(u, symbolic)
_simplified[r] = r
return r
def run(u: UOp) -> UOp: return graph_rewrite(u, symbolic + evaluate)
def to_python(u: UOp):
"""A fully evaluated program as nested python lists."""
if u.op is Ops.CONST: return u.arg
if u.op is Ops.STACK: return [to_python(s) for s in u.src]
raise ValueError(f"program did not evaluate to a value, stuck at {u.op}")
uopc/uir.py · 178 lines
"""Week 1: the .uir interchange format.
# gemm_256x256.uir
%0 = buffer 65536 : dtype=f32 slot=0
%1 = buffer 65536 : dtype=f32 slot=1
%2 = reshape %0, (256, 1, 256)
%3 = reshape %1, (1, 256, 256)
%4 = mul %2, %3
%5 = reduce %4 : op=add pop=1
Grammar (one statement per line, `#` starts a comment):
line := '%' NAME '=' opname operands? (':' attr*)?
operands := operand (',' operand)*
operand := '%' NAME | literal literal := int | float | true | false | tuple
attr := key '=' value
Integer/float literals in a *source* position become CONSTs (`add %x, 1.0`, `range 10`).
Tuple literals are the arg of movement ops. The last statement is the program.
Two pieces of sugar make hand-written programs readable:
* binary ops broadcast: size-1 axes are EXPANDed, missing leading axes RESHAPEd in
* a bare literal adopts the dtype of the other operand (`mul %f32, 2` means 2.0)
"""
from __future__ import annotations
import ast, math, re
from uopc.uop import UOp, Ops, ParamArg, CallArg, RangeArg, AxisType, BINARY, UOpError
from uopc.dtype import dtypes, AddrSpace
POSITIONAL_ARG = {Ops.BUFFER, Ops.ALLOC, Ops.CONST, Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.SHRINK}
LITERAL_WORDS = {"true": True, "false": False, "inf": math.inf, "-inf": -math.inf, "nan": math.nan}
class ParseError(Exception): pass
def _split_top(s: str) -> list[str]:
"""Split on commas that are not inside parentheses."""
out, depth, cur = [], 0, ""
for ch in s:
if ch == "," and depth == 0: out.append(cur.strip()); cur = ""; continue
depth += (ch == "(") - (ch == ")")
cur += ch
if cur.strip(): out.append(cur.strip())
return out
def _literal(tok: str):
if tok in LITERAL_WORDS: return LITERAL_WORDS[tok]
try: return ast.literal_eval(tok)
except (ValueError, SyntaxError): raise ParseError(f"bad literal {tok!r}") from None
def _attr_value(key: str, v: str):
if key == "dtype": return dtypes.from_name(v)
if key == "op": return Ops[v.upper()]
if key == "kind": return AxisType[v.upper()]
if key == "space": return AddrSpace[v.upper()]
if key == "name": return v
return _literal(v)
def _broadcast(a: UOp, b: UOp) -> tuple[UOp, UOp]:
if a.shape == b.shape: return a, b
n = max(len(a.shape), len(b.shape))
a = a.reshape((1,) * (n - len(a.shape)) + a.shape) if len(a.shape) < n else a
b = b.reshape((1,) * (n - len(b.shape)) + b.shape) if len(b.shape) < n else b
out = []
for x, y in zip(a.shape, b.shape):
if x != y and 1 not in (x, y): raise ParseError(f"can't broadcast {a.shape} with {b.shape}")
out.append(max(x, y))
return (a if a.shape == tuple(out) else a.expand(out)), (b if b.shape == tuple(out) else b.expand(out))
def parse(text: str) -> UOp:
env: dict[str, UOp] = {}
last = None
for lineno, raw in enumerate(text.splitlines(), 1):
line = raw.split("#", 1)[0].strip()
if not line: continue
try: last = _parse_line(line, env)
except (ParseError, UOpError, KeyError, ValueError, TypeError) as e:
raise ParseError(f"line {lineno}: {raw.strip()!r}: {e}") from None
if last is None: raise ParseError("empty program")
return last
def _parse_line(line: str, env: dict[str, UOp]) -> UOp:
lhs, rhs = (p.strip() for p in line.split("=", 1))
if not lhs.startswith("%"): raise ParseError("statement must start with %name =")
body, _, attrs_s = rhs.partition(":")
opname, _, ops_s = body.strip().partition(" ")
op = Ops[opname.upper()]
attrs = {}
for k, v in re.findall(r"(\w+)=(\([^=]*\)|\S+)", attrs_s):
attrs[k] = _attr_value(k, v)
toks = _split_top(ops_s)
lit = None
if op in POSITIONAL_ARG:
if op in (Ops.BUFFER, Ops.ALLOC, Ops.CONST): lit, toks = _literal(toks[0]), toks[1:]
else: lit, toks = _literal(toks[-1]), toks[:-1]
src = []
for t in toks:
if t.startswith("%"):
if t[1:] not in env: raise ParseError(f"undefined {t}")
src.append(env[t[1:]])
else: src.append(_literal(t))
env[lhs[1:]] = u = _build(op, src, lit, attrs)
return u
def _build(op: Ops, src: list, lit, attrs: dict) -> UOp:
# bare literals in source position become CONSTs, typed like their neighbours
# (a WHERE's condition is bool: its branches decide the literal's type)
like = next((s for s in (src[1:] if op is Ops.WHERE else src) if isinstance(s, UOp)), None)
def const(v):
if op in (Ops.INDEX, Ops.RANGE) or like is None or isinstance(v, bool): return UOp.const(v)
return UOp.const(v, like.dtype)
src = [s if isinstance(s, UOp) else const(s) for s in src]
match op:
case Ops.BUFFER | Ops.ALLOC:
return UOp(op, (), ParamArg(attrs.get("slot", 0), attrs.get("dtype", dtypes.f32), (lit,), attrs.get("space")))
case Ops.PARAM: return UOp(op, (), ParamArg(attrs["slot"], attrs.get("dtype", dtypes.f32), tuple(attrs.get("shape", ()))))
case Ops.CONST: return UOp.const(lit, attrs.get("dtype"))
case Ops.CALL: return UOp(op, tuple(src), CallArg(attrs.get("name", "f")))
case Ops.CAST: return src[0].cast(attrs["dtype"])
case Ops.RANGE: return UOp(op, tuple(src), RangeArg(attrs.get("idx", 0), attrs.get("kind", AxisType.LOOP)))
case Ops.REDUCE: return src[0].reduce(attrs.get("op", Ops.ADD), attrs.get("pop", 1))
case Ops.RESHAPE | Ops.EXPAND | Ops.PERMUTE: return UOp(op, (src[0],), tuple(lit))
case Ops.SHRINK: return src[0].shrink(lit)
if op in BINARY:
a, b = _broadcast(*src)
return UOp(op, (a, b))
if op is Ops.WHERE:
c, a, b = src
a, b = _broadcast(a, b); c, a = _broadcast(c, a); c, b = _broadcast(c, b)
return UOp(op, (c, a, b))
return UOp(op, tuple(src))
# ---------------------------------------------------------------------------
# printing
def _fmt_lit(v) -> str:
if isinstance(v, bool): return "true" if v else "false"
if isinstance(v, float):
if math.isnan(v): return "nan"
if math.isinf(v): return "inf" if v > 0 else "-inf"
return repr(v)
if isinstance(v, tuple): return "(" + ", ".join(_fmt_lit(x) for x in v) + ("," if len(v) == 1 else "") + ")"
return str(v)
def to_uir(root: UOp) -> str:
names: dict[UOp, str] = {}
lines = []
for u in root.toposort():
if u.op is Ops.CONST and u is not root: continue # constants are printed inline
names[u] = f"%{len(names)}"
ops = [_fmt_lit(s.arg) if s.op is Ops.CONST else names[s] for s in u.src]
attrs = []
a = u.arg
match u.op:
case Ops.BUFFER | Ops.ALLOC:
ops = [str(a.shape[0])]
attrs = [f"dtype={a.dtype.name}", f"slot={a.slot}"] + ([f"space={a.addrspace.name.lower()}"] if a.addrspace else [])
case Ops.PARAM: attrs = [f"slot={a.slot}", f"dtype={a.dtype.name}"] + ([f"shape={_fmt_lit(a.shape)}"] if a.shape else [])
case Ops.CONST: ops = [_fmt_lit(a)]
case Ops.CALL: attrs = [f"name={a.name}"]
case Ops.CAST: attrs = [f"dtype={a.name}"]
case Ops.RANGE: attrs = [f"idx={a.idx}", f"kind={a.kind.name.lower()}"]
case Ops.REDUCE: attrs = [f"op={a[0].name.lower()}", f"pop={a[1]}"]
case Ops.RESHAPE | Ops.EXPAND | Ops.PERMUTE | Ops.SHRINK: ops.append(_fmt_lit(a))
line = f"{names[u]} = {u.op.name.lower()}"
if ops: line += " " + ", ".join(ops)
if attrs: line += " : " + " ".join(attrs)
lines.append(line)
return "\n".join(lines) + "\n"
uopc/render_c.py · 143 lines
"""Weeks 2-3: rendering a linearized program as a C function.
Each UOp becomes at most one line of C. CONST, INDEX, AFTER and pure ALU ops
used once are rendered inline as expressions; everything else gets a named
variable. BUFFERs become function parameters, sorted by slot.
"""
from __future__ import annotations
import math
from uopc.uop import UOp, Ops, AxisType, ALU
from uopc.dtype import dtypes, DType, AddrSpace
from uopc.linearize import linearize
from uopc.symbolic import vmin
PRELUDE = """#include <math.h>
#include <stdbool.h>
static inline int floordiv_i32(int a, int b) { int q = a / b; return (q * b != a && ((a < 0) != (b < 0))) ? q - 1 : q; }
static inline int floormod_i32(int a, int b) { int r = a % b; return (r != 0 && ((r < 0) != (b < 0))) ? r + b : r; }
"""
def cconst(v, d: DType) -> str:
if d == dtypes.bool: return "1" if v else "0"
if d == dtypes.i32: return str(v) if v != -2**31 else "(-2147483647-1)"
if math.isnan(v): return "NAN"
if math.isinf(v): return "INFINITY" if v > 0 else "-INFINITY"
s = repr(float(v))
return s + "f" if d == dtypes.f32 else f"(({d.ctype}){s}f)"
def ctype(d: DType) -> str: return d.ctype
def _binop(op: Ops, a: str, b: str, ua: UOp, ub: UOp, d: DType) -> str:
match op:
case Ops.ADD: return f"({a}+{b})"
case Ops.MUL: return f"({a}*{b})"
case Ops.CMPLT: return f"({a}<{b})"
case Ops.CMPNE: return f"({a}!={b})"
case Ops.MAX: return f"(({a}>{b})?{a}:{b})"
case Ops.FLOORDIV:
# C division truncates; it equals floor division when both sides are non-negative
if vmin(ua) >= 0 and ub.op is Ops.CONST and ub.arg > 0: return f"({a}/{b})"
return f"floordiv_i32({a},{b})"
case Ops.FLOORMOD:
if vmin(ua) >= 0 and ub.op is Ops.CONST and ub.arg > 0: return f"({a}%{b})"
return f"floormod_i32({a},{b})"
raise NotImplementedError(op)
UNARY_FN = {Ops.EXP2: "exp2", Ops.LOG2: "log2", Ops.SQRT: "sqrt", Ops.SIN: "sin"}
def _unop(op: Ops, a: str, d: DType, arg) -> str:
if op is Ops.CAST: return f"(({arg.ctype})({a}))"
if op is Ops.RECIP: return f"(1.0f/{a})" if d == dtypes.f32 else f"(({d.ctype})1.0f/{a})"
return f"{UNARY_FN[op]}{'f' if d == dtypes.f32 else ''}({a})"
class CRenderer:
"""Render a SINK of stores into one C function. Subclassed for CUDA in week 7."""
kernel_prefix = "void"
buffer_prefix = ""
shared_prefix = ""
barrier = None
restrict = "restrict"
def render_range(self, u: UOp, name: str, end: str) -> list[str]:
pragma = ["#pragma omp parallel for"] if u.arg.kind is AxisType.THREAD else []
return pragma + [f"for (int {name} = 0; {name} < {end}; {name}++) {{"]
def render(self, sink: UOp, name: str = "kernel") -> str:
lin = linearize(sink)
uses: dict[UOp, int] = {}
for u in lin:
for s in u.src: uses[s] = uses.get(s, 0) + 1
r: dict[UOp, str] = {}
lines: list[str] = []
depth = 1
bufs = sorted({u for u in lin if u.op is Ops.BUFFER}, key=lambda u: u.arg.slot)
counter = 0
def var(prefix: str) -> str:
nonlocal counter
counter += 1
return f"{prefix}{counter}"
def line(s: str): lines.append(" " * depth + s)
for u in lin:
op = u.op
if op is Ops.CONST: r[u] = cconst(u.arg, u.dtype)
elif op is Ops.BUFFER: r[u] = f"data{u.arg.slot}"
elif op is Ops.ALLOC:
r[u] = f"{'smem' if u.addrspace is AddrSpace.SHARED else 'acc'}{u.arg.slot}"
pre = self.shared_prefix if u.addrspace is AddrSpace.SHARED else ""
if self._scalar(u): line(f"{ctype(u.dtype)} {r[u]};") # a size-1 accumulator is just a variable
else: line(f"{pre}{ctype(u.dtype)} {r[u]}[{u.arg.shape[0]}];")
elif op is Ops.AFTER: r[u] = r[u.src[0]]
elif op is Ops.INDEX:
base = u.src[0]
while base.op is Ops.AFTER: base = base.src[0]
r[u] = r[u.src[0]] if base.op is Ops.ALLOC and self._scalar(base) else f"{r[u.src[0]]}[{r[u.src[1]]}]"
elif op is Ops.LOAD:
r[u] = var("val")
line(f"{ctype(u.dtype)} {r[u]} = {r[u.src[0]]};")
elif op is Ops.STORE: line(f"{r[u.src[0]]} = {r[u.src[1]]};")
elif op is Ops.RANGE:
r[u] = self.range_name(u)
for l in self.render_range(u, r[u], r[u.src[0]]): line(l)
if self.is_loop(u): depth += 1
elif op is Ops.END:
for rng in u.src[1:]:
if self.is_loop(rng):
depth -= 1
line("}")
elif op is Ops.BARRIER: line(self.barrier + ";")
elif op is Ops.SINK: pass
elif op is Ops.WHERE:
expr = f"({r[u.src[0]]}?{r[u.src[1]]}:{r[u.src[2]]})"
self._assign(u, expr, uses, r, line, var)
elif op in ALU:
if len(u.src) == 1: expr = _unop(op, r[u.src[0]], u.dtype, u.arg)
else: expr = _binop(op, r[u.src[0]], r[u.src[1]], u.src[0], u.src[1], u.dtype)
self._assign(u, expr, uses, r, line, var)
elif op is Ops.WMMA: self.render_wmma(u, r, line)
else: raise NotImplementedError(f"C renderer can't handle {op}")
params = ", ".join(f"{self.buffer_prefix}{ctype(b.dtype)}* {self.restrict} data{b.arg.slot}" for b in bufs)
return self.header(sink, lin) + f"{self.kernel_prefix} {name}({params}) {{\n" + "\n".join(lines) + "\n}\n"
def header(self, sink: UOp, lin: list[UOp]) -> str: return PRELUDE
def _scalar(self, alloc: UOp) -> bool: return alloc.arg.shape == (1,) and alloc.addrspace is AddrSpace.LOCAL
def range_name(self, u: UOp) -> str: return f"ridx{u.arg.idx}"
def is_loop(self, u: UOp) -> bool:
"""On a CPU every range is a for loop -- GPU ranges too, which makes this an emulator for GPU kernels."""
return True
def _assign(self, u, expr, uses, r, line, var):
if uses.get(u, 0) <= 1: r[u] = expr # used once: inline it
else:
r[u] = var("alu")
line(f"{ctype(u.dtype)} {r[u]} = {expr};")
def render_wmma(self, u, r, line): raise NotImplementedError("WMMA needs the CUDA renderer")
def render_c(sink: UOp, name: str = "kernel") -> str: return CRenderer().render(sink, name)
uopc/runtime.py · 70 lines
"""Week 2: the runtime. Compile C with the system compiler, load it, call it.
Compiled shared objects are cached on disk by a hash of (source, flags), so a
program is only compiled once per machine.
"""
from __future__ import annotations
import ctypes, hashlib, os, subprocess, tempfile
from pathlib import Path
import numpy as np
from uopc.dtype import DType, dtypes
CC = os.environ.get("CC", "gcc")
# -fassociative-math lets gcc vectorize reductions (it reorders float adds, like every ML compiler does).
# NaN/inf semantics are kept. Set CFLAGS to override, e.g. CFLAGS="-O2 -fopenmp" for strict IEEE order.
CFLAGS = os.environ.get("CFLAGS", "-O3 -march=native -fopenmp -ffp-contract=fast -fassociative-math "
"-fno-signed-zeros -fno-trapping-math").split()
CACHE = Path(os.environ.get("UOPC_CACHE", Path.home() / ".cache" / "uopc"))
NP = {dtypes.f32: np.float32, dtypes.i32: np.int32, dtypes.bool: np.bool_, dtypes.f16: np.float16}
def np_dtype(d: DType): return NP[d]
class CompileError(Exception): pass
def compile_c(src: str, flags: list[str] | None = None) -> Path:
flags = CFLAGS if flags is None else flags
key = hashlib.sha256((src + "\0" + " ".join([CC, *flags])).encode()).hexdigest()[:24]
CACHE.mkdir(parents=True, exist_ok=True)
so = CACHE / f"{key}.so"
if so.exists(): return so
with tempfile.TemporaryDirectory(dir=CACHE) as td: # same filesystem, so os.replace is atomic
c = Path(td) / "k.c"
c.write_text(src)
tmp = Path(td) / "k.so"
p = subprocess.run([CC, *flags, "-shared", "-fPIC", "-o", str(tmp), str(c), "-lm"], capture_output=True, text=True)
if p.returncode != 0:
numbered = "\n".join(f"{i+1:4d} {l}" for i, l in enumerate(src.splitlines()))
raise CompileError(f"{CC} failed:\n{p.stderr}\n--- source ---\n{numbered}")
os.replace(tmp, so) # atomic: concurrent compiles of the same kernel are fine
return so
class Program:
"""A compiled C function taking one pointer per buffer slot."""
def __init__(self, src: str, name: str = "kernel", flags: list[str] | None = None):
self.src, self.name = src, name
self.lib = ctypes.CDLL(str(compile_c(src, flags)))
self.fn = getattr(self.lib, name)
self.fn.restype = None
def __call__(self, *bufs: np.ndarray):
for b in bufs:
if not b.flags.c_contiguous: raise ValueError("buffers must be contiguous")
self.fn(*(ctypes.c_void_p(b.ctypes.data) for b in bufs))
def run_c(sink, buffers: dict[int, np.ndarray], name: str = "kernel") -> str:
"""Render, compile and run a SINK of stores. buffers maps slot -> numpy array. Returns the C source."""
from uopc.render_c import render_c
from uopc.uop import Ops
src = render_c(sink, name)
bufs = sorted({u for u in sink.toposort() if u.op is Ops.BUFFER}, key=lambda u: u.arg.slot)
for b in bufs:
a = buffers[b.arg.slot]
if a.dtype != np_dtype(b.dtype) or a.size < b.arg.shape[0]:
raise TypeError(f"buffer slot {b.arg.slot}: kernel wants {b.arg.shape[0]} x {b.dtype.name}, got {a.size} x {a.dtype}")
Program(src, name)(*(buffers[b.arg.slot] for b in bufs))
return src
uopc/linearize.py · 128 lines
"""Weeks 2-3: turning a graph into a straight-line list of instructions with loops.
A graph has no order; C does. Linearizing means picking a topological order
that also respects *scope*: the set of loops a node must sit inside.
scope(RANGE r) = scope(end) | {r}
scope(END x, *rs) = scope(x) - rs
scope(anything else) = union of scope(srcs)
Loops nest in increasing RangeArg.idx order. If two ranges a < b ever appear in
one scope, b is nested inside a, so *every* node that uses b must also sit inside a.
Closing scopes under "encloses" turns each scope into a path in the loop tree:
a sorted tuple that is exactly the stack of open loops when the node runs.
Then we list-schedule:
1. emit a ready node whose scope is exactly the open loops
2. else open a loop that pending nodes need next (sibling loops whose END
another sibling's nodes read must go first)
3. else close the innermost loop
ALLOCs are declared in the deepest scope shared by all their users (so each CPU
thread gets its own private accumulator).
"""
from __future__ import annotations
import heapq
from uopc.uop import UOp, Ops
from uopc.dtype import AddrSpace
class LinearizeError(Exception): pass
def scopes(sink: UOp) -> dict[UOp, frozenset[UOp]]:
sc: dict[UOp, frozenset[UOp]] = {}
for u in sink.toposort():
s = frozenset().union(*(sc[x] for x in u.src))
if u.op is Ops.RANGE: s = s | {u}
if u.op is Ops.END: s = s - set(u.src[1:])
sc[u] = s
return sc
def _key(s) -> tuple[UOp, ...]: return tuple(sorted(s, key=lambda r: r.arg))
def loop_targets(topo: list[UOp], sink: UOp) -> dict[UOp, tuple[UOp, ...]]:
sc = scopes(sink)
# enclosing loops of each range: every smaller range it ever shares a scope with
enc: dict[UOp, set[UOp]] = {}
for s in set(sc.values()):
t = _key(s)
for j, r in enumerate(t): enc.setdefault(r, set()).update(t[:j])
def close(s):
out, todo = set(s), list(s)
while todo:
for a in enc.get(todo.pop(), ()):
if a not in out: out.add(a); todo.append(a)
return _key(out)
target = {u: close(sc[u]) for u in topo}
for u in topo:
if u.op is Ops.RANGE: target[u] = close(sc[u] - {u} | enc.get(u, set()))
return target
def linearize(sink: UOp) -> list[UOp]:
topo = sink.toposort()
order = {u: i for i, u in enumerate(topo)}
target = loop_targets(topo, sink)
users: dict[UOp, list[UOp]] = {u: [] for u in topo}
for u in topo:
for s in set(u.src): users[s].append(u)
for u in topo:
if u.op is Ops.ALLOC and u.addrspace is AddrSpace.SHARED: target[u] = () # shared memory: whole block, top of kernel
elif u.op is Ops.ALLOC and users[u]:
ts = [target[x] for x in users[u]]
n = 0
while all(len(t) > n and t[n] is ts[0][n] for t in ts): n += 1
target[u] = ts[0][:n]
waiting = {u: len(set(u.src)) for u in topo}
ready: set[UOp] = set() # ready RANGE / END nodes
heaps: dict[tuple, list[tuple[int, UOp]]] = {} # ready everything-else, bucketed by scope
pending, stack, out = set(topo), [], []
def make_ready(u: UOp):
if u.op in (Ops.RANGE, Ops.END): ready.add(u)
else: heapq.heappush(heaps.setdefault(target[u], []), (order[u], u))
def emit(u: UOp):
out.append(u); pending.discard(u); ready.discard(u)
for x in users[u]:
waiting[x] -= 1
if waiting[x] == 0: make_ready(x)
for u in topo:
if waiting[u] == 0: make_ready(u)
while pending:
cur = tuple(stack)
# 1. anything that belongs exactly here
if heaps.get(cur):
emit(heapq.heappop(heaps[cur])[1]); continue
# 2. open a loop that pending nodes need next
need = {u: target[u][len(cur)] for u in pending
if u.op is not Ops.RANGE and len(target[u]) > len(cur) and target[u][:len(cur)] == cur}
cands = {r for r in need.values() if r in ready and target[r] == cur}
if cands:
# if nodes needing loop A read (through an END) the results of sibling loop B, open B first
ended: dict[UOp, set[UOp]] = {}
blocked: set[UOp] = set()
for u in topo:
if u not in pending: continue
e = set(u.src[1:]) & cands if u.op is Ops.END else set()
for x in u.src:
if x in pending: e |= ended.get(x, set())
ended[u] = e
if u in need and e - {need[u]}: blocked.add(need[u])
r = min(cands - blocked or cands, key=lambda r: r.arg)
emit(r); stack.append(r); continue
# 3. close the innermost loop(s)
ends = [u for u in ready if u.op is Ops.END and target[u] == cur[:len(cur) - (len(u.src) - 1)]
and set(u.src[1:]) == set(cur[len(cur) - (len(u.src) - 1):])]
if ends:
e = min(ends, key=order.__getitem__)
emit(e); del stack[len(stack) - (len(e.src) - 1):]; continue
stuck = sorted(pending, key=order.__getitem__)[:5]
raise LinearizeError(f"can't schedule; open loops {[r.arg.idx for r in cur]}, stuck at "
f"{[(u.op, [r.arg.idx for r in target[u]]) for u in stuck]}")
if stack: raise LinearizeError(f"RANGE without END: loops {[r.arg.idx for r in stack]} are never closed")
return out
uopc/rangeify.py · 201 lines
"""Week 4: rangeify. From shapes to loops.
A shaped graph says *what* each element is. Rangeify says *how to compute one
element*: give every output axis a RANGE, then ask each node "what is your value
at index (i0, i1, ...)?" and push the index down the graph:
BUFFER value[i] -> LOAD(INDEX(buf, i))
ALU f(a, b)[idx] -> f(a[idx], b[idx])
RESHAPE x.reshape(s)[idx] -> x[unravel(ravel(idx, s), x.shape)]
EXPAND x.expand(s)[idx] -> x[0 where x had size 1]
PERMUTE x.permute(p)[idx] -> x[idx permuted back]
SHRINK x.shrink(b)[idx] -> x[idx + begin]
STACK stack(a, b)[i, *rest] -> where(i < 1, a[rest], b[rest])
REDUCE sum over the last n -> new RANGEs, an accumulator, a loop
After rangeify no op has a shape: everything is a scalar in some loop. Symbolic
simplification then cleans up the index arithmetic (ravel . unravel == id).
"""
from __future__ import annotations
from dataclasses import dataclass, field
from uopc.uop import UOp, Ops, ParamArg, AxisType, ALU, MOVEMENT, sink, range_
from uopc.dtype import dtypes, DType
from uopc.symbolic import simplify
def ravel(idx: tuple[UOp, ...], shape: tuple[int, ...]) -> UOp:
flat = UOp.const(0)
for i, s in zip(idx, shape): flat = flat * s + i
return flat
def unravel(flat: UOp, shape: tuple[int, ...]) -> tuple[UOp, ...]:
out, stride = [], 1
for s in reversed(shape):
out.append((flat // stride) % s if s != 1 else UOp.const(0))
stride *= s
return tuple(reversed(out))
def identity(op: Ops, d: DType):
if op is Ops.ADD: return 0.0 if d.is_float else 0
return -float("inf") if d.is_float else dtypes.min(d)
@dataclass
class Opts:
"""Week 6: how to map a kernel onto the machine. Chosen by uopc.opt, applied during rangeify."""
upcast: dict[int, int] = field(default_factory=dict) # output axis -> unroll factor (register tile)
thread: int | None = None # output axis run in parallel (OpenMP)
unroll: int = 1 # unroll factor of the innermost reduce axis
tile: dict[int, int] = field(default_factory=dict) # output axis -> cache block size (week 6, loop tiling)
gpu: bool = False # weeks 7-8: output axes become GLOBAL (blockIdx) ranges
local: dict[int, int] = field(default_factory=dict) # output axis -> threads per block on that axis (LOCAL)
@dataclass
class Ctx:
next_range: int = 0
next_alloc: int = 0
memo: dict = field(default_factory=dict)
reduce_ranges: dict = field(default_factory=dict)
opts: Opts = field(default_factory=Opts)
def rng(self, size: int, kind: AxisType = AxisType.LOOP) -> UOp:
self.next_range += 1
return range_(size, self.next_range - 1, kind)
def alloc(self, d: DType, n: int = 1) -> UOp:
self.next_alloc += 1
return UOp(Ops.ALLOC, (), ParamArg(self.next_alloc - 1, d, (n,)))
def ranges_in(*us: UOp) -> list[UOp]:
return sorted({r for u in us for r in u.toposort() if r.op is Ops.RANGE}, key=lambda r: r.arg)
def rangeify(u: UOp, idx: tuple[UOp, ...], ctx: Ctx) -> UOp:
"""The scalar value of u at index idx."""
key = (u, idx)
if (r := ctx.memo.get(key)) is not None: return r
assert len(idx) == len(u.shape), f"{u.op}: {len(idx)} indices for shape {u.shape}"
op = u.op
if op is Ops.CONST: r = u
elif op is Ops.BUFFER: r = u.index(idx[0]).load()
elif op is Ops.STAGE: r = rangeify(u.src[0], idx, ctx) # no kernelizer: compute it inline
elif op in ALU: r = u.replace(src=tuple(rangeify(s, idx, ctx) for s in u.src))
elif op is Ops.INDEX:
iv = tuple(rangeify(i, (), ctx) for i in u.src[1:])
r = rangeify(u.src[0], iv + idx, ctx)
elif op is Ops.STACK:
i, rest = idx[0], idx[1:]
r = rangeify(u.src[-1], rest, ctx)
for k in range(len(u.src) - 2, -1, -1):
r = (i < k + 1).where(rangeify(u.src[k], rest, ctx), r)
elif op is Ops.RESHAPE:
# simplify now: ((i*N + j)//N) % N mentions j but *is* i. Dependence analysis
# (which loops a reduce must sit in) and the memo both need the truth.
r = rangeify(u.src[0], tuple(simplify(i) for i in unravel(ravel(idx, u.shape), u.src[0].shape)), ctx)
elif op is Ops.EXPAND:
r = rangeify(u.src[0], tuple(UOp.const(0) if s == 1 and n != 1 else i for i, s, n in zip(idx, u.src[0].shape, u.shape)), ctx)
elif op is Ops.PERMUTE:
inv = [0] * len(idx)
for k, p in enumerate(u.arg): inv[p] = idx[k]
r = rangeify(u.src[0], tuple(inv), ctx)
elif op is Ops.SHRINK: r = rangeify(u.src[0], tuple(i + b if b else i for i, (b, _) in zip(idx, u.arg)), ctx)
elif op is Ops.REDUCE: r = _reduce(u, idx, ctx)
else: raise NotImplementedError(f"rangeify: {op}")
ctx.memo[key] = r
return r
def _axis_parts(ctx: Ctx, size: int, factor: int, kind: AxisType) -> tuple[UOp, int]:
"""An axis of `size` as (outer RANGE or CONST 0, unroll factor): index = outer * factor + c."""
if factor > 1 and size % factor == 0:
return (ctx.rng(size // factor, kind) if size > factor else UOp.const(0)), factor
return (ctx.rng(size, kind) if size > 1 else UOp.const(0)), 1
def _combos(parts: list[tuple[UOp, int]]) -> list[tuple[UOp, ...]]:
"""Every index tuple obtained by fixing the unrolled part of each axis to a constant.
parts[k] = (base, f): the index of axis k is base * f + c for c in 0..f-1."""
out = [()]
for b, f in parts:
known0 = b.op is Ops.CONST and b.arg == 0
out = [c + (UOp.const(k) if known0 else (b * f + k if f > 1 else b),) for c in out for k in range(f)]
return out
def _reduce(u: UOp, idx: tuple[UOp, ...], ctx: Ctx) -> UOp:
rop, n = u.arg
x = u.src[0]
rshape = x.shape[len(x.shape) - n:]
# unrolled copies of the same reduce (they differ only by constants in idx) share one loop
rkey = (u, frozenset(ranges_in(*idx)))
if (parts := ctx.reduce_ranges.get(rkey)) is None:
parts = [_axis_parts(ctx, s, ctx.opts.unroll if k == n - 1 else 1, AxisType.REDUCE) for k, s in enumerate(rshape)]
ctx.reduce_ranges[rkey] = parts
real = [r for r, _ in parts if r.op is Ops.RANGE]
vals = [rangeify(x, idx + c, ctx) for c in _combos(parts)]
while len(vals) > 1: # unrolled reduce axis: a balanced tree
vals = [UOp(rop, (vals[i], vals[i + 1])) if i + 1 < len(vals) else vals[i] for i in range(0, len(vals), 2)]
if not real: return vals[0]
acc = ctx.alloc(u.dtype)
# the accumulator lives inside the loops the output index uses...
acc0 = acc.after(*ranges_in(*idx))
init = acc0.index(0).store(UOp.const(identity(rop, u.dtype), u.dtype))
# ...is re-read every iteration of the reduce loops...
accl = acc0.after(init, *real)
upd = accl.index(0).store(UOp(rop, (accl.index(0).load(), vals[0])))
# ...and is read once more after the loops end
return acc0.after(upd.end(*real)).index(0).load()
def merge_ends(root: UOp) -> UOp:
"""Unrolled copies each END the same shared loop. Merge them into one END over a SINK."""
while True:
groups: dict[tuple, list[UOp]] = {}
for u in root.toposort():
if u.op is Ops.END: groups.setdefault(u.src[1:], []).append(u)
dup = next((g for g in groups.values() if len(g) > 1), None)
if dup is None: return root
merged = UOp(Ops.END, (UOp(Ops.SINK, tuple(e.src[0] for e in dup)), *dup[0].src[1:]))
root = root.substitute({e: merged for e in dup})
def _gpu_axis(ctx: Ctx, size: int, local: int, factor: int) -> tuple[UOp, int]:
"""size = blocks * local * factor: index = (gidx * local + lidx) * factor + c."""
assert size % (local * factor) == 0, f"axis {size} not divisible by local {local} * upcast {factor}"
g = ctx.rng(size // (local * factor), AxisType.GLOBAL) if size // (local * factor) > 1 else UOp.const(0)
if local == 1: return g, factor
return g * local + ctx.rng(local, AxisType.LOCAL), factor
def lower(root: UOp, out: UOp, ctx: Ctx | None = None) -> UOp:
"""Lower one shaped value into a SINK that stores it, element by element, into BUFFER `out`."""
ctx = ctx or Ctx()
o = ctx.opts
if o.gpu:
# each output axis gets a GLOBAL range and (optionally) a LOCAL one; neither becomes a loop on a GPU
gl = [_gpu_axis(ctx, s, o.local.get(k, 1), o.upcast.get(k, 1)) for k, s in enumerate(root.shape)]
plist = gl
else:
# the THREAD axis gets the lowest range index, so it is the outermost loop
order = ([o.thread] if o.thread is not None else []) + [k for k in range(len(root.shape)) if k != o.thread]
tiled = {k: t for k, t in o.tile.items() if 1 < t < root.shape[k] and root.shape[k] % t == 0 and t % o.upcast.get(k, 1) == 0}
parts: dict[int, tuple[UOp, int]] = {}
# tiled axes: first every block loop (outer), then the loops inside a block. A block of outputs reuses
# the same rows of each operand, so they stay in cache while the block is computed
outer = {k: ctx.rng(root.shape[k] // t, AxisType.THREAD if k == o.thread else AxisType.LOOP) for k, t in tiled.items()}
for k in order:
f = o.upcast.get(k, 1)
if k in tiled:
inner = tiled[k] // f
parts[k] = (outer[k] * inner + ctx.rng(inner) if inner > 1 else outer[k], f)
else:
parts[k] = _axis_parts(ctx, root.shape[k], f, AxisType.THREAD if k == o.thread else AxisType.LOOP)
plist = [parts[k] for k in range(len(root.shape))]
stores = []
for idx in _combos(plist):
val = rangeify(root, idx, ctx)
stores.append(out.index(ravel(idx, root.shape) if root.shape else UOp.const(0)).store(val))
loops = ranges_in(*(b for b, _ in plist))
body = stores[0] if len(stores) == 1 else UOp(Ops.SINK, tuple(stores))
return simplify(merge_ends(sink(body.end(*loops))))
uopc/composite.py · 92 lines
"""Week 4: composite movement ops, built only from RESHAPE/EXPAND/PERMUTE/SHRINK/STACK.
PAD and FLIP are not primitives in this course. Building them out of primitives is
a good exercise in thinking about index maps, and rangeify + symbolic turn the
result back into tight code (PAD becomes a bounds-checked WHERE).
"""
from __future__ import annotations
from uopc.uop import UOp, Ops, prod
def zeros_like(x: UOp) -> UOp:
z = UOp.const(0, x.dtype)
return z.reshape((1,) * len(x.shape)).expand(x.shape) if x.shape else z
def move_axis(x: UOp, src: int, dst: int) -> UOp:
order = [i for i in range(len(x.shape)) if i != src]
order.insert(dst, src)
return x.permute(order) if order != list(range(len(x.shape))) else x
def concat2(a: UOp, b: UOp, axis: int) -> UOp:
"""Concatenate two same-shaped tensors along `axis`: STACK, move the new axis next to `axis`, merge."""
assert a.shape == b.shape
s = UOp(Ops.STACK, (a, b)) # (2, *shape)
s = move_axis(s, 0, axis) # (..., 2, n, ...)
shape = list(a.shape)
shape[axis] *= 2
return s.reshape(tuple(shape))
def _shrink_axis(x: UOp, axis: int, b: int, e: int) -> UOp:
arg = [(0, s) for s in x.shape]
arg[axis] = (b, e)
return x.shrink(arg) if (b, e) != (0, x.shape[axis]) else x
def pad_axis(x: UOp, axis: int, before: int, after: int) -> UOp:
"""Zero-pad one axis by doubling with zeros and shrinking back."""
while after > 0:
n = x.shape[axis]
x = _shrink_axis(concat2(x, zeros_like(x), axis), axis, 0, n + min(after, n))
after -= min(after, n)
while before > 0:
n = x.shape[axis]
x = _shrink_axis(concat2(zeros_like(x), x, axis), axis, n - min(before, n), 2 * n)
before -= min(before, n)
return x
def pad(x: UOp, arg: tuple[tuple[int, int], ...]) -> UOp:
for axis, (b, a) in enumerate(arg): x = pad_axis(x, axis, b, a)
return x
def anti_eye(n: int, dtype) -> UOp:
"""M[i, j] = (i + j == n - 1), from one-hot + the expand/reshape/shrink "diagonal trick".
u is one-hot at n-1, length 2n-1. Expand u to (n+1, 2n-1), flatten, keep n*2n
elements, view as (n, 2n): element (i, j) is u[(2n*i + j) mod (2n-1)] = u[i + j].
"""
if n == 1: return UOp.const(1, dtype).reshape((1, 1))
one = UOp.const(1, dtype).reshape((1,))
u = pad_axis(one, 0, n - 1, n - 1) # (2n-1,)
m = u.reshape((1, 2 * n - 1)).expand((n + 1, 2 * n - 1)).reshape(((n + 1) * (2 * n - 1),))
m = m.shrink(((0, n * 2 * n),)).reshape((n, 2 * n))
return m.shrink(((0, n), (0, n)))
def flip_axis(x: UOp, axis: int) -> UOp:
"""Reverse one axis: y[i] = sum_j x[j] * [i + j == n-1]. O(n^2), but only primitives."""
n = x.shape[axis]
if n == 1: return x
xl = move_axis(x, axis, len(x.shape) - 1) # (..., n)
lead = xl.shape[:-1]
xe = xl.reshape(lead + (1, n)).expand(lead + (n, n))
m = anti_eye(n, x.dtype).reshape((1,) * len(lead) + (n, n)).expand(lead + (n, n))
y = (xe * m).reduce(Ops.ADD, 1) # (..., n)
return move_axis(y, len(x.shape) - 1, axis)
def sum_to(g: UOp, shape: tuple[int, ...]) -> UOp:
"""Reverse of EXPAND: sum g over the axes where `shape` has 1 and g doesn't."""
axes = [i for i, (s, t) in enumerate(zip(g.shape, shape)) if s != t]
if not axes: return g
keep = [i for i in range(len(shape)) if i not in axes]
return g.permute(keep + axes).reduce(Ops.ADD, len(axes)).reshape(shape)
def is_contiguous_buffer(u: UOp) -> bool:
while u.op is Ops.RESHAPE: u = u.src[0]
return u.op is Ops.BUFFER and prod(u.shape) == u.arg.shape[0]
uopc/kernelize.py · 220 lines
"""Week 5: kernel fission and fusion.
Rangeify can lower *any* shape graph into one giant loop nest, but that is often
terrible: a REDUCE consumed through an EXPAND would be recomputed for every
element it is broadcast to ((A@B)@C would recompute A@B for every output). So we cut the graph into
kernels at *materialization points*:
* every root (the tensors the user asked for)
* every STAGE (the user, or a rule, asked for a temp buffer)
* every REDUCE, unless fuse_reduce is on (then reductions nest, see flash attention)
Everything between materialization points is fused into one kernel. A kernel
reads BUFFERs and writes exactly one. The schedule is the list of kernels in
dependency order. Whole schedules are cached on the *structure* of the graph,
so a training step only pays for Python-side compilation once.
"""
from __future__ import annotations
import os, time
from dataclasses import dataclass
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, prod, _argkey
from uopc.dtype import DType
from uopc.rangeify import lower
from uopc.composite import is_contiguous_buffer
from uopc.runtime import Program, np_dtype
from uopc.rewrite import UPat, PatternMatcher, graph_rewrite
DEBUG = int(os.environ.get("DEBUG", "0"))
# ---------------------------------------------------------------------------
# buffers: a BUFFER UOp owns a numpy array. When the UOp dies, so does the memory.
_next_slot = 0
def new_buffer(dtype: DType, size: int, data: np.ndarray | None = None) -> UOp:
global _next_slot
_next_slot += 1
u = UOp(Ops.BUFFER, (), ParamArg(_next_slot, dtype, (max(size, 1),)))
if data is not None: u.__dict__["np"] = np.ascontiguousarray(data.reshape(-1), dtype=np_dtype(dtype))
return u
def buffer_data(u: UOp) -> np.ndarray:
if "np" not in u.__dict__: u.__dict__["np"] = np.zeros(u.arg.shape[0], np_dtype(u.dtype))
return u.__dict__["np"]
def base_buffer(u: UOp) -> UOp:
while u.op is Ops.RESHAPE: u = u.src[0]
return u
# ---------------------------------------------------------------------------
# the schedule
@dataclass
class Kernel:
name: str
src: str
refs: list[tuple[str, int]] # one per local slot: ("in", i) | ("tmp", j) | ("out", r)
program: Program | None = None
flops: int = 0
@dataclass
class Plan:
kernels: list[Kernel]
temps: list[tuple[DType, int]]
outs: list[tuple[str, int] | tuple[DType, int]] # per root: ("in", i) alias, or (dtype, size) fresh
temp_arrays: list | None = None # allocated on first run, then reused: no page faults on every step
def _materialize_points(roots: list[UOp], fuse_reduce: bool) -> list[UOp]:
topo = UOp(Ops.SINK, tuple(roots)).toposort()
m = []
for u in topo:
if u.op is Ops.SINK: continue
if u.op is Ops.STAGE or (u.op is Ops.REDUCE and not fuse_reduce) or (u in roots and not is_contiguous_buffer(u)):
m.append(u)
return m
# Tensor.detach() wraps a node in STAGE(arg="detach"): autodiff stops there, execution ignores it
strip_detach = PatternMatcher([(UPat(Ops.STAGE, arg="detach", src=(UPat.var("x"),)), lambda x: x)])
def stage_broadcast_compute(roots: list[UOp], factor: int = 4) -> list[UOp]:
"""Fusing is only free if nothing gets recomputed. A computed tensor that is EXPANDed by
`factor` or more (e.g. a LayerNorm output feeding a matmul) would be recomputed for every
copy, so we materialize it first. This is the fusion/recompute trade-off, as one rule."""
from uopc.uop import MOVEMENT, ALU
todo: set[UOp] = set()
for u in UOp(Ops.SINK, tuple(roots)).toposort():
if u.op is not Ops.EXPAND or prod(u.shape) < factor * prod(u.src[0].shape): continue
base = u.src[0]
while base.op in MOVEMENT: base = base.src[0]
if base.op in ALU and base.op is not Ops.CAST and base.shape and _fused_work(base) >= 4096: todo.add(base)
return UOp.wrap(roots, todo, lambda n: UOp(Ops.STAGE, (n,))) if todo else roots
def _fused_work(u: UOp) -> int:
"""ALU ops * elements that recomputing u would cost (stopping at materialized nodes)."""
from uopc.uop import ALU
n, seen, stack = 0, set(), [u]
while stack:
x = stack.pop()
if x in seen or x.op in (Ops.BUFFER, Ops.REDUCE, Ops.STAGE): continue
seen.add(x)
if x.op in ALU: n += 1
stack.extend(x.src)
return n * prod(u.shape)
def _kernel_name(root: UOp) -> str:
red = any(u.op is Ops.REDUCE for u in root.toposort())
return ("r_" if red else "e_") + "_".join(map(str, root.shape or (1,)))
def make_plan(roots: list[UOp], fuse_reduce: bool = False, opt=None) -> tuple[Plan, list[UOp]]:
"""Kernelize. Returns the plan and the canonical list of input BUFFERs it refers to as ("in", i)."""
inputs = [u for u in UOp(Ops.SINK, tuple(roots)).toposort() if u.op is Ops.BUFFER]
in_idx = {u: i for i, u in enumerate(inputs)}
roots = [graph_rewrite(r, strip_detach) for r in roots]
# fuse_reduce asks for "never write the big intermediate", so don't stage broadcasts either
if not fuse_reduce: roots = stage_broadcast_compute(roots)
if opt is not None: roots = opt.layout(roots)
points = _materialize_points(roots, fuse_reduce)
# every materialization point gets a placeholder buffer; consumers read it through a reshape
ref: dict[UOp, tuple[str, int]] = {}
temps: list[tuple[DType, int]] = []
outs: list = []
place: dict[UOp, UOp] = {}
for k, u in enumerate(points):
if u in roots: continue
ref[u] = ("tmp", len(temps))
temps.append((u.dtype, max(1, prod(u.shape))))
for r, u in enumerate(roots):
if is_contiguous_buffer(u): outs.append(("in", in_idx[base_buffer(u)]))
else:
ref[u] = ("out", r)
outs.append((u.dtype, max(1, prod(u.shape))))
for u in points:
ph = UOp(Ops.BUFFER, (), ParamArg(-1 - len(place), u.dtype, (max(1, prod(u.shape)),))) # negative slots: placeholders
place[u] = ph.reshape(u.shape) if u.shape != ph.shape else ph
ph_ref = {base_buffer(place[u]): ref[u] for u in points}
kernels = []
for u in points:
body = u.src[0] if u.op is Ops.STAGE else u
body = body.substitute({p: v for p, v in place.items() if p is not u})
# renumber buffers 0..n-1 locally (output last) so identical kernels render identical C
bufs = [b for b in body.toposort() if b.op is Ops.BUFFER]
local = {b: UOp(Ops.BUFFER, (), ParamArg(i, b.dtype, b.arg.shape)) for i, b in enumerate(bufs)}
out = UOp(Ops.BUFFER, (), ParamArg(len(bufs), u.dtype, (max(1, prod(u.shape)),)))
body = body.substitute(local)
kname = _kernel_name(body)
if opt is not None: ast = opt.lower(body, out)
else: ast = lower(body, out)
from uopc.render_c import render_c
src = render_c(ast, kname)
refs = [ph_ref[b] if b.arg.slot < 0 else ("in", in_idx[b]) for b in bufs] + [ref[u]]
# simplification can delete a read (x*0 -> 0): the C signature then has fewer parameters,
# so bind exactly the slots that survived, or every later argument shifts by one
alive = sorted({b.arg.slot for b in ast.toposort() if b.op is Ops.BUFFER})
refs = [refs[s] for s in alive]
kernels.append(Kernel(kname, src, refs, flops=_flops(body)))
return Plan(kernels, temps, outs), inputs
def _flops(u: UOp) -> int:
"""Rough FLOP count of a shaped kernel body: one per ALU element computed."""
from uopc.uop import ALU
return sum(prod(x.shape or (1,)) for x in u.toposort() if x.op in ALU) + \
sum(prod(x.src[0].shape) for x in u.toposort() if x.op is Ops.REDUCE)
# ---------------------------------------------------------------------------
# structural cache ("JIT"): same graph shape -> same plan, whatever the buffers hold
_plans: dict[tuple, Plan] = {}
def _structure_key(roots: list[UOp]) -> tuple:
ids: dict[UOp, int] = {}
key = []
for u in UOp(Ops.SINK, tuple(roots)).toposort():
ids[u] = len(ids)
arg = ("buf", u.dtype.name, u.arg.shape) if u.op is Ops.BUFFER else _argkey(u.arg)
key.append((u.op, arg, tuple(ids[s] for s in u.src)))
return tuple(key)
def realize(roots: list[UOp], fuse_reduce: bool = False, opt=None) -> list[UOp]:
"""Compute each root. Returns, per root, a BUFFER UOp holding the result."""
uniq = list(dict.fromkeys(roots))
t0 = time.perf_counter()
key = (_structure_key(uniq), fuse_reduce, opt)
inputs = [u for u in UOp(Ops.SINK, tuple(uniq)).toposort() if u.op is Ops.BUFFER]
if (plan := _plans.get(key)) is None:
plan, _ = make_plan(uniq, fuse_reduce, opt)
for k in plan.kernels:
k.program = Program(k.src, k.name)
if DEBUG >= 4: print(k.src)
_plans[key] = plan
t1 = time.perf_counter()
if plan.temp_arrays is None: plan.temp_arrays = [np.empty(n, np_dtype(d)) for d, n in plan.temps]
temps = plan.temp_arrays
outs = [inputs[o[1]] if o[0] == "in" else new_buffer(o[0], o[1]) for o in plan.outs]
def get(rf):
kind, i = rf
return buffer_data(inputs[i]) if kind == "in" else temps[i] if kind == "tmp" else buffer_data(outs[i])
for k in plan.kernels:
ts = time.perf_counter()
k.program(*(get(r) for r in k.refs))
if DEBUG >= 2:
dt = time.perf_counter() - ts
extra = ""
if DEBUG >= 3: # run again: the second run has warm caches and no first-touch page faults
ts = time.perf_counter(); k.program(*(get(r) for r in k.refs)); extra = f" (again: {(time.perf_counter()-ts)*1e3:.3f} ms)"
print(f" {k.name:30s} {dt*1e3:8.3f} ms {k.flops/max(dt, 1e-9)/1e9:7.2f} GFLOP/s{extra}")
if DEBUG >= 1: print(f"realize: {len(plan.kernels)} kernels, schedule {(t1-t0)*1e3:.1f} ms, run {(time.perf_counter()-t1)*1e3:.1f} ms")
by_root = dict(zip(uniq, outs))
return [by_root[r] for r in roots]
uopc/autodiff.py · 82 lines
"""Week 5: reverse-mode automatic differentiation on shape graphs.
Autodiff is just another graph transformation: given the graph for `loss` and
some target nodes, build the graph for d(loss)/d(target). It works on pure
shape graphs only -- no STORE, CALL, AFTER, RANGE or END -- which is why we do
it *before* rangeify.
Each rule maps (node, grad of node) -> grad of each source (None = no gradient).
"""
from __future__ import annotations
import math
from uopc.uop import UOp, Ops
from uopc.composite import sum_to, pad, zeros_like
LN2 = math.log(2)
def _full(like: UOp, v: float) -> UOp:
c = UOp.const(v, like.dtype)
return c.reshape((1,) * len(like.shape)).expand(like.shape) if like.shape else c
def _rule(u: UOp, g: UOp) -> tuple[UOp | None, ...]:
s = u.src
match u.op:
case Ops.ADD: return (g, g)
case Ops.MUL: return (g * s[1], g * s[0])
case Ops.MAX:
a, b = s
half = g * _full(g, 0.5)
zero = zeros_like(g)
return ((b < a).where(g, (a < b).where(zero, half)), (a < b).where(g, (b < a).where(zero, half)))
case Ops.RECIP: return (-(g * u * u),)
case Ops.EXP2: return (g * u * _full(u, LN2),)
case Ops.LOG2: return (g * s[0].recip() * _full(u, 1 / LN2),)
case Ops.SQRT: return (g * u.recip() * _full(u, 0.5),)
case Ops.SIN: return (g * UOp(Ops.SIN, (s[0] + _full(u, math.pi / 2),)),)
case Ops.CAST: return (g.cast(s[0].dtype) if s[0].dtype.is_float else None,)
case Ops.WHERE: return (None, s[0].where(g, zeros_like(g)), s[0].where(zeros_like(g), g))
case Ops.CMPLT | Ops.CMPNE | Ops.FLOORDIV | Ops.FLOORMOD: return (None,) * len(s)
case Ops.RESHAPE: return (g.reshape(s[0].shape),)
case Ops.EXPAND: return (sum_to(g, s[0].shape),)
case Ops.PERMUTE:
inv = [0] * len(u.arg)
for i, p in enumerate(u.arg): inv[p] = i
return (g.permute(inv),)
case Ops.SHRINK: return (pad(g, tuple((b, n - e) for (b, e), n in zip(u.arg, s[0].shape))),)
case Ops.STACK: return tuple(g.index(i) for i in range(len(s)))
case Ops.STAGE: return (None,) if u.arg == "detach" else (g,)
case Ops.REDUCE:
rop, n = u.arg
x = s[0]
ge = g.reshape(u.shape + (1,) * n).expand(x.shape)
if rop is Ops.ADD: return (ge,)
mx = u.reshape(u.shape + (1,) * n).expand(x.shape)
mask = x.eq(mx).cast(x.dtype) # ties share the gradient
cnt = mask.reduce(Ops.ADD, n).reshape(u.shape + (1,) * n).expand(x.shape)
return (ge * mask * cnt.recip(),)
case Ops.INDEX:
x, idx = s[0], s[1:]
if not all(i.op is Ops.CONST for i in idx): raise NotImplementedError("gradient of a data-dependent INDEX (use a one-hot matmul)")
gx = g
for i, n in zip(reversed(idx), reversed(x.shape[:len(idx)])):
z = zeros_like(gx)
gx = UOp(Ops.STACK, tuple(gx if k == i.arg else z for k in range(n)))
return (gx,) + (None,) * len(idx)
case Ops.CONST | Ops.BUFFER | Ops.PARAM: return ()
raise NotImplementedError(f"no gradient rule for {u.op} (autodiff only works on pure shape graphs)")
def gradient(root: UOp, targets: list[UOp]) -> list[UOp]:
topo = root.toposort()
# only walk nodes that lie on a path to some target
on_path = set(targets)
for u in topo:
if any(x in on_path for x in u.src): on_path.add(u)
grads: dict[UOp, UOp] = {root: _full(root, 1.0)}
for u in reversed(topo):
if u not in grads or u not in on_path or u in targets and not u.src: continue
for x, gx in zip(u.src, _rule(u, grads[u])):
if gx is None or x not in on_path: continue
grads[x] = grads[x] + gx if x in grads else gx
return [grads.get(t, zeros_like(t)) for t in targets]
uopc/tensor.py · 333 lines
"""Week 5: a lazy Tensor frontend.
A Tensor is a thin wrapper around a shaped UOp graph. Nothing runs until you call
.realize() / .numpy(): then the graph is kernelized, rangeified, rendered to C,
compiled, and executed. Every op below is a few lines of graph building on top of
the course's primitives -- there is no other math library underneath.
"""
from __future__ import annotations
import math, os
import numpy as np
from uopc.uop import UOp, Ops, prod
from uopc.dtype import DType, dtypes
from uopc.autodiff import gradient
from uopc.composite import pad as _pad, flip_axis
from uopc.kernelize import new_buffer, buffer_data, base_buffer, realize as _realize
from uopc.runtime import np_dtype
LOG2E, LN2 = 1 / math.log(2), math.log(2)
def _broadcast_shape(*shapes):
n = max(len(s) for s in shapes)
shapes = [(1,) * (n - len(s)) + tuple(s) for s in shapes]
out = []
for dims in zip(*shapes):
d = max(dims)
if any(x not in (1, d) for x in dims): raise ValueError(f"can't broadcast {shapes}")
out.append(d)
return tuple(out)
class Tensor:
rng = np.random.default_rng(0)
training = False
fuse_reduce = False
opt = None # week 6: kernel optimizer (set below)
def __init__(self, data, dtype: DType | None = None, requires_grad: bool = False):
if isinstance(data, UOp): self.uop = data
else:
arr = np.asarray(data)
if dtype is None: dtype = dtypes.f32 if arr.dtype.kind == "f" else dtypes.bool if arr.dtype.kind == "b" else dtypes.i32
arr = arr.astype(np_dtype(dtype))
buf = new_buffer(dtype, arr.size, arr)
self.uop = buf.reshape(arr.shape) if arr.shape != buf.shape else buf
self.requires_grad = requires_grad
self.grad: Tensor | None = None
# ---- basics -------------------------------------------------------------
@property
def shape(self) -> tuple[int, ...]: return self.uop.shape
@property
def dtype(self) -> DType: return self.uop.dtype
@property
def ndim(self) -> int: return len(self.shape)
def numel(self) -> int: return prod(self.shape)
def __repr__(self): return f"<Tensor {self.shape} {self.dtype.name}{' grad' if self.requires_grad else ''}>"
def realize(self) -> Tensor:
Tensor.realize_all(self)
return self
@staticmethod
def realize_all(*ts: Tensor):
todo = [t for t in ts if t.uop.op is not Ops.BUFFER and not (t.uop.op is Ops.RESHAPE and base_buffer(t.uop).op is Ops.BUFFER and prod(t.shape) == base_buffer(t.uop).arg.shape[0])]
if not todo: return
bufs = _realize([t.uop for t in todo], Tensor.fuse_reduce, Tensor.opt)
for t, b in zip(todo, bufs):
t._graph = t.uop # keep the graph: loss.item() before loss.backward() must still have gradients
t.uop = b.reshape(t.shape) if b.shape != t.shape else b
def numpy(self) -> np.ndarray:
self.realize()
return buffer_data(base_buffer(self.uop)).reshape(self.shape).copy()
def item(self): return self.numpy().item()
def detach(self) -> Tensor:
"""Same value, but gradients stop here. (Kernelize strips the marker, so it costs nothing to run.)"""
return Tensor(UOp(Ops.STAGE, (self.uop,), "detach"))
def assign(self, x: Tensor) -> Tensor:
"""Make this tensor (e.g. a parameter) hold x's value. Functional: rebinds to a new buffer."""
assert x.shape == self.shape and x.dtype == self.dtype, f"assign {x.shape} to {self.shape}"
self.uop = x.uop
return self
# ---- creation -----------------------------------------------------------
@staticmethod
def full(shape, v, dtype: DType = dtypes.f32) -> Tensor:
c = UOp.const(v, dtype)
return Tensor(c.reshape((1,) * len(shape)).expand(tuple(shape)) if shape else c)
@staticmethod
def zeros(*shape, dtype=dtypes.f32) -> Tensor: return Tensor.full(shape, 0, dtype)
@staticmethod
def ones(*shape, dtype=dtypes.f32) -> Tensor: return Tensor.full(shape, 1, dtype)
@staticmethod
def randn(*shape, requires_grad=False) -> Tensor:
return Tensor(Tensor.rng.standard_normal(shape).astype(np.float32), requires_grad=requires_grad)
@staticmethod
def uniform(*shape, low=-1.0, high=1.0, requires_grad=False) -> Tensor:
return Tensor(Tensor.rng.uniform(low, high, shape).astype(np.float32), requires_grad=requires_grad)
@staticmethod
def kaiming(*shape, requires_grad=True) -> Tensor:
bound = 1 / math.sqrt(prod(shape[1:]) if len(shape) > 1 else shape[0])
return Tensor.uniform(*shape, low=-bound, high=bound, requires_grad=requires_grad)
@staticmethod
def arange(n: int) -> Tensor: return Tensor(np.arange(n, dtype=np.int32))
@staticmethod
def manual_seed(s: int): Tensor.rng = np.random.default_rng(s)
# ---- elementwise ----------------------------------------------------------
def _const(self, x) -> Tensor: return x if isinstance(x, Tensor) else Tensor.full((), x, self.dtype)
def broadcast_to(self, shape) -> Tensor:
shape = tuple(shape)
if self.shape == shape: return self
u = self.uop
if len(u.shape) < len(shape): u = u.reshape((1,) * (len(shape) - len(u.shape)) + u.shape)
return Tensor(u.expand(shape) if u.shape != shape else u)
def _binop(self, op: Ops, x, reverse=False) -> Tensor:
# numpy-style promotion for python scalars: int tensor (op) 0.5 is a float computation
if isinstance(x, float) and not self.dtype.is_float: return self.float()._binop(op, x, reverse)
x = self._const(x)
a, b = (x, self) if reverse else (self, x)
shape = _broadcast_shape(a.shape, b.shape)
return Tensor(UOp(op, (a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))
def __add__(self, x): return self._binop(Ops.ADD, x)
def __radd__(self, x): return self._binop(Ops.ADD, x, True)
def __mul__(self, x): return self._binop(Ops.MUL, x)
def __rmul__(self, x): return self._binop(Ops.MUL, x, True)
def __neg__(self): return self * -1
def __sub__(self, x): return self + -self._const(x)
def __rsub__(self, x): return self._const(x) + -self
def __truediv__(self, x): return self * self._const(x).reciprocal()
def __rtruediv__(self, x): return self._const(x) * self.reciprocal()
def __lt__(self, x): return self._binop(Ops.CMPLT, x)
def __gt__(self, x): return self._binop(Ops.CMPLT, x, True)
def __ge__(self, x): return (self < x).logical_not()
def __le__(self, x): return (self > x).logical_not()
def eq(self, x): return self._binop(Ops.CMPNE, x).logical_not()
def ne(self, x): return self._binop(Ops.CMPNE, x)
def logical_not(self): return self._binop(Ops.CMPNE, True)
def maximum(self, x): return self._binop(Ops.MAX, x)
def minimum(self, x): return -((-self).maximum(-self._const(x)))
def where(self, a, b) -> Tensor:
a = a if isinstance(a, Tensor) else Tensor.full((), a)
b = b if isinstance(b, Tensor) else Tensor.full((), b, a.dtype)
shape = _broadcast_shape(self.shape, a.shape, b.shape)
return Tensor(UOp(Ops.WHERE, (self.broadcast_to(shape).uop, a.broadcast_to(shape).uop, b.broadcast_to(shape).uop)))
def cast(self, d: DType) -> Tensor: return Tensor(self.uop.cast(d))
def float(self) -> Tensor: return self.cast(dtypes.f32)
def reciprocal(self): return Tensor(self.uop.recip())
def exp2(self): return Tensor(UOp(Ops.EXP2, (self.uop,)))
def log2(self): return Tensor(UOp(Ops.LOG2, (self.uop,)))
def exp(self): return (self * LOG2E).exp2()
def log(self): return self.log2() * LN2
def sqrt(self): return Tensor(UOp(Ops.SQRT, (self.uop,)))
def rsqrt(self): return self.sqrt().reciprocal()
def sin(self): return Tensor(UOp(Ops.SIN, (self.uop,)))
def cos(self): return (self + math.pi / 2).sin()
def square(self): return self * self
def abs(self): return self.maximum(-self)
def relu(self): return (self > 0).where(self, 0.0)
def sigmoid(self): return ((-self).exp() + 1).reciprocal()
def tanh(self): return (self * 2).sigmoid() * 2 - 1
def gelu(self): return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self * self * self)).tanh())
def silu(self): return self * self.sigmoid()
# ---- movement -------------------------------------------------------------
def reshape(self, *shape) -> Tensor:
shape = tuple(shape[0]) if len(shape) == 1 and isinstance(shape[0], (tuple, list)) else shape
if -1 in shape:
known = prod(s for s in shape if s != -1)
shape = tuple(self.numel() // known if s == -1 else s for s in shape)
return self if shape == self.shape else Tensor(self.uop.reshape(shape))
def permute(self, *order) -> Tensor:
order = tuple(order[0]) if len(order) == 1 and isinstance(order[0], (tuple, list)) else order
order = tuple(o % self.ndim for o in order)
return self if order == tuple(range(self.ndim)) else Tensor(self.uop.permute(order))
def transpose(self, a=-1, b=-2) -> Tensor:
order = list(range(self.ndim))
order[a], order[b] = order[b], order[a]
return self.permute(order)
@property
def T(self) -> Tensor: return self.transpose()
def expand(self, *shape) -> Tensor:
shape = tuple(shape[0]) if len(shape) == 1 and isinstance(shape[0], (tuple, list)) else shape
return self.broadcast_to(tuple(s if s != -1 else o for s, o in zip(shape, self.shape)))
def flatten(self, start=0) -> Tensor: return self.reshape(self.shape[:start] + (-1,))
def unsqueeze(self, d) -> Tensor:
d = d % (self.ndim + 1)
return self.reshape(self.shape[:d] + (1,) + self.shape[d:])
def shrink(self, arg) -> Tensor:
arg = tuple((0, s) if a is None else a for a, s in zip(arg, self.shape))
return self if all(a == (0, s) for a, s in zip(arg, self.shape)) else Tensor(self.uop.shrink(arg))
def pad(self, arg) -> Tensor: return Tensor(_pad(self.uop, tuple(arg)))
def flip(self, axis: int) -> Tensor: return Tensor(flip_axis(self.uop, axis % self.ndim))
def __getitem__(self, idx) -> Tensor:
"""Ints and step-1 slices only."""
idx = idx if isinstance(idx, tuple) else (idx,)
arg, squeeze = [], []
for d, i in enumerate(idx):
n = self.shape[d]
if isinstance(i, int):
i = i % n
arg.append((i, i + 1)); squeeze.append(d)
else:
b, e, step = i.indices(n)
assert step == 1, "only step-1 slices"
arg.append((b, max(b, e)))
arg += [(0, s) for s in self.shape[len(arg):]]
t = self.shrink(arg)
return t.reshape(tuple(s for d, s in enumerate(t.shape) if d not in squeeze))
@staticmethod
def cat(ts: list[Tensor], dim: int = 0) -> Tensor:
"""Concatenate via pad + add: each piece is padded into place, then summed."""
dim = dim % ts[0].ndim
total = sum(t.shape[dim] for t in ts)
out, off = None, 0
for t in ts:
p = [(0, 0)] * t.ndim
p[dim] = (off, total - off - t.shape[dim])
piece = t.pad(p)
out = piece if out is None else out + piece
off += t.shape[dim]
return out
def stack_with(self, *others) -> Tensor: return Tensor(UOp(Ops.STACK, (self.uop, *(o.uop for o in others))))
# ---- reductions -----------------------------------------------------------
def _reduce(self, op: Ops, axis=None, keepdim=False) -> Tensor:
axes = list(range(self.ndim)) if axis is None else [a % self.ndim for a in (axis if isinstance(axis, (tuple, list)) else [axis])]
keep = [i for i in range(self.ndim) if i not in axes]
t = self.permute(keep + axes)
r = Tensor(t.uop.reduce(op, len(axes)))
if keepdim: r = r.reshape(tuple(1 if i in axes else s for i, s in enumerate(self.shape)))
return r
def sum(self, axis=None, keepdim=False): return self._reduce(Ops.ADD, axis, keepdim)
def max(self, axis=None, keepdim=False): return self._reduce(Ops.MAX, axis, keepdim)
def min(self, axis=None, keepdim=False): return -((-self).max(axis, keepdim))
def mean(self, axis=None, keepdim=False):
s = (self if self.dtype.is_float else self.float()).sum(axis, keepdim)
return s * (s.numel() / self.numel())
def var(self, axis=None, keepdim=False):
return (self - self.mean(axis, keepdim=True)).square().mean(axis, keepdim)
def softmax(self, axis=-1):
e = (self - self.max(axis, keepdim=True)).exp()
return e / e.sum(axis, keepdim=True)
def log_softmax(self, axis=-1):
m = self - self.max(axis, keepdim=True)
return m - m.exp().sum(axis, keepdim=True).log()
def logsumexp(self, axis=-1):
m = self.max(axis, keepdim=True)
return (self - m).exp().sum(axis, keepdim=True).log() + m
def argmax(self, axis=-1):
"""Index of the first maximum, as i32."""
axis = axis % self.ndim
n = self.shape[axis]
idx = Tensor.arange(n).reshape((n,) + (1,) * (self.ndim - axis - 1))
is_max = self.eq(self.max(axis, keepdim=True))
return is_max.where(n - idx, 0).max(axis) * -1 + n
# ---- linear algebra -------------------------------------------------------
def matmul(self, w: Tensor) -> Tensor:
"""(..., M, K) @ (..., K, N): broadcast to (..., M, N, K), multiply, reduce K."""
x = self if self.ndim > 1 else self.reshape(1, -1)
wt = w.reshape(-1, 1) if w.ndim == 1 else w
wt = wt.transpose() # (..., N, K)
x = x.reshape(x.shape[:-1] + (1, x.shape[-1])) # (..., M, 1, K)
wt = wt.reshape(wt.shape[:-2] + (1,) + wt.shape[-2:]) # (..., 1, N, K)
out = (x * wt).sum(-1)
if self.ndim == 1: out = out.reshape(out.shape[:-2] + out.shape[-1:])
if w.ndim == 1: out = out.reshape(out.shape[:-1])
return out
__matmul__ = matmul
def linear(self, w: Tensor, b: Tensor | None = None) -> Tensor:
out = self @ w
return out + b if b is not None else out
def conv2d(self, w: Tensor, b: Tensor | None = None, padding: int = 0, stride: int = 1) -> Tensor:
"""x (N, C, H, W), w (O, C, KH, KW). im2col from shrinks, then a matmul."""
x = self.pad(((0, 0), (0, 0), (padding, padding), (padding, padding))) if padding else self
N, C, H, W = x.shape
O, _, KH, KW = w.shape
OH, OW = (H - KH) // stride + 1, (W - KW) // stride + 1
cols = []
for ky in range(KH):
for kx in range(KW):
win = x.shrink(((0, N), (0, C), (ky, ky + stride * (OH - 1) + 1), (kx, kx + stride * (OW - 1) + 1)))
if stride > 1: win = win._stride2d(stride, OH, OW)
cols.append(win) # (N, C, OH, OW) each
patches = cols[0].stack_with(*cols[1:]) if len(cols) > 1 else cols[0].unsqueeze(0) # (KH*KW, N, C, OH, OW)
patches = patches.permute(1, 3, 4, 2, 0).reshape(N, OH, OW, C * KH * KW)
out = patches @ w.reshape(O, C * KH * KW).T # (N, OH, OW, O)
out = out.permute(0, 3, 1, 2)
return out + b.reshape(1, O, 1, 1) if b is not None else out
def _stride2d(self, s: int, OH: int, OW: int) -> Tensor:
N, C, H, W = self.shape
t = self.pad(((0, 0), (0, 0), (0, OH * s - H), (0, OW * s - W)))
return t.reshape(N, C, OH, s, OW, s)[:, :, :, 0, :, 0]
def max_pool2d(self, k: int = 2) -> Tensor:
N, C, H, W = self.shape
return self[:, :, :H // k * k, :W // k * k].reshape(N, C, H // k, k, W // k, k).max((3, 5))
# ---- losses ---------------------------------------------------------------
def one_hot(self, n: int) -> Tensor:
return self.unsqueeze(-1).eq(Tensor.arange(n)).float()
def cross_entropy(self, labels: Tensor) -> Tensor:
"""self: (B, C) logits, labels: (B,) i32."""
return -(self.log_softmax(-1) * labels.one_hot(self.shape[-1])).sum(-1).mean()
# ---- autodiff ------------------------------------------------------------
def backward(self, params: list[Tensor]) -> Tensor:
"""Set p.grad = d(self)/d(p) for each p. Gradients are lazy Tensors like everything else."""
assert self.shape == (), "backward needs a scalar"
root = getattr(self, "_graph", self.uop)
for p, g in zip(params, gradient(root, [p.uop for p in params])):
p.grad = Tensor(g) if p.grad is None else p.grad + Tensor(g)
return self
# ---- staging -------------------------------------------------------------
def contiguous(self) -> Tensor:
"""Ask for a materialization point: the kernelizer will write this into a temp buffer."""
return Tensor(UOp(Ops.STAGE, (self.uop,)))
if os.environ.get("OPT", "1") != "0":
from uopc.opt import HandOpt, SearchOpt
Tensor.opt = SearchOpt() if os.environ.get("OPT") == "2" else HandOpt()
uopc/nn.py · 100 lines
"""Week 5: layers and optimizers, written against the Tensor API only."""
from __future__ import annotations
import math
from uopc.tensor import Tensor
class Linear:
def __init__(self, n_in: int, n_out: int, bias: bool = True):
bound = 1 / math.sqrt(n_in)
self.weight = Tensor.uniform(n_in, n_out, low=-bound, high=bound, requires_grad=True)
self.bias = Tensor.zeros(n_out).realize() if bias else None
if self.bias is not None: self.bias.requires_grad = True
def __call__(self, x: Tensor) -> Tensor: return x.linear(self.weight, self.bias)
class Conv2d:
def __init__(self, c_in: int, c_out: int, k: int, padding: int = 0, stride: int = 1):
bound = 1 / math.sqrt(c_in * k * k)
self.weight = Tensor.uniform(c_out, c_in, k, k, low=-bound, high=bound, requires_grad=True)
self.bias = Tensor.zeros(c_out).realize()
self.bias.requires_grad = True
self.padding, self.stride = padding, stride
def __call__(self, x: Tensor) -> Tensor: return x.conv2d(self.weight, self.bias, self.padding, self.stride)
class LayerNorm:
def __init__(self, n: int, eps: float = 1e-5):
self.weight, self.bias, self.eps = Tensor.ones(n).realize(), Tensor.zeros(n).realize(), eps
self.weight.requires_grad = self.bias.requires_grad = True
def __call__(self, x: Tensor) -> Tensor:
mu = x.mean(-1, keepdim=True)
xc = x - mu
return xc * (xc.square().mean(-1, keepdim=True) + self.eps).rsqrt() * self.weight + self.bias
class Embedding:
"""A one-hot matmul: differentiable without a scatter op."""
def __init__(self, vocab: int, dim: int):
self.vocab, self.weight = vocab, Tensor.randn(vocab, dim) * 0.02
self.weight = self.weight.realize(); self.weight.requires_grad = True
def __call__(self, idx: Tensor) -> Tensor: return idx.one_hot(self.vocab) @ self.weight
def parameters(obj) -> list[Tensor]:
"""Every Tensor with requires_grad reachable through attributes, lists and dicts."""
out, seen = [], set()
def walk(o):
if id(o) in seen: return
seen.add(id(o))
if isinstance(o, Tensor):
if o.requires_grad: out.append(o)
elif isinstance(o, (list, tuple)): [walk(x) for x in o]
elif isinstance(o, dict): [walk(x) for x in o.values()]
elif hasattr(o, "__dict__"): [walk(x) for x in vars(o).values()]
walk(obj)
return out
class Optimizer:
def __init__(self, params: list[Tensor]): self.params = params
def zero_grad(self):
for p in self.params: p.grad = None
class SGD(Optimizer):
def __init__(self, params, lr=0.01, momentum=0.0):
super().__init__(params)
self.lr, self.momentum = lr, momentum
self.v = [Tensor.zeros(*p.shape).realize() for p in params] if momentum else []
def step(self, *extra: Tensor):
"""Update every parameter. Extra tensors (e.g. the loss) are realized in the same batch."""
news = []
for i, p in enumerate(self.params):
g = p.grad
if self.momentum:
self.v[i].assign(self.v[i] * self.momentum + g); g = self.v[i]
news.append(p.assign(p.detach() - g * self.lr))
Tensor.realize_all(*news, *self.v, *extra)
class AdamW(Optimizer):
def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=0.0):
super().__init__(params)
self.lr, self.b1, self.b2, self.eps, self.wd, self.t = lr, b1, b2, eps, wd, 0
self.m = [Tensor.zeros(*p.shape).realize() for p in params]
self.v = [Tensor.zeros(*p.shape).realize() for p in params]
def step(self, *extra: Tensor):
self.t += 1
# bias corrections are python floats baked into the graph as constants: each step's
# graph differs only in two CONSTs, so we pass them as tensors to keep the JIT cache warm
c1 = Tensor([self.lr / (1 - self.b1 ** self.t)])
c2 = Tensor([1 / (1 - self.b2 ** self.t)])
news = []
for p, m, v in zip(self.params, self.m, self.v):
g = p.grad
m.assign(m * self.b1 + g * (1 - self.b1))
v.assign(v * self.b2 + g.square() * (1 - self.b2))
upd = m * c1.reshape(*([1] * p.ndim)) / ((v * c2.reshape(*([1] * p.ndim))).sqrt() + self.eps)
news.append(p.assign(p.detach() * (1 - self.lr * self.wd) - upd))
Tensor.realize_all(*news, *self.m, *self.v, *extra)
uopc/opt.py · 163 lines
"""Week 6: making kernels fast on a CPU.
Three knobs, all applied while rangeifying (see rangeify.Opts):
UPCAST unroll output axes into a register tile. A 4x8 tile of a GEMM keeps 32
accumulators in registers and reuses each loaded A value 8 times and each
B value 4 times: arithmetic intensity goes from 1/2 to ~2.7 FLOP/load.
UNROLL unroll the innermost reduce axis; the loop body becomes a balanced tree.
THREAD run the outermost output loop on every core with OpenMP.
TILE split output axes into cache blocks (Opts.tile). Only search_opt tries it: with the layout
pass making the reduce axis contiguous, it measured no faster on the reference machine.
`hand_opt` picks them with simple rules. `search_opt` times candidates and keeps
the fastest (a tiny BEAM search).
"""
from __future__ import annotations
import itertools, os, time
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, ALU, MOVEMENT, prod
from uopc.rangeify import lower, rangeify, Ctx, Opts
from uopc.symbolic import simplify
from uopc.rewrite import UPat, PatternMatcher, graph_rewrite
NTHREADS = os.cpu_count() or 1
# ---------------------------------------------------------------------------
# layout: make the innermost reduce axis contiguous in memory
def _views(x: UOp):
"""The movement-op chains feeding x through ALU ops: (top of chain, materialized base)."""
seen, stack = set(), [x]
while stack:
u = stack.pop()
if u in seen: continue
seen.add(u)
if u.op in ALU: stack.extend(u.src)
elif u.op in MOVEMENT:
base = u
while base.op in MOVEMENT: base = base.src[0]
if base.op in (Ops.BUFFER, Ops.REDUCE, Ops.STAGE): yield u, base
def _last_axis_stride(view: UOp, base: UOp) -> int | None:
"""Distance in memory between view[..., k] and view[..., k+1], or None if not affine-looking."""
if view.shape[-1] == 1: return 0
fake = UOp(Ops.BUFFER, (), ParamArg(-999, base.dtype, (prod(base.shape),)))
v = view.substitute({base: fake.reshape(base.shape) if base.shape != fake.shape else fake})
ctx = Ctx()
rngs = tuple(ctx.rng(s) for s in v.shape)
val = rangeify(v, rngs, ctx)
if val.op is not Ops.LOAD or val.src[0].op is not Ops.INDEX: return None
e = val.src[0].src[1]
at = lambda k: simplify(e.substitute({r: UOp.const(k if i == len(rngs) - 1 else 0) for i, r in enumerate(rngs)}))
e0, e1 = at(0), at(1)
return e1.arg - e0.arg if e0.op is Ops.CONST and e1.op is Ops.CONST else None
# movement ops commute with elementwise ops: permute(a * b) == permute(a) * permute(b).
# Pushing them down to the leaves makes every view chain end at real data.
push_views = PatternMatcher([
(UPat(MOVEMENT, src=(UPat(ALU, name="a"),), name="m"),
lambda m, a: a.replace(src=tuple(m.replace(src=(s,)) for s in a.src))),
])
def layout_pass(roots: list[UOp]) -> list[UOp]:
"""Before broadcasting, copy operands whose reduce axis is strided into a contiguous temp (a STAGE).
The copy costs O(size of the operand); the reduce it feeds is usually much bigger.
This is what turns `A @ B` (B read down its columns) into a dot product of rows.
"""
roots = [graph_rewrite(r, push_views) for r in roots]
todo: set[UOp] = set()
for r in UOp(Ops.SINK, tuple(roots)).toposort():
if r.op is not Ops.REDUCE: continue
for view, base in _views(r.src[0]):
s = _last_axis_stride(view, base)
if s in (0, 1, None): continue
c = view
while c.op is Ops.EXPAND: c = c.src[0] # stage before the broadcast
# only stage operands that get broadcast into the reduce: copying them is cheap
if prod(c.shape) > prod(base.shape) or prod(c.shape) >= prod(r.src[0].shape) or c in todo: continue
staged = UOp(Ops.STAGE, (c,))
if _last_axis_stride(view.substitute({c: staged}), staged) == 1: todo.add(c)
return UOp.wrap(roots, todo, lambda n: UOp(Ops.STAGE, (n,))) if todo else roots
def _reduce_size(body: UOp) -> int:
return max((prod(u.src[0].shape[len(u.src[0].shape) - u.arg[1]:]) for u in body.toposort() if u.op is Ops.REDUCE), default=1)
def hand_opt_choice(body: UOp) -> Opts:
shape, rsize = body.shape, _reduce_size(body)
o = Opts()
if rsize > 1:
# register tile: 4x4 on the last two output axes. With AVX2's 16 vector registers, 16 accumulators plus
# the loaded operands don't quite fit (gcc spills a few), but measured it beats 2x8, 4x8 and 8x8 here
budget = 16
for axis, want in [(len(shape) - 1, 4), (len(shape) - 2, 4)]:
if axis < 0: continue
for f in (want, 4, 2):
if f <= budget and shape[axis] % f == 0 and shape[axis] >= f and f > 1:
o.upcast[axis] = f; budget //= f; break
if not o.upcast and rsize % 4 == 0 and rsize < 64: o.unroll = 4
# threads: only when there is enough work to amortize waking them up
work = prod(shape or (1,)) * rsize
if work >= 1 << 16:
for axis, s in enumerate(shape):
if s // o.upcast.get(axis, 1) >= 2:
o.thread = axis; break
return o
def hand_opt(body: UOp, out: UOp) -> UOp:
return lower(body, out, Ctx(opts=hand_opt_choice(body)))
def candidate_opts(body: UOp) -> list[Opts]:
shape, rsize = body.shape, _reduce_size(body)
outs = [a for a, s in enumerate(shape) if s > 1]
cands = [Opts(), hand_opt_choice(body)]
for a, b in itertools.product(outs[-2:], repeat=2):
for fa, fb in [(4, 4), (8, 4), (4, 8), (8, 8), (16, 2), (2, 16)]:
up = {a: fa} if a == b else {a: fa, b: fb}
if all(shape[k] % f == 0 for k, f in up.items()):
for u in (1, 4) if rsize % 4 == 0 else (1,):
cands.append(Opts(dict(up), hand_opt_choice(body).thread, u))
# cache blocking of the last two output axes (helps when operand rows don't fit in cache)
base = hand_opt_choice(body)
for t in (32, 64, 128):
tile = {a: t for a in outs[-2:] if shape[a] % t == 0}
if tile: cands.append(Opts(dict(base.upcast), base.thread, base.unroll, tile=tile))
return cands
def search_opt(body: UOp, out: UOp, inputs: list[np.ndarray] | None = None, reps: int = 3) -> UOp:
"""Time every candidate on random data, keep the fastest."""
from uopc.render_c import render_c
from uopc.runtime import Program, np_dtype
bufs = sorted({u for u in body.toposort() if u.op is Ops.BUFFER}, key=lambda u: u.arg.slot)
data = inputs or [np.random.randn(b.arg.shape[0]).astype(np_dtype(b.dtype)) for b in bufs]
res = np.zeros(out.arg.shape[0], np_dtype(out.dtype))
best, best_t = None, float("inf")
seen = set()
for o in candidate_opts(body):
key = (tuple(sorted(o.upcast.items())), o.thread, o.unroll, tuple(sorted(o.tile.items())))
if key in seen: continue
seen.add(key)
ast = lower(body, out, Ctx(opts=o))
prg = Program(render_c(ast, "k"), "k")
ts = []
for _ in range(reps):
t0 = time.perf_counter(); prg(*data, res); ts.append(time.perf_counter() - t0)
if min(ts) < best_t: best, best_t = ast, min(ts)
return best
class HandOpt:
"""The default optimizer kernelize uses: layout pass + hand-coded opts."""
layout = staticmethod(layout_pass)
lower = staticmethod(hand_opt)
class SearchOpt(HandOpt):
"""Like HandOpt, but times candidate opts for every new kernel (slow to compile, fast to run)."""
lower = staticmethod(search_opt)
uopc/gpu.py · 181 lines
"""Weeks 7-8: GPUs.
A GPU kernel is the body of a loop nest whose outer loops the hardware runs for
you: GLOBAL ranges become blockIdx, LOCAL ranges become threadIdx. Everything
else (reduce loops, upcasts, accumulators) is exactly what we already have.
* CUDARenderer -- renders the same linearized UOps as CUDA C
* gpu_opt -- picks blocks/threads/upcasts for a kernel
* emulate -- renders a GPU kernel with the *C* renderer, where GLOBAL and
LOCAL ranges are ordinary loops. Same index math, runs on a CPU,
so every GPU kernel is testable without a GPU.
* smem_gemm_ir -- a tiled GEMM written by hand in the loop IR, using SHARED memory
and BARRIER: the core trick of every fast GPU GEMM
* wmma_gemm -- tensor cores: pattern-match a GEMM and emit a WMMA kernel
"""
from __future__ import annotations
from uopc.uop import UOp, Ops, ParamArg, AxisType, prod, sink, range_
from uopc.dtype import dtypes, AddrSpace
from uopc.render_c import CRenderer
from uopc.rangeify import lower, Ctx, Opts
CUDA_PRELUDE = """#define INFINITY (__int_as_float(0x7f800000))
#define NAN (__int_as_float(0x7fffffff))
__device__ inline int floordiv_i32(int a, int b) { int q = a / b; return (q * b != a && ((a < 0) != (b < 0))) ? q - 1 : q; }
__device__ inline int floormod_i32(int a, int b) { int r = a % b; return (r != 0 && ((r < 0) != (b < 0))) ? r + b : r; }
"""
class CUDARenderer(CRenderer):
buffer_prefix = ""
shared_prefix = "__shared__ "
barrier = "__syncthreads()"
restrict = "__restrict__"
def render(self, sink: UOp, name: str = "kernel") -> str:
rngs = [u for u in sink.toposort() if u.op is Ops.RANGE]
# the last (fastest varying) axis gets .x, so neighbouring threads touch neighbouring memory
self.glob = [u for u in sorted(rngs, key=lambda r: r.arg, reverse=True) if u.arg.kind is AxisType.GLOBAL]
self.loc = [u for u in sorted(rngs, key=lambda r: r.arg, reverse=True) if u.arg.kind is AxisType.LOCAL]
assert len(self.glob) <= 3 and len(self.loc) <= 3, "fold axes: CUDA has 3 grid and 3 block dimensions"
self.kernel_prefix = f'extern "C" __global__ void __launch_bounds__({max(1, prod(launch_dims(sink)[1]))})'
# dtypes.f16.ctype is gcc's _Float16; CUDA spells it `half` (from cuda_fp16.h)
return super().render(sink, name).replace("_Float16", "half")
def header(self, sink, lin) -> str:
return ("#include <cuda_fp16.h>\n" if any(u.dtype == dtypes.f16 for u in lin) else "") + CUDA_PRELUDE
def is_loop(self, u: UOp) -> bool: return u.arg.kind not in (AxisType.GLOBAL, AxisType.LOCAL)
def range_name(self, u: UOp) -> str:
if u.arg.kind is AxisType.GLOBAL: return f"gidx{self.glob.index(u)}"
if u.arg.kind is AxisType.LOCAL: return f"lidx{self.loc.index(u)}"
return super().range_name(u)
def render_range(self, u: UOp, name: str, end: str) -> list[str]:
if u.arg.kind is AxisType.GLOBAL: return [f"int {name} = blockIdx.{'xyz'[self.glob.index(u)]}; /* {end} */"]
if u.arg.kind is AxisType.LOCAL: return [f"int {name} = threadIdx.{'xyz'[self.loc.index(u)]}; /* {end} */"]
return [f"for (int {name} = 0; {name} < {end}; {name}++) {{"]
def launch_dims(sink: UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
"""(grid, block) sizes from the GLOBAL and LOCAL ranges, in x, y, z order."""
rngs = sorted((u for u in sink.toposort() if u.op is Ops.RANGE), key=lambda r: r.arg, reverse=True)
g = tuple(r.src[0].arg for r in rngs if r.arg.kind is AxisType.GLOBAL)
l = tuple(r.src[0].arg for r in rngs if r.arg.kind is AxisType.LOCAL)
return (g + (1,) * (3 - len(g)), l + (1,) * (3 - len(l)))
def render_cuda(sink: UOp, name: str = "kernel") -> str: return CUDARenderer().render(sink, name)
# ---------------------------------------------------------------------------
# picking a launch configuration
def gpu_opt_choice(body: UOp, max_threads: int = 256) -> Opts:
"""Up to 3 output axes on the grid. Threads on the last axes (coalesced stores), 2x2 upcast for reduces."""
shape = body.shape
o = Opts(gpu=True)
reduces = any(u.op is Ops.REDUCE for u in body.toposort())
axes = [k for k, s in enumerate(shape) if s > 1][-3:]
if len(axes) < len(shape) - shape.count(1):
raise NotImplementedError("more than 3 non-trivial output axes: reshape first")
budget = max_threads
for k in reversed(axes):
if reduces and shape[k] % 2 == 0 and shape[k] >= 32: o.upcast[k] = 2
rest = shape[k] // o.upcast.get(k, 1)
t = 1
while t * 2 <= min(budget, 16 if len(axes) > 1 else budget) and rest % (t * 2) == 0: t *= 2
if t > 1: o.local[k] = t; budget //= t
return o
def gpu_lower(body: UOp, out: UOp, opts: Opts | None = None) -> UOp:
return lower(body, out, Ctx(opts=opts or gpu_opt_choice(body)))
def emulate(sink: UOp, name: str = "kernel") -> str:
"""The C source of a GPU kernel where blockIdx/threadIdx are loops. Run it with runtime.Program."""
if any(u.op is Ops.BARRIER for u in sink.toposort()):
raise NotImplementedError("emulate: kernels with BARRIER need all threads to run in lockstep; "
"running threads one after another (what loops do) would be wrong")
return CRenderer().render(sink, name)
# ---------------------------------------------------------------------------
# shared memory: a tiled GEMM in the loop IR
def smem_gemm_ir(N: int, T: int = 16) -> UOp:
"""C = A @ B, N x N, with T x T thread blocks and T x T tiles staged in shared memory.
Each block computes one T x T tile of C. For each of the N/T steps along K, the
block cooperatively copies a T x T tile of A and of B into shared memory (one
element per thread), waits at a barrier, multiplies the tiles out of shared
memory, and waits again before the next copy overwrites them. Global memory
traffic drops by a factor of T.
"""
f32 = dtypes.f32
A, B, C = (UOp(Ops.BUFFER, (), ParamArg(s, f32, (N * N,))) for s in range(3))
sA = UOp(Ops.ALLOC, (), ParamArg(0, f32, (T * T,), AddrSpace.SHARED))
sB = UOp(Ops.ALLOC, (), ParamArg(1, f32, (T * T,), AddrSpace.SHARED))
acc = UOp(Ops.ALLOC, (), ParamArg(2, f32, (1,)))
by, bx = range_(N // T, 0, AxisType.GLOBAL), range_(N // T, 1, AxisType.GLOBAL)
ty, tx = range_(T, 2, AxisType.LOCAL), range_(T, 3, AxisType.LOCAL)
row, col = by * T + ty, bx * T + tx
acc0 = acc.after(by, bx, ty, tx)
init = acc0.index(0).store(UOp.const(0.0))
kt = range_(N // T, 4) # step along K, one tile at a time
ld_a = sA.after(kt).index(ty * T + tx).store(A.index(row * N + kt * T + tx).load())
ld_b = sB.after(kt).index(ty * T + tx).store(B.index((kt * T + ty) * N + col).load())
sync1 = UOp(Ops.BARRIER, (ld_a, ld_b))
k = range_(T, 5)
accl = acc0.after(init, kt, k)
prod_ = sA.after(sync1).index(ty * T + k).load() * sB.after(sync1).index(k * T + tx).load()
inner = accl.index(0).store(accl.index(0).load() + prod_).end(k)
sync2 = UOp(Ops.BARRIER, (inner,)) # don't overwrite tiles others still read
loop = sync2.end(kt)
out = C.index(row * N + col).store(acc0.after(loop).index(0).load())
return sink(out.end(by, bx, ty, tx))
# ---------------------------------------------------------------------------
# tensor cores: pattern-match a GEMM, emit a WMMA kernel
def match_gemm(body: UOp) -> tuple[int, int, int] | None:
"""If body is (M,N) = sum_K A(M,K) * B(K,N) over contiguous f16/f32 buffers, return (M, N, K)."""
if body.op is not Ops.REDUCE or body.arg != (Ops.ADD, 1) or body.src[0].op is not Ops.MUL: return None
if len(body.src[0].shape) != 3 or not body.dtype.is_float: return None
M, N, K = body.src[0].shape
# the two operands must be broadcast along *different* axes: A over N (axis 1), B over M (axis 0).
# Without this check, x*x reduced over its last axis (a sum of squares) "is" a GEMM.
a, b = body.src[0].src
def bcast_axes(x: UOp) -> set[int]:
return {i for i, (s, t) in enumerate(zip(x.src[0].shape, x.shape)) if s == 1 and t > 1} if x.op is Ops.EXPAND else set()
if {frozenset(bcast_axes(a)), frozenset(bcast_axes(b))} != {frozenset({1}), frozenset({0})}: return None
return (M, N, K) if all(x % 16 == 0 for x in (M, N, K)) else None
def wmma_gemm(M: int, N: int, K: int) -> str:
"""C(f32, MxN) = A(f16, MxK) @ B(f16, KxN), one 16x16 output tile per warp, via nvcuda::wmma.
A warp (32 threads) cooperatively owns three "fragments": 16x16 tiles of A, B and C
spread across its registers in a hardware-defined layout. mma_sync does a
16x16x16 matrix multiply-accumulate in one instruction on the tensor cores.
"""
return f"""#include <mma.h>
using namespace nvcuda;
extern "C" __global__ void wmma_gemm(const half* __restrict__ A, const half* __restrict__ B, float* __restrict__ C) {{
int tile_m = blockIdx.y, tile_n = blockIdx.x; // one warp per 16x16 tile of C
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c;
wmma::fill_fragment(c, 0.0f);
for (int k = 0; k < {K}; k += 16) {{
wmma::load_matrix_sync(a, A + tile_m * 16 * {K} + k, {K});
wmma::load_matrix_sync(b, B + k * {N} + tile_n * 16, {N});
wmma::mma_sync(c, a, b, c);
}}
wmma::store_matrix_sync(C + tile_m * 16 * {N} + tile_n * 16, c, {N}, wmma::mem_row_major);
}}
/* launch: grid = ({N // 16}, {M // 16}), block = (32, 1, 1) */
"""
uopc/cuda.py · 90 lines
"""Weeks 7-8: a minimal CUDA runtime over ctypes: NVRTC to compile, the driver API to run.
NOTE: this file needs an NVIDIA GPU, a working driver (libcuda.so) and NVRTC
(libnvrtc.so, from the CUDA toolkit). The reference machine for this book had
neither working, so this module is exercised only by `tests/test_w07_gpu.py`
when CUDA is available; elsewhere GPU kernels are tested with gpu.emulate().
"""
from __future__ import annotations
import ctypes, ctypes.util
import numpy as np
def _lib(*names):
for n in names:
path = ctypes.util.find_library(n)
if path: return ctypes.CDLL(path)
raise OSError(f"none of {names} found")
class CUDAError(Exception): pass
class CUDA:
def __init__(self, device: int = 0):
self.cu = _lib("cuda")
self.rtc = _lib("nvrtc")
self._check(self.cu.cuInit(0))
self.dev = ctypes.c_int()
self._check(self.cu.cuDeviceGet(ctypes.byref(self.dev), device))
self.ctx = ctypes.c_void_p()
self._check(self.cu.cuCtxCreate_v2(ctypes.byref(self.ctx), 0, self.dev))
major, minor = ctypes.c_int(), ctypes.c_int()
self._check(self.cu.cuDeviceComputeCapability(ctypes.byref(major), ctypes.byref(minor), self.dev))
self.arch = f"sm_{major.value}{minor.value}"
def _check(self, status):
if status != 0: raise CUDAError(f"CUDA driver error {status}")
def compile(self, src: str) -> bytes:
prog = ctypes.c_void_p()
if self.rtc.nvrtcCreateProgram(ctypes.byref(prog), src.encode(), b"kernel.cu", 0, None, None):
raise CUDAError("nvrtcCreateProgram failed")
opts = [f"--gpu-architecture=compute_{self.arch[3:]}".encode(), b"--use_fast_math"]
status = self.rtc.nvrtcCompileProgram(prog, len(opts), (ctypes.c_char_p * len(opts))(*opts))
if status != 0:
n = ctypes.c_size_t()
self.rtc.nvrtcGetProgramLogSize(prog, ctypes.byref(n))
log = ctypes.create_string_buffer(n.value)
self.rtc.nvrtcGetProgramLog(prog, log)
raise CUDAError(f"nvrtc: {log.value.decode()}\n{src}")
n = ctypes.c_size_t()
self.rtc.nvrtcGetPTXSize(prog, ctypes.byref(n))
ptx = ctypes.create_string_buffer(n.value)
self.rtc.nvrtcGetPTX(prog, ptx)
return ptx.raw
def alloc(self, nbytes: int) -> ctypes.c_uint64:
p = ctypes.c_uint64()
self._check(self.cu.cuMemAlloc_v2(ctypes.byref(p), ctypes.c_size_t(nbytes)))
return p
def to_device(self, a: np.ndarray) -> ctypes.c_uint64:
p = self.alloc(a.nbytes)
self._check(self.cu.cuMemcpyHtoD_v2(p, a.ctypes.data_as(ctypes.c_void_p), ctypes.c_size_t(a.nbytes)))
return p
def to_host(self, p: ctypes.c_uint64, a: np.ndarray):
self._check(self.cu.cuMemcpyDtoH_v2(a.ctypes.data_as(ctypes.c_void_p), p, ctypes.c_size_t(a.nbytes)))
def free(self, p): self.cu.cuMemFree_v2(p)
def function(self, src: str, name: str):
mod = ctypes.c_void_p()
self._check(self.cu.cuModuleLoadData(ctypes.byref(mod), self.compile(src)))
fn = ctypes.c_void_p()
self._check(self.cu.cuModuleGetFunction(ctypes.byref(fn), mod, name.encode()))
return fn
def launch(self, fn, grid, block, *ptrs, shared: int = 0):
args = (ctypes.c_void_p * len(ptrs))(*(ctypes.cast(ctypes.byref(p), ctypes.c_void_p) for p in ptrs))
self._check(self.cu.cuLaunchKernel(fn, *grid, *block, shared, None, args, None))
self._check(self.cu.cuCtxSynchronize())
def run_cuda(src: str, name: str, grid, block, *arrays: np.ndarray):
"""Copy arrays to the GPU, run, copy every array back (they may be outputs)."""
cu = CUDA()
ptrs = [cu.to_device(a) for a in arrays]
cu.launch(cu.function(src, name), grid, block, *ptrs)
for p, a in zip(ptrs, arrays): cu.to_host(p, a); cu.free(p)
D.2 Tests: tests/
One file per week. helpers.py lowers and runs a shaped graph against numpy; conftest.py puts it on the path.
tests/conftest.py · 3 lines
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
tests/helpers.py · 15 lines
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, prod
from uopc.dtype import dtypes
from uopc.rangeify import lower
from uopc.runtime import run_c, np_dtype
def buf(slot, n, dtype=dtypes.f32): return UOp(Ops.BUFFER, (), ParamArg(slot, dtype, (n,)))
def run_shaped(root: UOp, inputs: list[np.ndarray]) -> np.ndarray:
"""Lower a shaped graph whose BUFFERs are slots 0..len(inputs)-1, run it, return the result."""
n = max(1, prod(root.shape))
out = buf(len(inputs), n, root.dtype)
res = np.zeros(n, np_dtype(root.dtype))
run_c(lower(root, out), {**dict(enumerate(inputs)), len(inputs): res})
return res.reshape(root.shape)
tests/test_w01_uops.py · 128 lines
"""Week 1: UOps, hash consing, rewriting, symbolic, evaluation by rewriting."""
import math, pytest
from uopc.uop import UOp, Ops, ParamArg, CallArg, UOpError, range_
from uopc.dtype import dtypes
from uopc.rewrite import UPat, PatternMatcher, graph_rewrite
from uopc.symbolic import run, simplify, to_python, vmin, vmax, exec_alu
from uopc.uir import parse, to_uir
def stack(*xs): return UOp(Ops.STACK, tuple(x if isinstance(x, UOp) else UOp.const(x) for x in xs))
def param(slot, dtype=dtypes.i32, shape=()): return UOp(Ops.PARAM, (), ParamArg(slot, dtype, shape))
def test_hash_consing():
a, b = UOp.const(1) + UOp.const(2), UOp.const(1) + UOp.const(2)
assert a is b
assert UOp.const(1) is not UOp.const(1.0) is not UOp.const(True)
assert UOp.const(0.0) is not UOp.const(-0.0)
def test_shape_dtype_inference():
s = stack(stack(1, 2, 3), stack(4, 5, 6))
assert s.shape == (2, 3) and s.dtype == dtypes.i32
assert s.index(UOp.const(1)).shape == (3,)
assert (UOp.const(1) < UOp.const(2)).dtype == dtypes.bool
with pytest.raises(UOpError): UOp.const(1) + UOp.const(1.0)
with pytest.raises(UOpError): stack(1, 2) + stack(1, 2, 3)
with pytest.raises(UOpError): stack(1, 2.0)
def test_const_fold_semantics():
assert exec_alu(Ops.FLOORDIV, dtypes.i32, [-7, 2]) == -4 # floor, not truncation
assert exec_alu(Ops.FLOORMOD, dtypes.i32, [-7, 2]) == 1
assert exec_alu(Ops.ADD, dtypes.i32, [2**31 - 1, 1]) == -2**31 # wraps like int32
assert exec_alu(Ops.CAST, dtypes.i32, [-2.7]) == -2 # C truncation
assert exec_alu(Ops.ADD, dtypes.f32, [0.1, 0.2]) == pytest.approx(0.3, abs=1e-7)
assert exec_alu(Ops.MUL, dtypes.f32, [16777216.0, 1.0]) == 16777216.0
assert exec_alu(Ops.ADD, dtypes.f32, [16777216.0, 1.0]) == 16777216.0 # f32 rounding
def test_run_stack_program():
a, b = stack(1, 2, 3), stack(10, 20, 30)
assert to_python(run(a * b + a)) == [11, 42, 93]
assert to_python(run(stack(a, b).index(UOp.const(1)).index(UOp.const(2)))) == 30
w = (a < stack(2, 2, 2)).where(a, b)
assert to_python(run(w)) == [1, 20, 30]
def test_call_inlining_and_scoping():
p0 = param(0)
sq = UOp(Ops.CALL, (p0 * p0, UOp.const(7)), CallArg("sq"))
assert run(sq).arg == 49
# nested: f(x) = sq(x) + x ; the inner PARAM 0 is a different scope from the outer one
inner = UOp(Ops.CALL, (p0 * p0, p0 + 1), CallArg("sq"))
outer = UOp(Ops.CALL, (inner + p0, UOp.const(3)), CallArg("f"))
assert run(outer).arg == 16 + 3
def test_call_type_checks():
with pytest.raises(UOpError): UOp(Ops.CALL, (param(0) + param(1), UOp.const(1)))
def test_symbolic_identities():
x = param(0)
assert simplify(x + 0) is x and simplify(x * 1) is x and simplify(x * 0).arg == 0
assert simplify((x + 3) + 4) is simplify(x + 7)
assert simplify((x * 4) // 4) is x
assert simplify((x * 6 + 4) % 2).arg == 0
def test_symbolic_ranges():
i, j = range_(4, 0), range_(8, 1)
flat = i * 8 + j
assert simplify(flat // 8) is i and simplify(flat % 8) is j
assert simplify((flat // 8) * 8 + flat % 8) is simplify(flat)
assert simplify(j // 8).arg == 0 and simplify(j % 8) is j
assert simplify(i < 4).arg is True and simplify(i < 0).arg is False
assert (vmin(flat), vmax(flat)) == (0, 31)
def test_custom_pattern_matcher():
pm = PatternMatcher([(UPat(Ops.MUL, src=[UPat.var("x"), UPat.cvar("c", dtype=dtypes.i32)]),
lambda x, c: x + x if c.arg == 2 else None)])
x = param(0)
assert graph_rewrite(x * 2, pm) is x + x
def test_deep_graph_no_recursion_error():
x = UOp.const(0)
for i in range(20000): x = x + 1
assert run(x).arg == 20000
def test_uir_roundtrip_and_syllabus_example():
src = """
%0 = buffer 65536 : dtype=f32 slot=0
%1 = buffer 65536 : dtype=f32 slot=1
%2 = reshape %0, (256, 256, 1)
%3 = reshape %1, (256, 1, 256)
%4 = mul %2, %3
%5 = reduce %4 : op=add pop=1
"""
g = parse(src)
assert g.shape == (256, 256)
assert parse(to_uir(g)) is g
def test_uir_evaluates():
g = parse("""
%0 = stack 1.0, 2.0, 3.0
%1 = stack %0, %0
%2 = mul %1, 2
%3 = permute %2, (1, 0)
%4 = reduce %3 : op=add pop=1
""")
assert to_python(run(g)) == [4.0, 8.0, 12.0]
def test_folding_matches_c_semantics():
# f32 inputs are rounded like the C compiler rounds the literal 16777217.0f
assert simplify(UOp.const(16777217.0) + UOp.const(2.0)).arg == 16777218.0
# constants made by reassociation wrap like int32 arithmetic: 65536*65536 == 0 (mod 2^32)
assert simplify((param(0) * 65536) * 65536).arg == 0
# INT_MIN // -1 traps in C: leave it to run time
assert exec_alu(Ops.FLOORDIV, dtypes.i32, [-2**31, -1]) is None
def test_inline_call_does_not_enter_nested_bodies(): # NOTES 2
from uopc.symbolic import inline_call
p0 = param(0)
inner = UOp(Ops.CALL, (p0 * p0, p0 + 1), CallArg("sq"))
outer = UOp(Ops.CALL, (inner + p0, UOp.const(3)), CallArg("f"))
body = inline_call(outer) # inline ONLY the outer call; inner must survive intact
assert body.src[0].op is Ops.CALL and body.src[0].src[0] is p0 * p0
assert run(body).arg == 19
def test_cast_inf_to_int_is_not_folded(): # NOTES 6
assert exec_alu(Ops.CAST, dtypes.i32, [math.inf]) is None
assert exec_alu(Ops.CAST, dtypes.i32, [math.nan]) is None
def test_float_constants_do_not_reassociate(): # NOTES 7
x = UOp(Ops.PARAM, (), ParamArg(0, dtypes.f32))
assert simplify((x + 1.0) + 2.0) is not simplify(x + 3.0)
tests/test_w02_render.py · 45 lines
"""Week 2: rendering loopless C, the runtime, memory (BUFFER/ALLOC/LOAD/STORE/AFTER)."""
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, sink
from uopc.dtype import dtypes
from uopc.render_c import render_c
from uopc.runtime import run_c
from helpers import buf
def test_elementwise_loopless():
a, b, c = buf(0, 4), buf(1, 4), buf(2, 4)
s = sink(*[c.index(i).store(a.index(i).load() * b.index(i).load() + 1.0) for i in range(4)])
A, B, C = np.arange(4, dtype=np.float32), np.full(4, 2, np.float32), np.zeros(4, np.float32)
run_c(s, {0: A, 1: B, 2: C})
np.testing.assert_allclose(C, A * B + 1)
def test_alloc_and_after():
# tmp[0] = a[0] + a[1]; out[0] = tmp[0] * tmp[0] -- AFTER orders the read after the write
a, out = buf(0, 2), buf(1, 1)
tmp = UOp(Ops.ALLOC, (), ParamArg(0, dtypes.f32, (1,)))
w = tmp.index(0).store(a.index(0).load() + a.index(1).load())
v = tmp.after(w).index(0).load()
O = np.zeros(1, np.float32)
run_c(sink(out.index(0).store(v * v)), {0: np.array([1, 2], np.float32), 1: O})
assert O[0] == 9
def test_floordiv_rendering_negative():
a, out = buf(0, 2, dtypes.i32), buf(1, 2, dtypes.i32)
x = a.index(0).load()
s = sink(out.index(0).store(x // a.index(1).load()), out.index(1).store(x % a.index(1).load()))
O = np.zeros(2, np.int32)
run_c(s, {0: np.array([-7, 2], np.int32), 1: O})
assert list(O) == [-4, 1]
assert "floordiv_i32" in render_c(s)
def test_casts_and_where():
a, out = buf(0, 3), buf(1, 3, dtypes.i32)
s = sink(*[out.index(i).store((a.index(i).load() < 0.0).where(UOp.const(-1), a.index(i).load().cast(dtypes.i32))) for i in range(3)])
O = np.zeros(3, np.int32)
run_c(s, {0: np.array([-2.5, 3.7, 0.2], np.float32), 1: O})
assert list(O) == [-1, 3, 0]
def test_run_c_checks_buffer_dtypes():
import pytest
out = buf(0, 4)
with pytest.raises(TypeError): run_c(sink(out.index(0).store(UOp.const(1.0))), {0: np.ones(4)}) # float64
tests/test_w03_loops.py · 54 lines
"""Week 3: RANGE/END. Hand-written gemm and conv in the loop language."""
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, sink, range_
from uopc.dtype import dtypes
from uopc.runtime import run_c
from uopc.linearize import linearize
from helpers import buf
def gemm(N, M, K):
A, B, C = buf(0, N * K), buf(1, K * M), buf(2, N * M)
i, j, k = range_(N, 0), range_(M, 1), range_(K, 2)
acc = UOp(Ops.ALLOC, (), ParamArg(0, dtypes.f32, (1,))).after(i, j)
init = acc.index(0).store(UOp.const(0.0))
acck = acc.after(init, k)
upd = acck.index(0).store(acck.index(0).load() + A.index(i * K + k).load() * B.index(k * M + j).load())
return sink(C.index(i * M + j).store(acc.after(upd.end(k)).index(0).load()).end(i, j))
def test_gemm_loops():
N, M, K = 17, 9, 33
a, b = np.random.randn(N, K).astype(np.float32), np.random.randn(K, M).astype(np.float32)
c = np.zeros((N, M), np.float32)
src = run_c(gemm(N, M, K), {0: a, 1: b, 2: c})
np.testing.assert_allclose(c, a @ b, atol=1e-4)
assert src.count("for (") == 3
def test_conv1d_loops():
# out[o] = sum_k x[o + k] * w[k]
n, kw = 20, 5
X, W, O = buf(0, n), buf(1, kw), buf(2, n - kw + 1)
o, k = range_(n - kw + 1, 0), range_(kw, 1)
acc = UOp(Ops.ALLOC, (), ParamArg(0, dtypes.f32, (1,))).after(o)
init = acc.index(0).store(UOp.const(0.0))
acck = acc.after(init, k)
upd = acck.index(0).store(acck.index(0).load() + X.index(o + k).load() * W.index(k).load())
s = sink(O.index(o).store(acc.after(upd.end(k)).index(0).load()).end(o))
x, w, out = np.random.randn(n).astype(np.float32), np.random.randn(kw).astype(np.float32), np.zeros(n - kw + 1, np.float32)
run_c(s, {0: x, 1: w, 2: out})
np.testing.assert_allclose(out, np.correlate(x, w, "valid"), atol=1e-5)
def test_loop_invariant_code_is_hoisted():
N = 8
A, C = buf(0, N * N), buf(1, N * N)
i, j = range_(N, 0), range_(N, 1)
row = A.index(i * N).load() # depends only on i
s = sink(C.index(i * N + j).store(row + A.index(i * N + j).load()).end(i, j))
lin = linearize(s)
pos = {u: n for n, u in enumerate(lin)}
assert pos[row] < pos[j] # emitted before the j loop opens
def test_missing_end_is_an_error():
import pytest
from uopc.linearize import LinearizeError
out, r = buf(0, 4), range_(4, 0)
with pytest.raises(LinearizeError, match="never closed"): linearize(sink(out.index(r).store(UOp.const(1.0))))
tests/test_w04_rangeify.py · 77 lines
"""Week 4: movement ops, broadcasting, reduction, rangeify, composite PAD/FLIP."""
import numpy as np, pytest
from uopc.uop import UOp, Ops
from uopc.uir import parse
from uopc.composite import pad, flip_axis, anti_eye
from uopc.symbolic import run, to_python
from helpers import buf, run_shaped
def test_gemm_from_movement_ops():
N = 16
a, b = np.random.randn(N, N).astype(np.float32), np.random.randn(N, N).astype(np.float32)
g = parse(f"""
%0 = buffer {N*N} : dtype=f32 slot=0
%1 = buffer {N*N} : dtype=f32 slot=1
%2 = reshape %0, ({N}, 1, {N})
%3 = reshape %1, ({N}, {N})
%4 = permute %3, (1, 0)
%5 = reshape %4, (1, {N}, {N})
%6 = mul %2, %5
%7 = reduce %6 : op=add pop=1""")
np.testing.assert_allclose(run_shaped(g, [a.ravel(), b.ravel()]), a @ b, atol=1e-4)
def test_syllabus_gemm_is_not_a_gemm():
# the example in the syllabus computes A * rowsum(B): a good reminder to test everything
N = 8
a, b = np.random.randn(N, N).astype(np.float32), np.random.randn(N, N).astype(np.float32)
g = parse(f"""
%0 = buffer {N*N} : dtype=f32 slot=0
%1 = buffer {N*N} : dtype=f32 slot=1
%2 = reshape %0, ({N}, {N}, 1)
%3 = reshape %1, ({N}, 1, {N})
%4 = mul %2, %3
%5 = reduce %4 : op=add pop=1""")
np.testing.assert_allclose(run_shaped(g, [a.ravel(), b.ravel()]), a * b.sum(1, keepdims=True), atol=1e-4)
@pytest.mark.parametrize("seed", range(5))
def test_random_movement_chains(seed):
rng = np.random.default_rng(seed)
x = rng.standard_normal(24).astype(np.float32)
u, ref = buf(0, 24).reshape((2, 3, 4)), x.reshape(2, 3, 4)
for _ in range(6):
k = rng.integers(4)
if k == 0:
p = tuple(rng.permutation(len(ref.shape))); u, ref = u.permute(p), ref.transpose(p)
elif k == 1:
arg = tuple((int(b), int(rng.integers(b + 1, s + 1))) for s in ref.shape for b in [rng.integers(0, s)])
u, ref = u.shrink(arg), ref[tuple(slice(b, e) for b, e in arg)]
elif k == 2:
u, ref = u.reshape((1,) + u.shape), ref.reshape((1,) + ref.shape)
u, ref = u.expand((2,) + u.shape[1:]), np.broadcast_to(ref, (2,) + ref.shape[1:])
else:
u, ref = u.reshape((int(np.prod(ref.shape)),)), ref.reshape(-1)
np.testing.assert_allclose(run_shaped(u, [x]), ref)
def test_reduce_max_and_nested():
x = np.random.randn(6, 5).astype(np.float32)
u = buf(0, 30).reshape((6, 5))
np.testing.assert_allclose(run_shaped(u.reduce(Ops.MAX, 1), [x.ravel()]), x.max(1))
# x - max(x) broadcast back: the reduce becomes a sibling loop, computed once per row (no kernelizer needed)
m = u.reduce(Ops.MAX, 1).reshape((6, 1)).expand((6, 5))
np.testing.assert_allclose(run_shaped(u + -m, [x.ravel()]), x - x.max(1, keepdims=True))
np.testing.assert_allclose(run_shaped(u.reduce(Ops.ADD, 2), [x.ravel()]), x.sum(), rtol=1e-5)
def test_stack_and_index_in_rangeify():
x = np.random.randn(4).astype(np.float32)
u = buf(0, 4)
s = UOp(Ops.STACK, (u, u * u, u + u))
np.testing.assert_allclose(run_shaped(s, [x]), np.stack([x, x * x, x + x]))
np.testing.assert_allclose(run_shaped(s.index(UOp.const(1)), [x]), x * x)
def test_pad_flip_composites():
x = np.random.randn(3, 4).astype(np.float32)
u = buf(0, 12).reshape((3, 4))
np.testing.assert_allclose(run_shaped(pad(u, ((1, 5), (2, 0))), [x.ravel()]), np.pad(x, ((1, 5), (2, 0))))
np.testing.assert_allclose(run_shaped(flip_axis(u, 1), [x.ravel()]), x[:, ::-1])
np.testing.assert_allclose(run_shaped(flip_axis(u, 0), [x.ravel()]), x[::-1])
assert to_python(run(anti_eye(3, u.dtype))) == [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
tests/test_w05_tensor.py · 111 lines
"""Week 5: the Tensor frontend, kernelization (fusion/fission), autodiff."""
import math
import numpy as np, pytest
from uopc.tensor import Tensor
from uopc.kernelize import make_plan
from uopc.nn import Linear, SGD, AdamW, parameters
rng = np.random.default_rng(0)
def r(*s): return rng.standard_normal(s).astype(np.float32)
@pytest.mark.parametrize("name,fn,ref", [
("add_broadcast", lambda a, b: a + b[0], lambda a, b: a + b[0]),
("mul_sub_div", lambda a, b: (a * b - 3) / (b.abs() + 1), lambda a, b: (a * b - 3) / (np.abs(b) + 1)),
("exp_log", lambda a, b: (a.exp() + 1).log(), lambda a, b: np.log(np.exp(a) + 1)),
("sum_axes", lambda a, b: a.sum((0, 2)), lambda a, b: a.sum((0, 2))),
("max_keepdim", lambda a, b: a.max(1, keepdim=True), lambda a, b: a.max(1, keepdims=True)),
("mean_var", lambda a, b: a.var(-1), lambda a, b: a.var(-1)),
("softmax", lambda a, b: a.softmax(-1), lambda a, b: np.exp(a - a.max(-1, keepdims=True)) / np.exp(a - a.max(-1, keepdims=True)).sum(-1, keepdims=True)),
("matmul_batched", lambda a, b: a @ b.transpose(), lambda a, b: a @ b.transpose(0, 2, 1)),
("where", lambda a, b: (a > b).where(a, b * 2), lambda a, b: np.where(a > b, a, b * 2)),
("getitem_cat", lambda a, b: Tensor.cat([a[:, 1:], b[:, :1]], 1), lambda a, b: np.concatenate([a[:, 1:], b[:, :1]], 1)),
("pad_flip", lambda a, b: a.pad(((0, 0), (1, 2), (0, 1))).flip(1), lambda a, b: np.pad(a, ((0, 0), (1, 2), (0, 1)))[:, ::-1]),
])
def test_ops(name, fn, ref):
a, b = r(2, 3, 4), r(2, 3, 4)
np.testing.assert_allclose(fn(Tensor(a), Tensor(b)).numpy(), ref(a, b), rtol=1e-4, atol=1e-5)
def numgrad(f, x, eps=1e-3):
g = np.zeros_like(x, dtype=np.float64)
for i in np.ndindex(x.shape):
xp, xm = x.astype(np.float64), x.astype(np.float64)
xp[i] += eps; xm[i] -= eps
g[i] = (f(xp) - f(xm)) / (2 * eps)
return g
@pytest.mark.parametrize("name,fn,npfn", [
("tanh_sum", lambda t: t.tanh().sum(), lambda x: np.tanh(x).sum()),
("softmax_xent", lambda t: t.log_softmax(-1)[:, 1].sum(), lambda x: (x - np.log(np.exp(x).sum(-1, keepdims=True)))[:, 1].sum()),
("max", lambda t: (t.max(-1) * t.sum(0).max()).sum(), lambda x: (x.max(-1) * x.sum(0).max()).sum()),
("pad_shrink", lambda t: t.pad(((1, 1), (0, 2)))[1:3, 1:].square().sum(),
lambda x: (np.pad(x, ((1, 1), (0, 2)))[1:3, 1:] ** 2).sum()),
("sqrt_recip", lambda t: (t.square() + 1).sqrt().reciprocal().sum(), lambda x: (1 / np.sqrt(x ** 2 + 1)).sum()),
("matmul", lambda t: (t @ t.T).sin().sum(), lambda x: np.sin(x @ x.T).sum()),
])
def test_gradients(name, fn, npfn):
x = r(3, 4)
t = Tensor(x, requires_grad=True)
fn(t).backward([t])
np.testing.assert_allclose(t.grad.numpy(), numgrad(npfn, x), rtol=2e-2, atol=2e-3)
def test_kernelize_fuses_elementwise_and_splits_reductions():
a = Tensor(r(64, 64))
plan, _ = make_plan([((a * 2 + 1).relu() * 3).uop])
assert len(plan.kernels) == 1 # a chain of elementwise ops: one kernel
plan, _ = make_plan([a.softmax(-1).uop])
assert len(plan.kernels) == 3 # max, sum, and the elementwise output
def test_training_reduces_loss():
Tensor.manual_seed(1)
X = r(64, 8); Y = (X.sum(1) > 0).astype(np.int32)
l1, l2 = Linear(8, 16), Linear(16, 2)
opt = AdamW(parameters([l1, l2]), lr=1e-2)
losses = []
for _ in range(40):
loss = l2(l1(Tensor(X)).relu()).cross_entropy(Tensor(Y))
opt.zero_grad(); loss.backward(opt.params); opt.step(loss)
losses.append(loss.item())
assert losses[-1] < 0.3 * losses[0]
def test_sgd_momentum_step():
w = Tensor(np.array([1.0, -2.0], np.float32), requires_grad=True)
opt = SGD([w], lr=0.1, momentum=0.9)
for _ in range(2):
loss = w.square().sum()
opt.zero_grad(); loss.backward(opt.params); opt.step()
# v1 = 2w0, w1 = w0 - .1*v1 = .8 w0; v2 = .9 v1 + 2 w1 = 3.4 w0, w2 = .8 w0 - .34 w0 = .46 w0
np.testing.assert_allclose(w.numpy(), np.array([0.46, -0.92]), rtol=1e-5)
def test_deleted_reads_do_not_shift_kernel_arguments():
t = Tensor(np.array([1, 2, 3], np.int32))
assert (t * 0).numpy().tolist() == [0, 0, 0] # simplifies to a constant: t is no longer read...
assert t.numpy().tolist() == [1, 2, 3] # ...and must not be overwritten
def test_python_float_promotes_int_tensors():
t = Tensor(np.array([1, 2, 3], np.int32))
np.testing.assert_allclose((t * 0.5).numpy(), [0.5, 1.0, 1.5])
assert t.mean().item() == 2.0
def test_item_before_backward_and_detach():
w = Tensor(np.array([1.0, 2.0], np.float32), requires_grad=True)
loss = (w * w).sum()
loss.item() # realizing the loss must keep its graph
loss.backward([w])
np.testing.assert_allclose(w.grad.numpy(), [2.0, 4.0])
w.grad = None
(w.detach() * w).sum().backward([w]) # d/dw (c * w) with c = w held constant
np.testing.assert_allclose(w.grad.numpy(), [1.0, 2.0])
def test_adamw_steps_reuse_one_plan(): # NOTES 17
from uopc import kernelize
Tensor.manual_seed(0)
lin = Linear(4, 2)
opt = AdamW([lin.weight, lin.bias], lr=1e-2)
x = Tensor(np.random.randn(8, 4).astype(np.float32))
def step():
loss = lin(x).square().mean()
opt.zero_grad(); loss.backward(opt.params); opt.step(loss)
step(); step()
n = len(kernelize._plans)
for _ in range(3): step()
assert len(kernelize._plans) == n # no new plans: nothing recompiled
tests/test_w06_opt.py · 72 lines
"""Week 6: upcast, unroll, threads, layout. Optimized kernels must compute the same thing."""
import numpy as np, pytest
from uopc.uop import Ops
from uopc.rangeify import lower, Ctx, Opts
from uopc.render_c import render_c
from uopc.runtime import Program
from uopc.opt import hand_opt_choice, layout_pass
from helpers import buf
def gemm_graph(M, N, K, bt=False):
A, B = buf(0, M * K), buf(1, K * N)
b = B.reshape((N, K)) if bt else B.reshape((K, N)).permute((1, 0))
return (A.reshape((M, 1, K)).expand((M, N, K)) * b.reshape((1, N, K)).expand((M, N, K))).reduce(Ops.ADD, 1)
@pytest.mark.parametrize("opts", [
Opts(), Opts(thread=0), Opts(upcast={1: 4}), Opts(upcast={0: 4, 1: 4}, thread=0),
Opts(unroll=4), Opts(upcast={0: 2, 1: 8}, unroll=4, thread=1), Opts(upcast={0: 32})])
def test_gemm_opts_correct(opts):
M, N, K = 32, 48, 64
a, b = np.random.randn(M, K).astype(np.float32), np.random.randn(K, N).astype(np.float32)
c = np.zeros(M * N, np.float32)
src = render_c(lower(gemm_graph(M, N, K), buf(2, M * N), Ctx(opts=opts)), "k")
Program(src, "k")(a, b, c)
np.testing.assert_allclose(c.reshape(M, N), a @ b, rtol=1e-4, atol=1e-4)
if opts.thread is not None: assert "#pragma omp parallel for" in src
def test_upcast_shares_reduce_loop():
src = render_c(lower(gemm_graph(8, 8, 8), buf(2, 64), Ctx(opts=Opts(upcast={0: 2, 1: 2}))), "k")
assert src.count("for (") == 3 # 2 output loops + ONE shared reduce loop
assert src.count("acc") >= 4
def test_reduce_unroll_has_no_inner_loop():
x = buf(0, 16)
src = render_c(lower(x.reduce(Ops.ADD, 1), buf(1, 1), Ctx(opts=Opts(unroll=16))), "k")
assert "for (" not in src
def test_layout_pass_stages_transposed_operand():
g = gemm_graph(64, 64, 64)
(g2,) = layout_pass([g])
assert any(u.op is Ops.STAGE for u in g2.toposort())
(g3,) = layout_pass([gemm_graph(64, 64, 64, bt=True)]) # already contiguous: nothing to do
assert not any(u.op is Ops.STAGE for u in g3.toposort())
def test_hand_opt_choice():
o = hand_opt_choice(gemm_graph(64, 64, 64))
assert o.upcast and o.thread is not None
def max_loop_depth(src: str) -> int:
return max((len(l) - len(l.lstrip())) // 2 for l in src.splitlines() if l.lstrip().startswith("for ("))
def test_fused_attention_reduces_are_not_nested(): # NOTES 12
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent / "examples"))
from flash_attention import attention, FlashOpt
from uopc.tensor import Tensor
from uopc.kernelize import make_plan
N, D = 64, 16
q, k, v = (Tensor(np.random.randn(N, D).astype(np.float32)) for _ in range(3))
plan, _ = make_plan([attention(q, k, v).uop], fuse_reduce=True, opt=FlashOpt)
# loops: i (thread), then a reduce over keys j, inside it a dot product over D: depth 3, never 4
assert max(max_loop_depth(k.src) for k in plan.kernels) <= 3
@pytest.mark.parametrize("tile", [{0: 8, 1: 16}, {1: 32}, {0: 16}])
def test_cache_blocking_correct(tile):
M, N, K = 32, 64, 48
a, b = np.random.randn(M, K).astype(np.float32), np.random.randn(K, N).astype(np.float32)
c = np.zeros(M * N, np.float32)
src = render_c(lower(gemm_graph(M, N, K), buf(2, M * N), Ctx(opts=Opts(upcast={0: 4, 1: 4}, thread=0, tile=tile))), "k")
Program(src, "k")(a, b, c)
np.testing.assert_allclose(c.reshape(M, N), a @ b, rtol=1e-4, atol=1e-4)
assert src.count("for (") == 3 + len(tile) # one extra block loop per tiled axis
tests/test_w07_gpu.py · 66 lines
"""Weeks 7-8: GPU mapping. Every kernel is checked on the CPU through the emulator;
the CUDA test only runs when a GPU and NVRTC are available."""
import numpy as np, pytest
from uopc.uop import UOp, Ops, AxisType
from uopc.dtype import dtypes
from uopc.runtime import Program
from uopc.rangeify import Opts
from uopc.gpu import gpu_lower, gpu_opt_choice, emulate, render_cuda, launch_dims, smem_gemm_ir, wmma_gemm, match_gemm
from helpers import buf
def gemm_graph(M, N, K):
A, B = buf(0, M * K), buf(1, K * N)
return (A.reshape((M, 1, K)).expand((M, N, K)) * B.reshape((K, N)).permute((1, 0)).reshape((1, N, K)).expand((M, N, K))).reduce(Ops.ADD, 1)
@pytest.mark.parametrize("opts", [Opts(gpu=True), Opts(gpu=True, local={0: 4, 1: 8}),
Opts(gpu=True, local={1: 16}, upcast={0: 2, 1: 2})])
def test_gpu_gemm_emulated(opts):
M, N, K = 32, 64, 48
a, b = np.random.randn(M, K).astype(np.float32), np.random.randn(K, N).astype(np.float32)
c = np.zeros(M * N, np.float32)
ast = gpu_lower(gemm_graph(M, N, K), buf(2, M * N), opts)
Program(emulate(ast), "kernel")(a, b, c)
np.testing.assert_allclose(c.reshape(M, N), a @ b, atol=1e-4)
src = render_cuda(ast)
assert "__global__" in src and ("threadIdx" in src) == bool(opts.local)
def test_gpu_softmax_like_emulated():
x = np.random.randn(16, 32).astype(np.float32)
X = buf(0, 16 * 32).reshape((16, 32))
m = X.reduce(Ops.MAX, 1).reshape((16, 1)).expand((16, 32))
body = UOp(Ops.EXP2, (X + -m,))
out = np.zeros(16 * 32, np.float32)
ast = gpu_lower(body, buf(1, 16 * 32))
Program(emulate(ast), "kernel")(x, out)
np.testing.assert_allclose(out.reshape(16, 32), np.exp2(x - x.max(1, keepdims=True)), rtol=1e-5)
def test_launch_dims_and_auto_choice():
o = gpu_opt_choice(gemm_graph(64, 64, 64))
ast = gpu_lower(gemm_graph(64, 64, 64), buf(2, 64 * 64), o)
grid, block = launch_dims(ast)
assert np.prod(grid) * np.prod(block) * np.prod(list(o.upcast.values()) or [1]) == 64 * 64
def test_smem_gemm_renders():
src = render_cuda(smem_gemm_ir(64, 16), "smem_gemm")
assert src.count("__syncthreads();") == 2 and "__shared__ float smem0[256];" in src
assert launch_dims(smem_gemm_ir(64, 16)) == ((4, 4, 1), (16, 16, 1))
def test_wmma():
assert match_gemm(gemm_graph(32, 32, 32)) == (32, 32, 32)
assert "mma_sync" in wmma_gemm(32, 32, 32)
def _has_cuda():
try:
from uopc.cuda import CUDA
CUDA(); return True
except Exception: return False
@pytest.mark.skipif(not _has_cuda(), reason="no CUDA device / NVRTC")
def test_cuda_smem_gemm_on_gpu():
from uopc.cuda import run_cuda
N = 64
a, b, c = np.random.randn(N, N).astype(np.float32), np.random.randn(N, N).astype(np.float32), np.zeros(N * N, np.float32)
ast = smem_gemm_ir(N, 16)
grid, block = launch_dims(ast)
run_cuda(render_cuda(ast, "smem_gemm"), "smem_gemm", grid, block, a, b, c)
np.testing.assert_allclose(c.reshape(N, N), a @ b, atol=1e-3)
tests/test_w09_models.py · 38 lines
"""Weeks 9-10: a tiny GPT trains (loss goes down) and its pieces match numpy."""
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).parent.parent / "examples"))
from gpt import GPT, Attention
from uopc.tensor import Tensor
from uopc.nn import LayerNorm, AdamW, parameters
def test_layernorm_matches_numpy():
x = np.random.randn(4, 8).astype(np.float32)
y = LayerNorm(8)(Tensor(x)).numpy()
ref = (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + 1e-5)
np.testing.assert_allclose(y, ref, atol=1e-5)
def test_attention_is_causal():
Tensor.manual_seed(0)
att = Attention(8, 2, 6)
x = np.random.randn(1, 6, 8).astype(np.float32)
y1 = att(Tensor(x)).numpy()
x2 = x.copy(); x2[0, 4:] += 10 # change the future
y2 = att(Tensor(x2)).numpy()
np.testing.assert_allclose(y1[0, :4], y2[0, :4], atol=1e-5) # the past doesn't see it
def test_tiny_gpt_learns():
Tensor.manual_seed(0)
V, T = 11, 8
model = GPT(V, T, 16, 2, 1)
opt = AdamW(parameters(model), lr=1e-2)
seq = np.arange(64) % V # a trivially predictable sequence
x = Tensor(np.stack([seq[i:i + T] for i in range(4)]).astype(np.int32))
y = Tensor(np.stack([seq[i + 1:i + T + 1] for i in range(4)]).astype(np.int32))
losses = []
for _ in range(30):
loss = model(x).reshape(-1, V).cross_entropy(y.reshape(-1))
opt.zero_grad(); loss.backward(opt.params); opt.step(loss)
losses.append(loss.item())
assert losses[-1] < 0.25 * losses[0]
D.3 Examples: examples/
examples/datasets.py · 29 lines
"""Tiny dataset loaders. Files are cached in ~/.cache/uopc/data."""
import gzip, urllib.request
from pathlib import Path
import numpy as np
DATA = Path.home() / ".cache" / "uopc" / "data"
MNIST_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"
SHAKESPEARE_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
def fetch(url: str, name: str) -> Path:
DATA.mkdir(parents=True, exist_ok=True)
p = DATA / name
if not p.exists():
print(f"downloading {url}")
urllib.request.urlretrieve(url, p)
return p
def mnist():
"""(X_train (60000, 784) f32 in [0,1], Y_train (60000,) i32, X_test, Y_test)"""
def images(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=16).reshape(-1, 784).astype(np.float32) / 255
def labels(n): return np.frombuffer(gzip.open(fetch(MNIST_URL + n, n)).read(), np.uint8, offset=8).astype(np.int32)
return (images("train-images-idx3-ubyte.gz"), labels("train-labels-idx1-ubyte.gz"),
images("t10k-images-idx3-ubyte.gz"), labels("t10k-labels-idx1-ubyte.gz"))
def shakespeare() -> str:
return fetch(SHAKESPEARE_URL, "tinyshakespeare.txt").read_text()
examples/mnist_mlp.py · 43 lines
"""Week 5 milestone: train an MLP on MNIST with our own compiler.
python examples/mnist_mlp.py # ~94% test accuracy in 600 steps (~15 ms/step)
DEBUG=2 python examples/mnist_mlp.py # print every kernel with its time
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
import numpy as np
from datasets import mnist
from uopc.tensor import Tensor
from uopc.nn import Linear, AdamW, parameters
class MLP:
def __init__(self):
self.l1, self.l2 = Linear(784, 128), Linear(128, 10)
def __call__(self, x: Tensor) -> Tensor: return self.l2(self.l1(x).relu())
def main(steps=600, bs=128):
Xtr, Ytr, Xte, Yte = mnist()
Tensor.manual_seed(0)
model = MLP()
opt = AdamW(parameters(model), lr=1e-3)
rng = np.random.default_rng(0)
t0 = time.perf_counter()
for step in range(steps):
idx = rng.integers(0, len(Xtr), bs)
x, y = Tensor(Xtr[idx]), Tensor(Ytr[idx])
loss = model(x).cross_entropy(y)
opt.zero_grad()
loss.backward(opt.params)
opt.step(loss)
if step % 100 == 0 or step == steps - 1:
print(f"step {step:4d} loss {loss.item():.4f} {(time.perf_counter() - t0) / (step + 1) * 1e3:.1f} ms/step")
pred = np.concatenate([model(Tensor(Xte[i:i + 1000])).argmax(-1).numpy() for i in range(0, len(Xte), 1000)])
acc = (pred == Yte).mean()
print(f"test accuracy {acc * 100:.2f}%")
return acc
if __name__ == "__main__": main()
examples/bench_gemm.py · 59 lines
"""Week 6: GEMM GFLOP/s, step by step, vs numpy (which calls an optimized BLAS).
python examples/bench_gemm.py 1024
B strided C[i,j] = sum_k A[i,k] B[k,j]: the reduce walks down a column of B (stride N)
B^T same math, B stored transposed: both operands are contiguous along k
"""
import sys, time
import numpy as np
from uopc.uop import UOp, Ops, ParamArg
from uopc.dtype import dtypes
from uopc.rangeify import lower, Ctx, Opts
from uopc.render_c import render_c
from uopc.runtime import Program
from uopc.tensor import Tensor
def gemm_graph(N, b_transposed=False):
A = UOp(Ops.BUFFER, (), ParamArg(0, dtypes.f32, (N * N,)))
B = UOp(Ops.BUFFER, (), ParamArg(1, dtypes.f32, (N * N,)))
a = A.reshape((N, 1, N)).expand((N, N, N))
bt = B.reshape((N, N)) if b_transposed else B.reshape((N, N)).permute((1, 0))
return (a * bt.reshape((1, N, N)).expand((N, N, N))).reduce(Ops.ADD, 1)
def bench(fn, reps=5):
fn(); best = float("inf")
for _ in range(reps):
t0 = time.perf_counter(); fn(); best = min(best, time.perf_counter() - t0)
return best
if __name__ == "__main__":
N = int(sys.argv[1]) if len(sys.argv) > 1 else 1024
a, b = np.random.randn(N, N).astype(np.float32), np.random.randn(N, N).astype(np.float32)
bt = np.ascontiguousarray(b.T)
c = np.zeros(N * N, np.float32)
out = UOp(Ops.BUFFER, (), ParamArg(2, dtypes.f32, (N * N,)))
flop = 2 * N ** 3
print(f"N={N} ({flop / 1e9:.2f} GFLOP per GEMM)")
rows = [("B strided", False, "naive", Opts()),
("B strided", False, "threads", Opts(thread=0)),
("B strided", False, "4x4 tile + threads", Opts(upcast={0: 4, 1: 4}, thread=0)),
("B^T", True, "naive", Opts()),
("B^T", True, "threads", Opts(thread=0)),
("B^T", True, "4x4 tile", Opts(upcast={0: 4, 1: 4})),
("B^T", True, "4x4 tile + threads", Opts(upcast={0: 4, 1: 4}, thread=0)),
("B^T", True, "+ 64x64 cache blocks", Opts(upcast={0: 4, 1: 4}, thread=0, tile={0: 64, 1: 64}))]
for layout, transposed, name, o in rows:
prg = Program(render_c(lower(gemm_graph(N, transposed), out, Ctx(opts=o)), "gemm"), "gemm")
rhs = bt if transposed else b
t = bench(lambda: prg(a, rhs, c))
err = np.abs(c.reshape(N, N) - a @ b).max()
print(f" {layout:10s} {name:20s} {flop / t / 1e9:8.1f} GFLOP/s max err {err:.1e}")
A, B = Tensor(a), Tensor(b)
t = bench(lambda: (A @ B).realize())
print(f" {'Tensor A@B':10s} {'(layout pass+opt)':20s} {flop / t / 1e9:8.1f} GFLOP/s (end to end, includes the transpose copy)")
t = bench(lambda: a @ b)
print(f" {'numpy':31s} {flop / t / 1e9:8.1f} GFLOP/s")
examples/flash_attention.py · 107 lines
"""Week 6: three ways to compute attention, out = softmax(Q K^T / sqrt(D)) V.
1. unfused -- the Tensor code, kernelized normally: S = QK^T is materialized (N*N floats)
2. fused -- the same Tensor code with fuse_reduce=True: reduces nest inside each other,
nothing N*N ever touches memory; QK^T is recomputed by every pass (max, sum,
output) and, in the output pass, once per 16-wide tile of the head dimension
3. online -- FlashAttention's online softmax, written by hand in the week-3 loop IR:
one pass over keys, running max m, running sum l, rescaled accumulator o[D]
python examples/flash_attention.py 1024 64
"""
import math, sys, time
import numpy as np
from uopc.uop import UOp, Ops, ParamArg, AxisType, sink, range_
from uopc.dtype import dtypes
from uopc.render_c import render_c
from uopc.runtime import Program
from uopc.rangeify import lower, Ctx, Opts
from uopc.tensor import Tensor
LOG2E = 1 / math.log(2)
def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
return (q @ k.T * (1 / math.sqrt(q.shape[-1]))).softmax(-1) @ v
def online_attention_ir(N: int, D: int) -> UOp:
"""FlashAttention-style online softmax as a UOp loop program. One query row per thread."""
f32 = dtypes.f32
Q, K, V, OUT = (UOp(Ops.BUFFER, (), ParamArg(s, f32, (N * D,))) for s in range(4))
alloc = lambda slot, n: UOp(Ops.ALLOC, (), ParamArg(slot, f32, (n,)))
M, L, O, S = alloc(0, 1), alloc(1, 1), alloc(2, D), alloc(3, 1)
scale = (1 / math.sqrt(D)) * LOG2E # work in base 2: exp(x) = exp2(x * log2 e)
i = range_(N, 0, AxisType.THREAD)
Mi, Li, Oi = M.after(i), L.after(i), O.after(i)
init_m = Mi.index(0).store(UOp.const(-math.inf))
init_l = Li.index(0).store(UOp.const(0.0))
d0 = range_(D, 1)
init_o = Oi.index(d0).store(UOp.const(0.0)).end(d0)
j = range_(N, 2)
# s = (q_i . k_j) * scale
Sj = S.after(i, j)
init_s = Sj.index(0).store(UOp.const(0.0))
e = range_(D, 3)
Se = Sj.after(init_s, e)
dot = Se.index(0).store(Se.index(0).load() + Q.index(i * D + e).load() * K.index(j * D + e).load()).end(e)
s = Sj.after(dot).index(0).load() * scale
# online softmax update
Mj, Lj = Mi.after(init_m, j), Li.after(init_l, j)
m_old = Mj.index(0).load()
m_new = m_old.maximum(s)
corr = UOp(Ops.EXP2, (m_old - m_new,)) # rescale what we accumulated so far
p = UOp(Ops.EXP2, (s - m_new,))
st_l = Lj.index(0).store(Lj.index(0).load() * corr + p)
d = range_(D, 4)
Od = Oi.after(init_o, j)
st_o = Od.index(d).store(Od.index(d).load() * corr + p * V.index(j * D + d).load()).end(d)
st_m = Mj.index(0).store(m_new)
loop_j = UOp(Ops.SINK, (st_l, st_o, st_m)).end(j)
# normalize
d2 = range_(D, 5)
out = OUT.index(i * D + d2).store(Oi.after(loop_j).index(d2).load() * Li.after(loop_j).index(0).load().recip())
return sink(out.end(d2, i))
class FlashOpt:
"""fuse_reduce kernels: unroll the head dim (up to 16) so each QK^T value is shared by 16 accumulators."""
from uopc.opt import layout_pass as layout
@staticmethod
def lower(body, out):
return lower(body, out, Ctx(opts=Opts(upcast={len(body.shape) - 1: min(body.shape[-1], 16)}, thread=0)))
def bench(fn, reps=3):
fn(); best = math.inf
for _ in range(reps):
t0 = time.perf_counter(); fn(); best = min(best, time.perf_counter() - t0)
return best
if __name__ == "__main__":
N = int(sys.argv[1]) if len(sys.argv) > 1 else 1024
D = int(sys.argv[2]) if len(sys.argv) > 2 else 64
rng = np.random.default_rng(0)
q, k, v = (rng.standard_normal((N, D)).astype(np.float32) for _ in range(3))
s = q @ k.T / math.sqrt(D)
p = np.exp(s - s.max(-1, keepdims=True))
ref = (p / p.sum(-1, keepdims=True)) @ v
def run_tensor():
return attention(Tensor(q), Tensor(k), Tensor(v)).numpy()
t = bench(run_tensor)
print(f"unfused {t*1e3:8.2f} ms err {np.abs(run_tensor() - ref).max():.1e} (materializes {N*N*4/1e6:.1f} MB of scores)")
Tensor.fuse_reduce, old = True, Tensor.opt
Tensor.opt = FlashOpt
t = bench(run_tensor)
print(f"fused {t*1e3:8.2f} ms err {np.abs(run_tensor() - ref).max():.1e} (no N*N buffer, QK^T recomputed)")
Tensor.fuse_reduce, Tensor.opt = False, old
prg = Program(render_c(online_attention_ir(N, D), "flash"), "flash")
out = np.zeros(N * D, np.float32)
t = bench(lambda: prg(q, k, v, out))
print(f"online {t*1e3:8.2f} ms err {np.abs(out.reshape(N, D) - ref).max():.1e} (one pass, running max/sum)")
examples/gpt.py · 106 lines
"""Weeks 9-10: train a GPT on Tiny Shakespeare with our compiler, then sample from it.
python examples/gpt.py # small model, a few minutes on a laptop CPU
python examples/gpt.py --steps 2000 --layers 4 --dim 128
Everything -- embeddings, causal attention, LayerNorm, GELU, AdamW, the backward pass --
is graph building on the Tensor API. Every kernel it runs was generated by uopc.
"""
import argparse, math, sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
import numpy as np
from datasets import shakespeare
from uopc.tensor import Tensor
from uopc.nn import Linear, LayerNorm, Embedding, AdamW, parameters
class Attention:
def __init__(self, dim: int, heads: int, ctx: int):
self.qkv, self.proj, self.heads = Linear(dim, 3 * dim), Linear(dim, dim), heads
mask = np.triu(np.full((ctx, ctx), -1e9, np.float32), 1) # -inf above the diagonal: no peeking
self.mask = Tensor(mask)
def __call__(self, x: Tensor) -> Tensor:
B, T, C = x.shape
H, D = self.heads, C // self.heads
qkv = self.qkv(x).reshape(B, T, 3, H, D)
q, k, v = (qkv[:, :, i].permute(0, 2, 1, 3) for i in range(3)) # (B, H, T, D)
att = (q @ k.transpose()) * (1 / math.sqrt(D)) + self.mask[:T, :T].reshape(1, 1, T, T)
y = att.softmax(-1) @ v # (B, H, T, D)
return self.proj(y.permute(0, 2, 1, 3).reshape(B, T, C))
class Block:
def __init__(self, dim, heads, ctx):
self.ln1, self.attn, self.ln2 = LayerNorm(dim), Attention(dim, heads, ctx), LayerNorm(dim)
self.fc, self.out = Linear(dim, 4 * dim), Linear(4 * dim, dim)
def __call__(self, x: Tensor) -> Tensor:
x = x + self.attn(self.ln1(x))
return x + self.out(self.fc(self.ln2(x)).gelu())
class GPT:
def __init__(self, vocab, ctx, dim, heads, layers):
self.tok, self.pos = Embedding(vocab, dim), Tensor.randn(ctx, dim) * 0.02
self.pos = self.pos.realize(); self.pos.requires_grad = True
self.blocks = [Block(dim, heads, ctx) for _ in range(layers)]
self.ln, self.head = LayerNorm(dim), Linear(dim, vocab)
def __call__(self, idx: Tensor) -> Tensor:
B, T = idx.shape
x = self.tok(idx) + self.pos[:T].reshape(1, T, -1)
for b in self.blocks: x = b(x)
return self.head(self.ln(x))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=600)
ap.add_argument("--batch", type=int, default=16)
ap.add_argument("--ctx", type=int, default=64)
ap.add_argument("--dim", type=int, default=64)
ap.add_argument("--heads", type=int, default=4)
ap.add_argument("--layers", type=int, default=2)
ap.add_argument("--lr", type=float, default=3e-3)
a = ap.parse_args()
text = shakespeare()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in text], np.int32)
split = int(0.9 * len(data))
train, val = data[:split], data[split:]
rng = np.random.default_rng(0)
def batch(src):
ix = rng.integers(0, len(src) - a.ctx - 1, a.batch)
return Tensor(np.stack([src[i:i + a.ctx] for i in ix])), Tensor(np.stack([src[i + 1:i + a.ctx + 1] for i in ix]))
Tensor.manual_seed(0)
model = GPT(len(chars), a.ctx, a.dim, a.heads, a.layers)
params = parameters(model)
print(f"{sum(p.numel() for p in params) / 1e3:.1f}k parameters, vocab {len(chars)}")
opt = AdamW(params, lr=a.lr, wd=0.01)
t0 = time.perf_counter()
for step in range(a.steps):
x, y = batch(train)
loss = model(x).reshape(-1, len(chars)).cross_entropy(y.reshape(-1))
opt.zero_grad()
loss.backward(opt.params)
opt.step(loss)
if step % 50 == 0 or step == a.steps - 1:
vx, vy = batch(val)
vloss = model(vx).reshape(-1, len(chars)).cross_entropy(vy.reshape(-1)).item()
print(f"step {step:5d} train {loss.item():.3f} val {vloss:.3f} {(time.perf_counter() - t0) / (step + 1) * 1e3:6.0f} ms/step")
# sample: feed the last ctx characters, draw the next one from the softmax
idx = [stoi["\n"]]
for _ in range(300):
ctx = idx[-a.ctx:]
ctx = ctx + [0] * (a.ctx - len(ctx)) # fixed shape -> one compiled graph
logits = model(Tensor(np.array([ctx], np.int32))).numpy()[0, min(len(idx), a.ctx) - 1]
p = np.exp(logits - logits.max()); p /= p.sum()
idx.append(int(rng.choice(len(chars), p=p)))
print("".join(chars[i] for i in idx))
if __name__ == "__main__": main()
D.4 Project files
pyproject.toml · 22 lines
[project]
name = "uopc"
version = "0.1.0"
description = "A machine learning compiler built from scratch, one week at a time"
requires-python = ">=3.11"
dependencies = ["numpy"]
[project.optional-dependencies]
dev = ["pytest", "markdown", "pygments"]
[build-system]
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[tool.setuptools]
packages = ["uopc", "minicc", "tensorjit", "labs", "labs.kernels"]
[tool.setuptools.package-data]
"labs.kernels" = ["*.cuh"]
[tool.pytest.ini_options]
testpaths = ["tests"]