Loop Reordering and Unrolling

Efficient kernel execution relies on aligning software instructions with the physical realities of hardware. While loop tiling optimizes for cache capacity, it does not inherently address how data is retrieved from a cache line or how efficiently the CPU pipeline is utilized. Loop reordering and unrolling are transformations that target spatial locality and instruction-level parallelism (ILP).

These transformations change the execution order and structure without altering the semantic result. For a deep learning compiler, the objective is to permute the iteration space so that memory accesses are contiguous and the instruction pipeline remains saturated.

Memory Layout and Spatial Locality

To understand why loop order matters, we must look at how multidimensional tensors are flattened into linear memory. In C/C++ and most deep learning frameworks (like PyTorch or TVM default), tensors are stored in row-major order. This means the last dimension changes the fastest in memory.

Consider a tensor $A$ of shape $(N, M)$. The element $A_{i,j}$ is stored at address base + i*M + j.

• Accessing $A_{i,j}$ then $A_{i,j+1}$ (incrementing the inner index) results in a stride-1 access pattern. This is optimal because loading $A_{i,j}$ brings adjacent elements into the cache line.
• Accessing $A_{i,j}$ then $A_{i+1,j}$ (incrementing the outer index) results in a stride-$M$ access pattern. If $M$ is large, every access might trigger a cache miss.

The following diagram illustrates the difference in memory traversal between stride-1 access and strided access.

Memory access patterns dictate cache efficiency. Stride-1 access utilizes the full cache line, whereas strided access effectively wastes memory bandwidth.

Loop Permutation in Matrix Multiplication

The classic example of loop reordering is Matrix Multiplication (MatMul). The mathematical definition involves three indices: $i$, $j$, and $k$.

$C_{i,j} += A_{i,k} \times B_{k,j}$

A naive implementation typically uses the $i-j-k$ order:

for (int i = 0; i < N; i++) {
    for (int j = 0; j < M; j++) {
        for (int k = 0; k < K; k++) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

In this configuration:

1. $C[i][j]$ is invariant in the inner loop (Good).
2. $A[i][k]$ is accessed with stride-1 (Good).
3. $B[k][j]$ is accessed with stride-$M$ (Bad).

Because $B$ is accessed column-wise while stored row-wise, this schedule suffers from severe cache thrashing. By changing the loop order to $i-k-j$, we alter the access pattern:

for (int i = 0; i < N; i++) {
    for (int k = 0; k < K; k++) {
        for (int j = 0; j < M; j++) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

In the $i-k-j$ permutation:

1. $A[i][k]$ is invariant in the inner loop.
2. $B[k][j]$ is accessed with stride-1 (Good).
3. $C[i][j]$ is accessed with stride-1 (Good).

This simple reordering can result in significant speedups, often 5x to 10x on standard CPUs, merely by respecting the hardware's preference for linear access. Compilers model this using Polyhedral analysis to determine dependency-free permutations that maximize spatial locality.

Loop Unrolling

While reordering optimizes data movement, loop unrolling optimizes the instruction stream. Unrolling involves replicating the loop body multiple times within a single iteration and decreasing the loop counter frequency.

Consider a standard dot product loop:

for (int i = 0; i < N; i++) {
    sum += A[i] * B[i];
}

A CPU executing this code encounters overhead: incrementing i, comparing i to N, and a branch instruction. Additionally, scalar operations often fail to fill the processor's multiple execution units.

An unrolled version might look like this:

for (int i = 0; i < N; i += 4) {
    sum += A[i] * B[i];
    sum += A[i+1] * B[i+1];
    sum += A[i+2] * B[i+2];
    sum += A[i+3] * B[i+3];
}

Unrolling provides three specific advantages:

1. Reduction of Loop Overhead: The comparison and branch instructions are executed $N/4$ times instead of $N$ times.
2. Instruction Level Parallelism (ILP): Modern CPUs are superscalar and out-of-order. Independent instructions within the unrolled block can be dispatched simultaneously to different arithmetic logic units (ALUs).
3. Vectorization Opportunities: Compilers can easily map unrolled scalar operations to SIMD vector instructions (e.g., AVX-512 or NEON), processing the entire block in a single cycle.

The Optimization Space and Trade-offs

Unrolling is not a monotonic optimization where "more is better." It introduces a trade-off between pipeline saturation and register pressure.

1. Register Pressure: Each unrolled operation may require temporary registers to hold intermediate values. If the compiler unrolls too aggressively, the CPU runs out of architectural registers. This forces "register spilling," where variables are moved to the stack (L1 cache/RAM), causing a severe performance penalty that negates the benefits of unrolling.
2. Instruction Cache (I-Cache): Unrolling increases the binary size of the program. If the inner loop becomes too large, it may no longer fit in the instruction cache, causing pipeline stalls as the CPU fetches instructions from main memory.

The following chart demonstrates the performance characteristics of different unroll factors. Notice the diminishing returns and eventual regression as register pressure increases.

Normalized throughput increases with the unroll factor as loop overhead diminishes and ILP improves. However, past a factor of 8, performance degrades due to register spilling and instruction cache misses.

Compiler Implementation Strategy

In frameworks like TVM or MLIR, these transformations are applied during the lowering phase. The compiler does not guess randomly; it uses cost models (discussed in Chapter 6) to select parameters.

A typical transformation recipe for a dense kernel involves a composite strategy:

1. Reorder the loops to ensure the innermost loop accesses the largest tensor with stride-1.
2. Split the loops (tiling) to fit the working set into L1 cache.
3. Unroll the innermost tile loops to a factor determined by the target architecture's register count (e.g., unroll factor 4 or 8 for standard float32 on x86).

By combining reordering to fix memory access patterns and unrolling to saturate the instruction pipeline, the compiler generates a kernel that approaches the theoretical peak performance of the hardware.