Repository navigation
Perf/tolk compile time - #235
Draft
ghennequin wants to merge 6 commits into
Draft
ghennequin wants to merge 6 commits into
ghennequin wants to merge 6 commits into
Conversation
A beam search timed a candidate by running it, so every sample compiled, linked and allocated the candidate's program again, and a kernel an earlier round compiled was compiled anew when a later round's action path made it again. Search.beam_search takes ~prepare in place of ~measure: a function that prepares a program once and returns the sampler the search times it with, up to three times. Tolk_engine.prepare compiles and links the program and allocates its buffers once; Tolk_engine.measure is the first sample of a prepared program. The search keeps a candidate's program across rounds and reuses it when its kernel is reached again, while the compute filter and the binaries already timed are decided anew each round. The search suite's add_3d case compiles 58 programs for 50 distinct kernels, where compiling each round's candidates anew takes 87.
A graph of index arithmetic and comparisons, as a jitted iLQR solve builds, compiled in minutes while its eager run takes milliseconds. Two passes walked a slice per node where a property of the slice answers. - fold_where_closure asked whether a branch reaches its condition before its own INDEX gate, so every selection paid a walk of the branch that the gate then rejected. The cheap guards run first: the gate is the operation set a node reaches, kept per node and shared with a source that adds none. - remove_bufferize toposorted the slice of a reduction's sources for every stage it considered, to ask whether it reads a buffer. Each node now keeps that property, shared with a source that reads one. A differentiable iLQR solve of 20 steps unrolled six times compiled in 15 s where it took 56 s, its implicit gradient in 8 s where it took 52 s; a pendulum solve of 10 steps unrolled six times in 40 s where it took 76 s, its gradient in 29 s where it took 79 s.
A candidate of BEAM_UOPS_MAX instructions or more was rendered, compiled and only then dropped, and the candidates that reach the limit are the unrolls whose compilation is slowest: comgr spent seconds on a candidate the search had already decided against, while the round that timed it waited on the compile. Codegen.with_uops_cap caps a program's linearized instructions, domain local since candidates compile on worker domains, and Codegen.do_linearize raises Too_many_uops once linearizing reaches it, so neither the renderer nor the compiler sees the program. Search.try_compile sets the cap from BEAM_UOPS_MAX around to_program and drops the candidate the exception names, where it before dropped it after the fact. The lorenz example's step, cold: the compile work spent on candidates that were dropped fell from 154 s to 30 s over its 97 searches, and the rounds whose time was one candidate's compilation, from 12.3 s to about 2 s.
A candidate's sample ran its program again and again until the runs took 10 us in all, and every run paid a profile session of its own: its stamps, a synchronization, and a calibration of the device's clock against the host's that takes five stamped round trips of the device. A microsecond kernel paid tens of sessions for one measurement. A run's span is stamped by the device itself, so one run measures as well as the mean of many: Tolk_engine.prepare returns one run's span as the sample, and the search keeps the least of the three it takes. A clock's calibration is kept for two seconds, per device: a span's length needs the clock's rate, which a stale offset does not skew, and a stop whose last calibration is recent calibrates nothing. The lorenz example's step, cold: the candidates' timing fell from 66 s to 60 s over its 97 searches.
…only in buffers A graph that reaches the same kernel twice — a forward pass run on two samples, a layer a model uses once per block — searched and compiled it once per reach: the beam cache's key was the kernel's own, and two kernels that name different buffers differ in it though every shape, variable and operation is the same. The key now renames a kernel's parameters by their first appearance, so kernels that differ only in which buffers they name share an answer. A kernel of other sizes, variables or operations keeps its own key: the renaming is a bijection on the parameters' slots, and the rest of the kernel is the key. Of the first 48 searches of the lorenz example's step, 3 are answered by an earlier one; a graph that repeats a layer shares one per repeat.
…that are noise A search spends seconds per kernel — a compile, a link and a load of the device's code object for every candidate it tries — to save microseconds on a kernel that is microseconds long already, and it keeps a round going for an improvement smaller than the jitter of the timings it compares. BEAM_SKIP_UNDER_US=n times the kernel itself in the first round, which then holds it as a candidate the search can answer with, and stops after that round when the kernel is under n microseconds. BEAM_MIN_PROGRESS_REL=f stops a search whose best candidate improves on the best so far by less than f of its own time, where BEAM_MIN_PROGRESS is that same bound in seconds. Both are off by default: searching every kernel and every improvement is what the beam has always done. The lorenz example's step, cold: 20 and 0.005 take its jit from 457 s to 296 s, with steps of 1.6 to 1.9 ms and the same drift-norm trajectory; a warm cache takes 3.2 s.
tmattio
added a commit
that referenced
this pull request
Oct 5, 2026
A beam search paid most of its time before it timed anything. Each of a candidate's three samples called Tolk_engine.measure, which rebuilt scratch buffers from the program's ELF signature, linked a schedule, and averaged up to 1000 runs, printing a DEBUG=2 line for each. Equal kernels reached by two sequences of actions were compiled twice, and a candidate past BEAM_UOPS_MAX was rendered and compiled before being dropped. PR #235 measured these costs on AMD: 12 to 16% duplicate compiles, and 154 s of a lorenz_simple search compiling over-cap candidates. Tolk_engine.timer replaces measure and is staged as a search uses it: applied to the kernel it allocates one buffer per storage parameter, once per search and never for a cached one; applied to a program it links it once over those buffers; each sample is one silent run. Building the call from the kernel's parameters also removes the Option.get on sparse globals. beam_search takes ~time, applies it to its kernel before compiling, keeps each kernel's compilation by Scheduler.ast, filled before a round is spread over domains, and checks the uop cap on Codegen.linearize, a new stage that to_program completes. ?cold goes: every caller asked for it. The search still picks what tinygrad's picks: searches.golden and the generator check are unchanged. On an M1 Max under load, a candidate's setup on Metal is about 1.1-1.3 ms, from 3.5-4.3 ms. New search rows in the bench record the link. The replay of a compiled call allocates 823 words, from 828, since a run walks its calls without a closure. D115 and D116 record the compile changes.
tmattio
added a commit
that referenced
this pull request
Oct 5, 2026
tinygrad's search starts from an untimed kernel of infinite time and compares the least of each program's timings. The least over hundreds of noisy candidates is biased low, so a round continued, and a candidate was kept alone at the end, on a gain within the noise. A search could answer a kernel slower than the one it was given. PR #235 measured trailing rounds on 2 to 4% gains on AMD (25.96 to 25.42 us), and first rounds that searched kernels already at their floor. The search now compiles and times its kernel before its first round, with no early stop, and starts the beam from it. A round progresses iff the greatest sample of its fastest candidate plus BEAM_MIN_PROGRESS is below the least sample of the beam's first kernel; otherwise the search answers the beam's first kernel. One comparison replaces tinygrad's three exits and the fastest kept alone. The cost is the kernel's own compilation and three samples per search. The golden generator applies the same rules, and searches.golden was regenerated: every search chooses what it chose, with three more samples. On the host, a search of a kernel as small as the exec suite's answers the kernel itself, so that suite scales optimised programs' samples down to keep checking what a searched program computes. D117 records the departure.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
WIP