Skip to content

TensorIR Scheduling and Memory Tiling

This document explores the theoretical concepts behind Loop Tiling and Scheduling within TVM's TensorIR (Step 4B).

The Limitations of Sequential Fusion

In 05_architecture_clarifications, we established that FuseTIR groups loops into a single kernel to avoid flushing intermediate data to VRAM (Global Memory). However, simply placing loops side-by-side in a function is not enough to guarantee performance.

Consider an image of size 224x224 processed by a Convolution layer. If the Conv loop executes fully, it produces $224 \times 224 = 50,176$ intermediate float32 values.

GPU's ultra-fast L1 Cache (Shared Memory) is extremely small (often around 64KB per Streaming Multiprocessor). The entire 50,176-value tensor simply cannot fit. Once the L1 Cache is full, the hardware has no choice but to spill the excess data back out to VRAM, completely defeating the purpose of the fusion.

The Solution: Loop Tiling

To solve this capacity limit, compilers use a technique called Loop Tiling (or Loop Blocking).

Instead of processing the entire 224x224 image at once, we chop the computation into smaller, bite-sized "Tiles" that are guaranteed to fit comfortably inside the L1 Cache.

The Mechanics of Tiling

Tiling involves two primary AST transformations:

  1. Loop Splitting: A single loop iterating from 0 to 224 is mathematically split into two nested loops:
  2. An Outer Loop iterating over the tiles (e.g., 14 steps).
  3. An Inner Loop iterating within the tile (e.g., 16 steps).
  4. Loop Reordering: The nested loops are reordered so that the outer loops of different operations (like Conv and Add) are fused together.

compute_at (Physical Interleaving)

In TVM TensorIR, the primitive compute_at (or reverse_compute_at) is the command that physically interleaves the loops.

When we tell the compiler to compute Add at the spatial axis of Conv2d, we are instructing it to:

"Run the Conv loop for exactly one 16x16 tile. Then immediately pause the Conv loop, and run the Add loop for that exact same 16x16 tile. Once both are done, move to the next tile."

This ensures that the intermediate data generated by the Conv loop is immediately consumed by the Add loop while it is still hot in the L1 Cache, maximizing cache hits and minimizing VRAM bandwidth usage.