Episode 6: Integration with Machine Learning Frameworks

What this lecture covers

This lecture connects the abstractions built in earlier episodes to the outside world: how do you take a model written in an existing framework like PyTorch and bring it into the machine learning compilation flow? It answers this by introducing two building tools, the tensor expression (TE) API for generating primitive tensor functions and the BlockBuilder API for incrementally constructing a computational graph function, and then using TorchFX to trace a PyTorch model's graph and translate it, node by node, into an IRModule.

The lecture works through the running Fashion MNIST two-layer network example again, first building its IRModule by hand with BlockBuilder, then automating the process with a generic FX-to-Relax translator, and finally comparing translating operations into TensorIR calls versus higher-level Relax built-in operators. After watching, you should be able to explain what BlockBuilder and emit_te do, and how a computational graph node in one framework maps to one in another.

Key ideas

• Tensor expression (TE): a domain-specific language (te.placeholder, te.compute) for concisely describing a tensor computation, such as matrix multiplication, that can be converted into a TensorIR PrimFunc via te.create_prim_func.
• Fused TE functions: te.compute can chain operations (for example matmul followed by ReLU) into a single generated tensor function, useful for testing or later fusion transformations.
• BlockBuilder: a TVM API for incrementally constructing a Relax function, using scopes (function, dataflow block) and emit/emit_te calls that correspond closely to lines of generated TVMScript.
• emit_te: a BlockBuilder method that, given a TE function and its arguments, automatically creates the placeholder, builds the PrimFunc, inserts it into the IRModule, and emits a call_tir to it, returning a DataflowVar referencing the result.
• DataflowVar: a Relax variable that is internal to a dataflow block and not visible outside it, distinct from outputs marked visible via emit_output.
• TorchFX symbolic tracing: a PyTorch utility that captures a model's forward computation as a graph of nodes (placeholders, call_function, call_module, get_attr, output), which can be walked in topological order for translation, though it cannot represent control flow.
• Generic FX-to-Relax translation: iterating over an FX graph's nodes, mapping each PyTorch parameter to a Relax constant and each computation node to a Relax variable via a dictionary of per-operator translation functions.
• High-level Relax operators vs. TensorIR calls: an operation like matmul can be translated either into a low-level call_tir to a generated TensorIR function, or into a high-level built-in operator like relax.op.add, deferring the choice of implementation (library call or generated code) to a later stage.

Walkthrough

Recap and the tensor expression API (1:01)

The lecture reviews te.placeholder and te.compute, building a generic matrix multiplication function that works for any compatible shapes, and shows converting it to a TensorIR PrimFunc with te.create_prim_func and displaying it with show.

Generic and fused TE functions (7:09)

The lecture demonstrates a generic ReLU function built with a lambda that works across different tensor shapes, then shows fusing matmul and ReLU into a single TE function, with the option to keep or expose the intermediate buffer as an explicit function argument.

Building a Relax function with BlockBuilder (11:11)

The lecture introduces Relax variables as computational graph node placeholders, then uses BlockBuilder with function and dataflow-block scopes, calling emit_te to insert matmul and ReLU as call_tir operations and emit_output to mark the visible return value, producing an IRModule closely matching what was written by hand in Episode 4.

From PyTorch model to TorchFX graph (21:22)

Starting from a simple nn.Module implementing matmul plus ReLU, the lecture uses torch.fx.symbolic_trace to capture the model's computation as a graph of nodes, noting the limitation that FX tracing cannot represent control flow.

A generic FX-to-IRModule translator (24:22)

The lecture builds a from_fx function that walks the FX graph in topological order, handling placeholder nodes (model inputs), get_attr nodes (parameters, mapped to Relax constants), call_function nodes (dispatched through a user-supplied translation dictionary), and the final output node, using a node map to track FX nodes to their translated Relax variables.

Applying the translator to Fashion MNIST (36:30)

The lecture applies the translation flow to the two-layer Fashion MNIST model, defining a call_module map for nn.Linear and nn.ReLU (using TOPI's predefined topi.nn.dense tensor expression), builds the resulting IRModule, and confirms it produces the same prediction as the original PyTorch model.

High-level operators versus TensorIR calls (41:35)

The lecture shows an alternative mapping that translates PyTorch operations into Relax's built-in high-level operators (relax.op.add, relax.op.matmul) instead of generated TensorIR functions, explaining that deferring to a high-level operator keeps options open for later mapping to a library call or generated code, while direct TensorIR calls offer more flexibility for custom or fused operations.

Before you watch

• Watch Episode 4 first, since this lecture reuses its Fashion MNIST two-layer model and BlockBuilder-generated IRModule as reference points.
• Familiarity with basic PyTorch (nn.Module, forward) is expected for the FX tracing sections.
• Understanding call_tir and dataflow blocks from Episode 4 helps when reading the BlockBuilder output.

Check your understanding

1. What does emit_te do inside BlockBuilder, and what steps does it perform under the hood?
2. Why can TorchFX's symbolic tracing not represent a model containing control flow?
3. How does the generic FX-to-Relax translator decide what to do with a call_function node versus a call_module node?
4. What is the difference between translating an operation into a call_tir to a TensorIR function and translating it into a high-level Relax operator like relax.op.add?
5. Why does the translator need a separate map_param step for handling module weights?