tinygrad-notes

Pattern matching and graph rewriting

All tutorials · Updated September 21, 2026 · tinygrad 8ad8f73

tinygrad expresses many compiler transformations as a pattern plus a replacement function. A pattern describes a graph shape; the function returns a replacement UOp or None when the rule does not apply.

The current imports are from tinygrad.uop.ops. Ops names operations, UOp constructs graphs, UPat constructs patterns, and PatternMatcher groups rules.

A complete rewrite

from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite

x = UOp.variable("x", 0, 10)
expression = (x + 0) * 2
remove_zero = PatternMatcher([
  (UPat(Ops.ADD, src=(UPat.var("value"), UPat.const(0))),
   lambda value: value),
])
rewritten = graph_rewrite(expression, remove_zero)
assert rewritten.src[0] is x
for value in range(11):
  assert rewritten.sym_infer({"x": value}) == 2 * value
print(rewritten.render())

The named pattern value binds a matched UOp to the callback argument. The rule removes addition by zero from within a larger graph; it does not replace the multiplication at the root.

This rule is deliberately restricted to the expression shown. A production algebraic rewrite must account for dtypes, overflow, floating-point signed zero/NaN behavior, validity, and any relevant side effects. Passing an integer example is not a proof that the same identity is safe for every operation/type combination.

Match a node versus rewrite a graph

PatternMatcher.rewrite attempts a replacement for the node supplied to it. graph_rewrite traverses a graph, propagates rewritten sources, and applies the matcher according to the requested traversal mode.

A rule can therefore become applicable after one of its inputs changes. Traversal order, repetition, and termination matter: a pair of inverse rules can repeatedly undo one another. Matchers are organized into compiler passes so each pass has a defined purpose and expected input representation.

Do not mutate shared nodes in place. Construct a replacement or use replace; multiple consumers may refer to the same interned UOp.

Understand what your pattern allows

UPat can constrain operation, dtype, arguments, and sources. Named variables capture structure; they do not automatically prove semantic properties such as positivity or alignment.

A callback may inspect its captures and decline the rewrite by returning None. A context object can carry pass-specific state. Both mechanisms are used by real compiler passes, but stateful behavior should be understood before assuming a matcher is an order-independent collection of equations.

Read an actual pass

In schedule/indexing.py, pm_apply_rangeify includes rules for reductions, padding, stacks, source indexing, and movement removal. These rules transform tensor-level structure after range information has been computed. Applying them to an arbitrary UOp without the required context is not meaningful.

In codegen/init.py, full_rewrite_to_sink sequences symbolic cleanup, optimization, expansion, reduction lowering, and later target work. A rule’s correct location depends on which invariants exist at that point.

Use VIZ to inspect a real before/after rewrite and trace a surprising node to the pass that introduced it.

Original chapter by Di Zhu: historical version.