Say I have 2 + 3 as two constants and an ADD. What does the pattern matcher do with those objects?
from tinygrad.uop.ops import UOp, Ops
two = UOp.const(2)
three = UOp.const(3)
add = UOp(Ops.ADD, src=(two, three))
There are three nodes here. The ADD points to the two constants.
CONST(2) ─┐
├──→ ADD
CONST(3) ─┘
The addition has not happened yet. add.src holds references to two and three. Its arg is None. For two, the value 2 is in arg, and src is empty. So the ADD holds the inputs to a calculation. The constants already hold values. [0]
Calling the matcher on the ADD.
from tinygrad.uop.symbolic import symbolic_simple
folded = symbolic_simple.rewrite(add)
This returns a CONST holding 5. The matcher has a list of patterns paired with functions. The ADD matches this rule.
(
UPat(GroupOp.ALU - {Ops.THREEFRY}, src=bare_const, name="a"),
fold_const_alu,
)
ADD belongs to the ALU group, and both inputs match bare_const. name="a" binds the whole ADD node. The matched object goes into fold_const_alu as the argument a, with its references to both constants still there.
fold_const_alu(a=add)
The function runs in Python. Its body is
def fold_const_alu(a: UOp) -> UOp:
return a.const_like(exec_alu(a.op, a.dtype, [const_arg(s) for s in a.src], False))
Splitting that line up so I can follow the values.
values = [const_arg(s) for s in a.src]
value = exec_alu(a.op, a.dtype, values, False)
folded = a.const_like(value)
a.src contains the two constants. const_arg takes the value from each one, giving [2, 3]. Those are ordinary Python numbers now. The ADD object supplies the operation to exec_alu, and the list supplies its operands.
How does it know what ADD means? The table has an entry for it.
import operator
# Just the ADD entry from the table.
python_alu = {
Ops.ADD: operator.add,
}
The call in exec_alu is python_alu[op](*operands). With these values, that becomes
function = python_alu[Ops.ADD]
alu = function(*[2, 3])
So operator.add(2, 3) gives back the Python integer 5. The addition happens here, while the graph is being processed.
There is still a.const_like(5). That calls
UOp.const(5, a.dtype)
The new constant has the ADD’s type, weakint in this example. For this scalar integer, the constructor gives us
UOp(Ops.CONST, arg=5, src=())
The number goes into arg. There are no inputs. This is the UOp returned by fold_const_alu, and then by the matcher. folded points to it. The original ADD still refers to two and three.
What happens to anything that was pointing at the ADD? Returning a constant from the function doesn’t change those references by itself. I can put a parent above the ADD and follow that too.
from tinygrad.uop.ops import graph_rewrite
output = UOp.sink(add)
rewritten = graph_rewrite(output, symbolic_simple)
The SINK’s input is the ADD. The graph walker keeps a mapping from old nodes to their replacements. For this ADD, the replacement is the CONST holding 5. When the walker reaches the SINK, its new inputs contain that constant. It constructs another SINK around them.
old output → ADD(CONST(2), CONST(3))
new output → CONST(5)
There are two roots now. output.src still contains add. The returned rewritten SINK holds the constant in its src. I can compare the two roots and check that the original references are still there.
assert output.src[0] is add
assert rewritten.src[0] is folded
This is the file I ran.
from tinygrad.uop.ops import UOp, Ops, graph_rewrite
from tinygrad.uop.symbolic import symbolic_simple
two = UOp.const(2)
three = UOp.const(3)
add = UOp(Ops.ADD, src=(two, three))
print("before:", [(u.op.name, u.arg) for u in add.toposort()])
folded = symbolic_simple.rewrite(add)
print("after: ", [(u.op.name, u.arg) for u in folded.toposort()])
assert len(add.toposort()) == 3 and len(folded.toposort()) == 1
assert folded.op is Ops.CONST and folded.arg == 5
output = UOp.sink(add)
rewritten = graph_rewrite(output, symbolic_simple)
print("old output points to:", output.src[0].op.name)
print("new output points to:", rewritten.src[0].op.name, rewritten.src[0].arg)
assert output.src[0] is add
assert rewritten.src[0] is folded
before: [('CONST', 2), ('CONST', 3), ('ADD', None)]
after: [('CONST', 5)]
old output points to: ADD
new output points to: CONST 5
Both inputs were constants, so the calculation could happen while processing the graph. Python added the numbers, const_like put the result into a UOp, and the walker built the parent’s inputs using that returned object. The old graph remains. To use the folded graph, I have to continue from rewritten.
I ran this example against tinygrad 868263e04.
[0] The UOp spec checks allowed node forms and types using pattern matchers. Its shared rules include checks for CONST and ALU nodes. The constant-folding rule traced here is in symbolic.py. ↩︎