Ep 8: GPU and Specialized Hardware, Part 2

What this lecture covers

This lecture continues the specialized-hardware discussion from the previous episode, asking how machine learning compilation can target accelerators built around bulk tensor instructions rather than scalar or vector ones. It works through a hypothetical accelerator that computes matrix multiplication in 16x16 chunks, first as low-level NumPy code, then as TensorIR.

By the end, you can explain the three ingredients an MLC backend needs to support a specialized accelerator: special memory scopes, tensorized compute blocks, and a tensorization process that swaps a block's body for a hardware-specific implementation. The lecture builds directly on the block and scheduling primitives introduced in earlier CPU and GPU lectures.

Key ideas

• Tensor computing: modern accelerators (TPU, tensor cores, matrix cores) operate on bulk regions of data at once instead of one element at a time.
• Low-level NumPy model: a worked example represents a hypothetical accelerator's registers and instructions as plain NumPy code, using temporary arrays for a register, b register, and an accumulator.
• Special memory scopes: tags like a.register, b.register, and accumulator mark which allocations map to fast on-chip memory rather than global memory.
• Tensorized block: a block whose body computes a subregion of the result (e.g. a 16x16x16 matrix multiply) rather than a single scalar, while still exposing clear read/write regions to the surrounding loops.
• Blockize: a scheduling primitive that takes an ordinary scalar loop nest, splits and reorders it, and groups a chunk of iterations into a new tensorized block.
• Tensor intrinsic: a registered pair of functions - a description (what the block computes) and an implementation (how to compute it, e.g. a call to a microkernel) - that the compiler matches against a block and substitutes in.
• Strided buffers: match_buffer declarations record strides so an intrinsic can operate on a non-contiguous sub-region of a larger array.
• Tensorization: the scheduling step, analogous to vectorization, that replaces a scalar-form block body with a call to the registered intrinsic implementation.

Walkthrough

Why specialized hardware needs new abstractions (0:00)

The lecture opens by tracing the progression from scalar computing to vector units to today's specialized accelerators (Google's TPU, NVIDIA's tensor cores, AMD's matrix cores). It explains why this is hard for compilation: each accelerator has different instruction shapes (say, 16x16 matrix multiply versus a two-vector dot product), so a general MLC process needs abstractions that generalize across backends rather than one custom flow per chip.

A low-level NumPy model of an accelerator (5:03)

The lecture builds a hypothetical accelerator model in NumPy for a 1024x1024 transposed matrix multiplication (c = a @ b.T). Three temporary arrays stand in for hardware registers: a.register, b.register, and an accumulator. The big matrix is split into 16x16 sub-matrices; each iteration uses a dma_copy instruction to move a sub-matrix into a register, then an acc_tmm_add instruction to multiply-accumulate onto the accumulator, running 64 times before writing the result back to global memory. This example establishes the three recurring elements: special memory scopes, specialized instructions, and surrounding loops that a programmer can still reorder.

Representing bulk computation as a tensorized block (13:12)

The lecture translates the NumPy model into TensorIR. Instead of a block computing one output element, a new block computes a 16x16 sub-region in one step, using spatial and reduction block iterators the same way scalar blocks do. An initialization region zeroes the accumulator on the first reduction step, and an update region performs the multiply-accumulate. Because the block declares its read and write regions, code outside the block can be transformed (split, reordered) without knowing whether the block body is a plain loop or an accelerated instruction - this isolation is what lets blockize convert an ordinary loop nest into this block form automatically.

Tensor intrinsics: description and implementation (27:22)

To avoid writing a custom backend for every accelerator, MLC registers pairs of functions called tensor intrinsics: a description function that specifies the computation a hardware instruction performs (reads from two 16x16 regions, writes to one), and an implementation function that specifies how to carry it out, such as calling an external microkernel. Buffers in the intrinsic declare memory scopes and stride constraints (for example, requiring 16-aligned dimensions), so the compiler can match a tensorized block against the description and substitute the implementation.

Running tensorization end to end (33:29)

The lecture decomposes the reduction into its init and update steps, calls tensorize to replace the update step with a call to the registered intrinsic, and shows the block body turning into a call to an external function backed by a small C microkernel imported through LLVM. The resulting program still runs on CPU in this demo (no real accelerator memory or instruction is used), but the same transformation pipeline applies unchanged when targeting real hardware such as an NVIDIA tensor core, which the lecture points to as optional follow-up material.

Before you watch

• Be comfortable with TensorIR blocks, block iterators, and scheduling primitives such as split and reorder from the CPU and GPU lectures earlier in this course.
• Review how the previous episode (GPU and Specialized Hardware, Part 1) builds an MLC flow for GPUs, since this lecture extends that same flow to specialized accelerators.

Check your understanding

1. What three elements does the lecture identify as common to programming any specialized accelerator, based on the low-level NumPy example?
2. Why does a tensorized block need to declare its read and write regions, and how does that help the blockize and scheduling transformations?
3. What is the difference between the description and implementation halves of a tensor intrinsic, and why are both needed?
4. Why does the lecture say specialization instructions still leave room to reorder outer loops?