Lecture 24: ML Compilation and Deployment Implementation

What this lecture covers

This lecture answers how the machine learning compilation concepts from the previous lecture look in practice, using Apache TVM (the material is adapted from TVM's own tutorials). It moves from writing a small loop-level program in TVM's tensor IR, through applying schedule transformations to speed up matrix multiplication, to compiling and running a full pretrained ResNet-18 model end to end.

After watching, you can write a simple TVM tensor IR program, build and invoke it as a runnable function, apply schedule transformations such as loop splitting, reordering, and tiling to change a loop nest's performance, and describe the steps of taking a PyTorch model through TVM's relay pipeline to a running, benchmarked deployment.

Key ideas

• Tensor IR (TVM script): a Python-like domain-specific language for describing loop-level computation directly, with blocks, axis annotations, and explicit iteration domains, used here to define a simple vector-add and then a matrix multiplication.
• Build and run cycle: an IR module is compiled with tvm.build against a target (such as LLVM for CPU) to produce a runtime module containing a callable packed function, which reads from input arrays and writes to an output array in place.
• Schedule: a helper object created from an IR module that exposes transformation APIs (split, reorder, parallel, and others) to restructure a loop nest programmatically, without hand-editing code, producing a new module with different loop structure but equivalent computation.
• Loop splitting and reordering change performance, not results: splitting the reduction/spatial loops of a matrix multiplication and reordering them (a form of register/cache blocking similar to earlier hardware acceleration lectures) can substantially reduce runtime for the same mathematical output.
• Benchmarking with an evaluator: TVM's time_evaluator repeatedly runs a compiled function to measure execution time, which the lecture uses to directly compare the untransformed and transformed matrix multiplication kernels.
• Automated search workflow: generating many schedule variants, benchmarking each on the target device, and picking the best is the manual version of the automated search process (schedule generation, benchmarking, cost model) described in the previous lecture.
• End-to-end deployment with relay: TVM's relay module takes a traced PyTorch model (via TorchScript) and input shapes, builds a computational graph of operators (convolutions, batch norm, ReLU, and so on) with their parameters, lowers and compiles it, and produces a GraphExecutorModule that can be fed an input image and run to get predictions.

Walkthrough

A minimal tensor IR program (1:03)

The lecture defines a vector-add computation (C = A + B over length-128 vectors) using TVM's tensor expression DSL, then inspects the generated TVM script: a function taking A, B, C as buffer parameters with a loop over the 128 elements, annotated to mark the spatial iteration.

Building and invoking the function (4:04)

The IR module is compiled with tvm.build(mod, target="llvm") to produce a runtime module; the main function is extracted as a packed function. NumPy-backed TVM arrays for A, B, and an uninitialized C are created, and calling the function with these arrays fills C with the computed sum, confirming the compiled kernel runs correctly.

Schedule transformations on the vector-add (8:12)

A Schedule object is created from the module, the target block is fetched, and its surrounding loop is retrieved. Calling split on the loop with a factor divides it into two nested loops; calling reorder changes their nesting order; calling parallel marks a loop for parallel execution. Each transformation is shown to change the printed TVM script's loop structure while leaving the underlying computation equivalent, illustrating that schedule transformations refactor the loop nest without hand-editing code.

Matrix multiplication and blocked scheduling (15:24)

A 1024x1024 matrix multiplication is defined in tensor IR, with an explicit reduction axis and an initialization branch for the first iteration of the accumulation. Built and benchmarked untransformed, it takes on the order of a few seconds per run. After splitting and reordering the y, x, and k loops into a tiled form (an approach the lecture compares to the register and cache blocking covered in earlier hardware acceleration lectures), the benchmarked runtime drops several times over, and the lecture experiments with different tile sizes to show performance varies by configuration.

End-to-end deployment of ResNet-18 (30:43)

A pretrained ResNet-18 is downloaded from torchvision, traced with TorchScript, and an input image is preprocessed into a tensor. TVM's relay.frontend.from_pytorch converts the traced module and input shapes into a relay module, a graph of operators such as convolution, batch norm, add, and ReLU with their parameters. relay.build compiles this into a runnable library, which is wrapped in a GraphExecutorModule that accepts the input image, runs the forward pass, and returns output predictions. Mapping the output through ImageNet class labels via argmax produces a prediction (a tabby cat) that matches the input image, demonstrating the full pipeline from a trained PyTorch model to a running deployed prediction.

Before you watch

• Watch the preceding lecture on model deployment and machine learning compilation, since this lecture implements the concepts (IR modules, loop-nest transformations, automated search) it introduces conceptually.
• Review the hardware acceleration lectures on loop tiling and cache blocking, since the matrix multiplication scheduling here directly parallels those techniques.
• Familiarity with basic PyTorch model loading and TorchScript tracing is helpful for following the ResNet-18 deployment section.

Check your understanding

1. What is the difference between writing a loop-level program in tensor IR and applying a schedule transformation to an existing IR module?
2. Why does splitting and reordering the loops of the matrix multiplication reduce its measured runtime without changing its output?
3. What role does the time_evaluator benchmarking play in choosing between different schedule transformations?
4. What are the main stages relay goes through to turn a traced PyTorch model into a runnable prediction on an input image?