Episode 2: Tensor Program Abstraction

What this lecture covers

This lecture zooms into a single piece of the machine learning compilation pipeline: how one primitive step of a model's computation, such as an add or a matrix multiply, is represented and optimized. It answers the question of what abstraction lets a compiler both express a tensor computation and systematically transform it into a faster version, rather than requiring an engineer to hand-write every optimized variant.

Building on Episode 1's concepts of tensors and tensor functions, the lecture introduces primitive tensor functions as the single units of computation inside a larger model, then presents the tensor program abstraction as a structured way to represent loops, buffers and compute statements so that transformations can be applied programmatically. A live TVM walkthrough shows splitting, reordering and parallelizing a loop, and compiling the result into a runnable function. After watching, you should be able to describe the three elements of a tensor program and explain why loop-carry dependencies restrict which transformations are valid.

Key ideas

• Primitive tensor function: a single unit of computation within a larger model execution, such as one add or linear layer, sometimes called a tensor operator.
• Multiple abstractions for the same function: the same add operation can be represented as a library call (torch.add), a plain Python loop, or low-level C code, each a different abstraction of the same computation.
• Tensor program abstraction: a representation built from three elements: multi-dimensional input, output and temporary buffers; a loop nest over the computation's dimensions; and a compute statement executed at each loop iteration.
• Program-based transformation: instead of hand-writing every optimized variant, a tensor program can be transformed by a tool through operations like loop splitting, loop reordering and thread or block binding.
• Loop splitting: dividing one loop into two nested loops, for example splitting a 128-length loop into chunks that let several elements be processed together.
• Loop-carry dependency: some loop reorderings are invalid because a later iteration depends on a value computed in an earlier one; tensor programs can be annotated (e.g. marking an axis "spatial") to declare when iterations are independent and safe to reorder or parallelize.
• Schedule: a TVM helper object used to apply a sequence of transformations, such as split, reorder and parallelize, to an IRModule's tensor program.
• Build: the TVM step that compiles a (possibly transformed) IRModule to a target such as LLVM, producing a runnable packed function.

Walkthrough

Recap and primitive tensor functions (0:00)

The lecture recaps tensors and tensor functions from Episode 1, then narrows in on primitive tensor functions: the single-step operations, like an add or a linear layer, that a model execution is built from. It frames the day's question as what representations exist for a primitive function and how each can be optimized.

Representing the same primitive function three ways (2:02)

Using a length-128 vector add as the running example, the lecture shows the same computation expressed as a torch.add library call, as an explicit Python for-loop, and as low-level C code. Each is a different abstraction, and the lecture frames the core MLC question as which abstraction to choose and how to transform between them.

Primitive function transformation and library mapping (7:12)

The lecture demonstrates transforming a plain for-loop add into a parallelized, four-element-at-a-time version, showing this can be more efficient than the naive loop. It contrasts two strategies frameworks use: simple library-call mapping (dispatching directly to a target's native library, like CUDA or ARM NEON) versus fine-grained program transformation, noting the trade-off between the engineering cost of manual mappings and the reusability of general transformations.

The tensor program abstraction's three elements (14:18)

The lecture defines a tensor program as containing multi-dimensional buffers (including temporary ones for intermediate results), a loop nest over the computation's dimensions, and compute statements tied to each loop position. It explains that this abstraction exists specifically to make program-based transformations possible, rather than requiring an engineer to manually write each optimized variant.

Loop splitting, reordering and dependencies (20:22)

Working through pseudocode, the lecture demonstrates loop splitting (dividing one loop into two), loop reordering, and binding an outer loop to GPU threads. It then explains loop-carry dependency: reordering is invalid when a later iteration reads a value written by an earlier one, and shows how annotating an axis as "spatial" in TVMScript tells the system the iterations are independent and safe to parallelize.

Live TVM walkthrough: building and transforming a schedule (29:26)

The lecture installs the mlc package, constructs an IRModule from TVMScript, and inspects it as a collection of tensor functions. Using tvm.tir.Schedule, it retrieves the compute block and its loop, applies split to divide the loop into three nested loops of lengths derived from a splitting factor, applies reorder to swap loop order, and applies parallel to mark the outer loop for parallel execution, printing the updated TVMScript after each step.

Compiling and running the transformed program (39:41)

The lecture builds the transformed module for the LLVM (CPU) target with tvm.build, constructs input and output NumPy-backed arrays with tvm.nd.array, and runs the compiled function, confirming the result matches the expected sum. It closes by summarizing the flow: start from a tensor program, apply transformations such as reordering, splitting and parallelization, then compile to a runnable function.

Before you watch

• Watch Episode 1 first, since this lecture builds directly on its definitions of tensor and tensor function.
• Basic familiarity with loops and array indexing in Python or a similar language helps with following the loop-transformation examples.
• No prior TVM experience is assumed; the live coding walkthrough introduces the API as it goes.

Check your understanding

1. What is a primitive tensor function, and how does it relate to an end-to-end model execution?
2. List the three key elements that make up a tensor program.
3. Why can a plain library-call mapping approach struggle to scale across many operators and hardware targets?
4. What is a loop-carry dependency, and why does it block certain loop reorderings?
5. What roles do TVM's Schedule and build functions play in the tensor program optimization flow?