All tutorials · Updated September 21, 2026 · tinygrad 8ad8f73
BEAM chooses among legal implementations of an already identified kernel by compiling and timing candidates. It does not replace the Tensor scheduler’s job of identifying dependencies and kernel boundaries.
OptOps currently contains TC, SPLIT, PADTO, and SWAP. Old examples using OptOps.UPCAST, LOCAL, or UNROLL as separate operation kinds no longer match this API. The axis role is now part of a split argument.
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.uop.ops import AxisType
opt = Opt(OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST))
assert opt.arg == (4, AxisType.UPCAST)
print(opt)
This constructs a proposed transformation; it does not prove that the transformation is legal for any particular kernel. postrange.Scheduler.apply_opt checks applicability. Splitting requires compatible axis types and extents; tensor-core matching has additional constraints.
The kernel Scheduler here is an optimization object. It is different from the high-level scheduler described in scheduling.
search.py builds an action set, copies candidate schedulers, applies legal options, and limits dimensions such as unrolled work and local work. Candidate kernels are compiled with to_program, duplicate binaries are filtered, and execution timings rank survivors.
The search retains a bounded set of promising candidates and expands them again until its stopping condition is reached. Larger beam width can explore more alternatives, increasing compilation and measurement cost without guaranteeing a faster final application.
The search can time a reduced global workload and scale its estimate. That is useful for candidate selection, but an estimated kernel timing is not a substitute for benchmarking the complete original workload.
from tinygrad import Tensor
a = Tensor.ones(16, 16).contiguous().realize()
b = Tensor.ones(16, 16).contiguous().realize()
out = (a @ b).realize()
assert out.tolist() == [[16.0] * 16 for _ in range(16)]
print("16x16 matmul result verified")
Save this as matmul.py. Compare fresh runs with the same backend:
DEV=CPU BEAM=0 DEBUG=2 python3 matmul.py
DEV=CPU BEAM=2 DEBUG=2 python3 matmul.py
The numerical example is checked independently of search. Search performance is device-dependent; this chapter supplies no portable speedup claim. On a GPU, choose its supported backend and repeat the measurement there.
Report cold compilation/search time, warmed kernel time, and end-to-end time separately. Search results can be cached; compare equivalent cache conditions. A faster kernel may still make a short-lived application slower if tuning costs dominate.
Use representative data when validating correctness, particularly after changing indexing, precision, padding, or reduction transformations. A constant-input timing workload is not sufficient regression coverage.
Original chapter by Di Zhu: historical version.