Episode 3: TensorIR Case Study

What this lecture covers

This lecture is a deep, hands-on dive into TensorIR, TVM's tensor program abstraction, using a matrix-multiplication-plus-ReLU example throughout. It answers a practical question left open in Episode 2: what does a real tensor program abstraction look like in detail, and how does its extra structure (the "block" concept) let a compiler safely and automatically transform loop code into faster variants?

The lecture builds a low-level NumPy reference implementation first, then shows the equivalent TVMScript/TensorIR code side by side, explaining blocks and block axis properties. It then performs live schedule transformations, split, reorder, compute_at, and decompose_reduction, building toward a version that runs more than five times faster, and explains why through CPU cache behavior. After watching, you should be able to read TensorIR code, understand what block axes declare, and explain why loop order affects performance.

Key ideas

• Low-level NumPy: a restricted style of NumPy using explicit loops and pre-allocated arrays instead of vectorized calls, used here as a bridge to understanding low-level tensor program code.
• TensorIR / TVMScript: TVM's tensor program abstraction, written in a Python-embedded syntax, that represents the same computation as low-level NumPy but with explicit buffers, loops and blocks.
• Block: a TensorIR construct (T.block) that encapsulates a self-contained unit of computation, declaring which regions of memory it reads and writes.
• Block axis (vi, vj, vk): declarations binding loop iterators to a block's computation, each marked as either spatial (independent across iterations, like output row/column indices) or reduction (accumulated across iterations, like the summed dimension in a matmul).
• T.axis.remap: shorthand (vi, vj, vk = T.axis.remap("SSR", [i, j, k])) for declaring multiple block axes and their spatial/reduction properties at once.
• Schedule: a TVM object (tvm.tir.Schedule) used to apply transformations such as split, reorder, compute_at and decompose_reduction to a block's loops without rewriting the program by hand.
• IRModule: a container holding a collection of tensor primitive functions (PrimFuncs), inspected and transformed as a unit.
• Cache locality: reordering loops changes memory access strides; accessing nearby elements in sequence keeps data in the CPU's L1/L2 cache, which the lecture shows can produce a measured 5x-plus speedup.

Walkthrough

Recap and the running example (1:03)

The lecture recaps the MLC process as transforming tensor functions across abstractions, then introduces the example used throughout: matrix multiply Y = A @ B followed by C = relu(Y), chosen because it resembles a linear-plus-ReLU layer in a neural network.

Building a low-level NumPy reference (5:44)

Starting from plain NumPy (A @ B, then np.maximum), the lecture asks what happens "under the hood" and writes an explicit low-level NumPy version: pre-allocated buffers, explicit triple-nested loops for the matmul reduction, and a separate loop for the ReLU, validated against the NumPy result with np.testing.assert_allclose.

TensorIR and the block concept (13:17)

Placing the TVMScript implementation next to the low-level NumPy code, the lecture maps each element: T.Buffer arguments correspond to the NumPy function arguments, T.alloc_buffer to the intermediate array, and T.grid loops to the explicit for-loops. It then focuses on the new construct with no NumPy equivalent: the block, which declares its axes (vi, vj, vk) as spatial or reduction, encoding which iterations are independent (safe to reorder or parallelize) and which must be accumulated in a fixed order.

Inspecting IRModules and PrimFunc (31:28)

The lecture shows that a TensorIR module is typed as IRModule, a container that can hold multiple PrimFuncs (for example, separate matmul and ReLU functions), and demonstrates retrieving individual functions by name.

Transforming the schedule step by step (39:33)

Using tvm.tir.Schedule, the lecture retrieves the compute block and its loops, applies split to divide the j loop into j0/j1, applies reorder to move j0 between the outer loops and k, then applies compute_at to move the ReLU computation inside the j0 loop, and finally decompose_reduction to separate the matmul's zero-initialization from its accumulation update. Each step is compared against a matching low-level NumPy variant.

Building, running and measuring performance (53:50)

The transformed IRModule is compiled with tvm.build(mod, target="llvm"), run on tvm.nd.array inputs, and validated for correctness. Timing both the original and transformed schedules shows the transformed version runs more than five times faster despite only reordering loops, which the lecture attributes to memory access patterns.

Why loop order matters: cache locality (56:06)

The lecture explains CPU cache hierarchy (L1, L2, main memory) and their large latency differences, then shows that the transformed loop order accesses Y and B in contiguous four-element stripes, keeping recently used data in cache, while the original order jumps across memory with a stride of 128 elements, causing frequent cache misses.

Constructing TensorIR from tensor expressions (1:01:08)

As an alternative to hand-writing TVMScript, the lecture shows the tensor expression (TE) domain-specific language, using te.placeholder and te.compute with a lambda describing each output element, then te.create_prim_func to generate the same TensorIR function automatically.

Before you watch

• Watch Episodes 1 and 2 first; this lecture assumes familiarity with tensors, tensor functions and the general tensor program abstraction.
• Comfort reading NumPy array code and nested loops is necessary to follow the low-level NumPy comparisons.
• A basic idea of CPU cache hierarchy is helpful but not required, since the lecture explains it directly.

Check your understanding

1. What does it mean for a block axis to be marked "spatial" versus "reduction," and why does that distinction matter for parallelization?
2. In the TVMScript code, what information does a block declare that a plain low-level NumPy loop does not?
3. What sequence of schedule operations transformed the original matmul-plus-ReLU program into the faster version, and what did each step change?
4. Why did reordering the loops produce a measured performance improvement, according to the cache locality explanation?
5. What is te.compute used for, and how does it relate to writing TensorIR by hand in TVMScript?