Now, we have learned how to optimize matrix multiplication on GPUs using tiling. Let's see how we can implement it with actual code.

For these optimized operations, we usually implement them as GPU kernels. A GPU kernel is a function that runs on the GPU, where we can control parallelization and memory access more directly. After writing the kernel, the code is compiled into a GPU binary and deployed to the GPU for execution.

There are multiple languages to write GPU kernels, like CUDA, Triton, and Helion. This time, we will use Helion to implement the matrix multiplication kernel for its simplicity.

The key abstraction in Helion is hl.tile. To tile a vector of size 10, we call hl.tile(10), and it picks an optimal tile size automatically. If the tile size is 4, it produces the slices 0:4, 4:8, and 8:10 in sequence. Each slice can be used to index into a tensor to get the corresponding chunk. For two-dimensional tiling, hl.tile([m, n]) yields a pair of slices, one per dimension.

Here is what the matmul kernel looks like written using Helion. hl.tile([num_rows, num_cols]) iterates over tiles of the output C, and hl.tile(k) iterates over tiles of the inner dimension, which is where the memory-optimized tiling from the previous sections happens. hl.zeros([rows, cols]) creates a zero-initialized accumulator for each output tile in fast memory. Each time a new tile of A and B is loaded, we call torch.addmm, which computes the inner products from those tiles and accumulates them into the output tile.

import torch, helion, helion.language as hl

@helion.kernel()
def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    # Initialize C in global memory.
    num_rows, k = a.size()
    k, num_cols = b.size()
    c = torch.empty([num_rows, num_cols], dtype=a.dtype, device=a.device)

    # Iterate the tiles of C.
    for rows, cols in hl.tile([num_rows, num_cols]):
        c_tile = hl.zeros([rows, cols], dtype=torch.float32)
        # Iterate the row splits of A and column splits of B.
        for k_range in hl.tile(k):
            # Update c_tile by computing inner products from a_tile, b_tile.
            c_tile = torch.addmm(c_tile, a[rows, k_range], b[k_range, cols])
        # Store c_tile back to global memory.
        c[rows, cols] = c_tile
    return c

In a modern kernel language like Helion, we do not need to tune the tile size or manually load and store data between global memory and shared memory. These are all managed by the compiler and runtime, which tunes the tile size automatically and pipelines memory access and computation for us.