Episode 4: Build End to End Models

What this lecture covers

Where Episode 3 focused on optimizing a single primitive tensor function, this lecture asks how those primitive functions get stitched together into a full model execution. Using a two-layer neural network trained on Fashion MNIST, it builds the same prediction first in plain NumPy, then in low-level NumPy, then in TVM's Relax abstraction, so you can see exactly what a high-level model call expands into underneath.

The core of the lecture is the computational graph view of a model and the call_tir construct that connects graph-level operations to low-level, destination-passing tensor functions. It also introduces dataflow blocks, which mark which regions of a program can safely be treated as a pure computational graph. After watching, you should be able to explain what call_tir does, why destination-passing style matters for memory management, and how an IRModule can mix TensorIR functions with externally registered library calls.

Key ideas

• End-to-end model execution: assembling multiple primitive tensor functions (like linear layers and ReLU) into a single runnable model, as opposed to optimizing one function in isolation.
• Computational graph (dataflow graph): a graph view where each node is a tensor operation and edges represent data dependencies between operations; this is the abstraction behind frameworks like PyTorch and TensorFlow.
• Destination-passing style: the low-level convention where a function receives pre-allocated output memory as an argument instead of returning newly allocated results, so memory management stays with the higher-level framework.
• call_tir: a Relax construct that wraps a destination-passing-style primitive function (which has side effects, since it writes into pre-allocated memory) so it behaves like a pure, side-effect-free node in a computational graph.
• Dataflow block: a marked region of a Relax function where every operation can be treated as part of a computational graph; operations outside a dataflow block must be treated as ordered, non-graph code.
• IRModule mixing library and TensorIR calls: a Relax function can call some operators through TensorIR (call_tir) and others through externally registered library functions (for example, PyTorch operations exposed via DLPack), using the same call interface for both.
• DLPack: a standard for exchanging tensor memory between frameworks (NumPy, PyTorch, TVM) without copying data, used here to implement environment functions by calling into PyTorch.
• Parameter binding: attaching a model's weights as constants inside the IRModule (via relax.transform.BindParams) so the compiled function takes only the input data, not every weight as a separate argument.

Walkthrough

Recap and imports (1:02)

The lecture recaps the MLC process as transforming development form to deployment form, notes prior lectures focused on single primitive functions, and introduces Relax, TVM's high-level computational graph abstraction, as the topic for assembling primitive functions into full models.

The Fashion MNIST example model (3:02)

The lecture loads the Fashion MNIST dataset and a pretrained two-layer MLP (linear, ReLU, linear, with softmax omitted since only the argmax prediction is needed). It walks through the model's shapes: a 784-element flattened image input, a 128-dimension hidden layer, and a 10-class output.

NumPy and low-level NumPy implementations (8:04)

The model is first implemented directly in NumPy (matmul, bias add, ReLU, matmul, bias add), correctly predicting "T-shirt/top." It's then reimplemented in low-level NumPy, following the same explicit-loop, explicit-allocation conventions from Episode 3, with three primitive functions (linear0, relu0, linear1) tied together by an outer function that manages intermediate buffers.

Computational graphs and call_tir (17:07)

Comparing the TVMScript Relax function to the low-level NumPy version side by side, the lecture shows a near one-to-one correspondence, then explains the computational graph view where each line becomes a graph node with explicit input/output dependencies. It introduces call_tir, showing through a NumPy reimplementation that it wraps a destination-passing-style function (explicit input plus pre-allocated output) so the call instead returns a value, preserving the side-effect-free property computational graphs need for reordering and parallelizing operations.

Dataflow blocks (32:20)

The lecture explains that not every part of a program can be treated as a pure computational graph (for example, file reads or explicit allocations), so Relax uses dataflow blocks to mark regions that are safe to treat as graphs, with an output marker showing which variables are visible outside the block. Most lectures in the course will use a single dataflow block per function.

Building and running the model (37:25)

Using relax.vm.build, the lecture compiles the IRModule into a virtual machine executable, initializes it on CPU, converts the input image and weights to TVM NDArrays, and runs the model's main function, getting the same prediction ("T-shirt/top") as the NumPy version.

Calling external library functions (41:28)

The lecture shows an alternative IRModule where operators are called by string name (for example env.linear) instead of by TensorIR function, and registers implementations for those names in the TVM runtime using PyTorch operations connected via DLPack (a zero-copy tensor exchange standard), demonstrating that TensorIR calls and library calls can be mixed within the same module.

Binding parameters (48:32)

To avoid passing every weight as a separate argument, the lecture uses a parameter-binding transform to embed the model's weights as constants inside the IRModule, so the compiled function's main entry point takes only the input image.

Before you watch

• Watch Episodes 2 and 3 first, since this lecture assumes familiarity with primitive tensor functions, TensorIR and TVM schedules.
• Basic knowledge of a simple feedforward neural network (linear layers, ReLU, softmax) is expected.
• Familiarity with NumPy indexing helps with following the low-level NumPy implementation.

Check your understanding

1. What problem does call_tir solve, and why does a destination-passing-style function break the assumptions of a pure computational graph?
2. What is a dataflow block, and why might a Relax function need operations outside one?
3. How does the lecture implement external "environment" functions like env.linear, and what role does DLPack play?
4. Why do most low-level tensor functions take pre-allocated output buffers as arguments instead of returning newly allocated results?
5. What does binding parameters into an IRModule change about how the compiled function is called?