Episode 5: Automated Program Optimization

What this lecture covers

Earlier lectures showed how to transform a tensor program by hand, choosing specific loop splits and orderings, and how to assemble primitive functions into an end-to-end model. This lecture asks how to automate the search for a good transformation instead of hand-tuning every choice. It introduces stochastic schedule transformations, where some decisions (like a loop split factor) are left as random variables to be sampled rather than fixed values.

The lecture shows how a search space defined this way can be explored by a simple random search, then by TVM's meta_schedule.tune_tir using evolutionary search, and finally by a default, generic search space that requires no hand-written schedule at all. It closes by replacing one primitive function in an end-to-end model with its tuned version. After watching, you should understand what a stochastic transformation and a trace are, and how automated tuning turns a search space into an optimized program.

Key ideas

• Schedule trace: a record of the sequence of transformation calls (get_block, get_loops, split, reorder, decompose_reduction, and so on) applied to a schedule, which can be printed and reused.
• Stochastic schedule transformation: a schedule where some parameters, such as a loop split factor, are sampled randomly (for example with sample_perfect_tile) instead of fixed, so each run can produce a different but valid program.
• Search space: the set of possible programs a stochastic schedule can generate; every valid combination of sampled decisions is one candidate program in that space.
• Random search: repeatedly sampling a stochastic schedule, building and benchmarking each result, and keeping the fastest one found so far.
• meta_schedule.tune_tir: TVM's tuning API that searches a given space (or a built-in default space) for a fast schedule, using evolutionary search rather than pure random sampling.
• Evolutionary search: a tuning strategy that mutates the recorded decisions in a trace and favors traces that produce better-performing programs, rather than resampling from scratch each time.
• Default (auto-scheduling) search space: a built-in, generic set of stochastic transformation rules that apply to common workloads like matrix multiplication or convolution without a hand-written schedule function.
• Replacing a primitive function: after tuning finds a faster implementation of one primitive function, it can be substituted into an existing end-to-end IRModule in place of the original, speeding up the full model.

Walkthrough

Recap of manual schedule transformation (1:01)

The lecture reviews the manual schedule transformation from Episode 3 on a 128x128 matrix multiplication: getting blocks and loops, splitting, reordering, and decomposing the reduction, and shows the resulting program runs faster than the untransformed version.

Traces (6:07)

The lecture introduces schedule.trace, showing it records the exact sequence of transformation calls used to reach the current schedule, and demonstrates comparing a trace side by side with the code that produced it.

From fixed to stochastic transformations (10:34)

Comparing a schedule_mm function to a new stochastic_schedule_mm, the lecture shows the only difference is a sample_perfect_tile call that draws random factors (which multiply to 128) instead of a hard-coded split factor of 4. Running the stochastic version repeatedly produces different, equally valid loop splits (such as 8x16 or 64x2) each time.

Anatomy of stochastic transformations (15:31)

The lecture explains that stochastic transformations add two elements to ordinary scheduling: sampling operations that produce symbolic random variables, and subsequent transformation calls that depend on those variables. It walks through running sample_perfect_tile first (which only records a decision in the trace) and then split/reorder (which visibly change the generated code based on the sampled values).

Random search over the space (21:13)

A simple Python loop repeatedly calls the stochastic schedule function, builds and benchmarks each resulting program, and keeps the best one found, illustrating a "poor man's" automated search before introducing TVM's built-in tuner.

Meta-schedule tuning (24:16)

The lecture calls meta_schedule.tune_tir, supplying the module, target, number of CPU cores, a trial budget, and the stochastic search space, and explains that under the hood it uses evolutionary search over trace decisions rather than plain random sampling, producing a schedule close to or better than the hand-tuned one.

The default search space (28:16)

Removing the hand-specified search space and letting meta-schedule use its built-in generic rules produces a program with multi-level loop tiling, vectorization and parallelization, running noticeably faster than the manually tuned version, at the cost of the schedule being less specific to this exact computation.

Replacing a primitive function end to end (35:19)

Returning to the two-layer MLP from Episode 4, the lecture isolates the model's linear layer, runs the stochastic-transformation-and-tune workflow on it in isolation, and then substitutes the tuned implementation back into the original end-to-end IRModule, showing the model still predicts correctly and can run faster.

Before you watch

• Watch Episode 3 first, since this lecture directly extends the manual schedule transformations (split, reorder, decompose_reduction) taught there.
• Watch Episode 4 for context on how a tuned primitive function fits back into an end-to-end model.
• Basic familiarity with the TVM Schedule API and traces is assumed.

Check your understanding

1. What is the difference between a fixed schedule transformation and a stochastic one?
2. What does a schedule trace record, and why is it useful for search and tuning?
3. How does evolutionary search in meta-schedule differ from pure random sampling of a stochastic schedule?
4. What is the advantage of meta-schedule's default search space over a hand-written stochastic schedule, and what is the trade-off?
5. Once tuning finds a faster implementation of one primitive function, how is it incorporated back into an end-to-end model?