Loop Nest Optimization (LNO)

Loop tiling (or blocking) partitions a large iteration space into smaller, regular blocks or tiles. This transformation is fundamental for fitting the working set of data into faster, limited memory hierarchies like an NPU's shared memory or register file, thereby drastically reducing accesses to slower global memory.

• Goal: Exploit temporal and spatial locality.
• Mechanism: Introduces new tile loops outside the original point loops.
• Key Parameter: Tile size, which is often auto-tuned based on hardware constraints (e.g., shared memory size).
• Example: Tiling a matrix multiplication loop to keep sub-matrices in shared memory for repeated use.

Loop interchange swaps the order of nested loops to improve memory access patterns and enable vectorization. On architectures with caches or coalesced memory access (like NPUs), accessing data in the order it is stored in memory is critical for performance.

• Goal: Transform memory access to stride-1 (contiguous) patterns.
• Impact: Enables memory coalescing on parallel architectures, where concurrent threads access contiguous memory locations combined into a single wide transaction.
• Example: Changing a column-major traversal to row-major in a C-style array to match memory layout, reducing cache misses.

Loop fusion combines two or more adjacent loops that have the same iteration space into a single loop. This reduces total loop overhead (branch instructions, index updates) and, more importantly, improves data locality by keeping producer and consumer data in fast memory (registers/cache) between operations.

• Goal: Reduce kernel launch overhead and improve data reuse.
• Relation to Kernel Fusion: Loop fusion is often a prerequisite for the more advanced kernel fusion, where separate GPU/NPU kernels are merged.
• Constraint: Loops must be control-flow equivalent and data-independent (or have manageable dependencies).

Loop unrolling replicates the body of a loop multiple times, decreasing the number of iteration checks and branch instructions. This reduces control overhead and increases opportunities for instruction-level parallelism (ILP) and register-level data reuse.

• Types: Full unrolling (loop is eliminated) and partial unrolling (factor N).
• Benefits: Exposes more independent operations for the scheduler, enables common subexpression elimination (CSE) across unrolled iterations.
• Trade-off: Increases register pressure and code size; excessive unrolling can lead to register spilling.

Loop fission (or distribution) is the inverse of fusion: it splits a single loop containing multiple statements into separate loops. This is used to isolate parallelizable sections, reduce register pressure within a loop body, or enable other transformations (like fusing a resulting sub-loop with another).

• Goal: Improve parallelism or manage resource constraints.
• Use Case: Separating a memory-intensive section from a compute-intensive section within a loop to optimize each independently.
• Challenge: Can negatively impact data locality by separating producer and consumer statements, potentially increasing memory traffic.

Loop unrolling and jam (or unroll-and-jam) is a combined transformation. First, an outer loop is unrolled. Then, the resulting inner loops (from different outer loop iterations) are fused (jammed) together. This powerful technique increases the granularity of data reuse in the inner loop.

• Goal: Enhance data locality and enable inner-loop vectorization.
• Mechanism: Creates larger inner loop bodies that operate on multiple elements from different outer loop iterations, increasing arithmetic intensity.
• Example: Critically used in optimizing linear algebra kernels (e.g., GEMM) by allowing multiple elements of a row to be reused from registers.

Loop tiling partitions a large iteration space into smaller, regular blocks or tiles. This transformation is critical for fitting the working set of data into faster, limited memory hierarchies like shared memory (on GPUs/NPUs) or CPU caches.

• Primary Goal: Improve data locality and reduce accesses to slower global/DRAM memory.
• AI Application: Enables efficient execution of large matrix multiplications (GEMM) and convolutions by keeping tile-sized sub-matrices in on-chip memory.
• Parameter Tuning: Optimal tile size is hardware-dependent and is often discovered via auto-tuning.

Loop interchange swaps the order of nested loops to improve memory access patterns. This is essential for aligning loop iterations with the data layout in memory.

• Key Benefit: Enables memory coalescing on parallel architectures, where concurrent threads access contiguous memory addresses, maximizing effective bandwidth.
• Example: Transforming a loop nest from column-major access to row-major access (or vice-versa) to match the storage order of a multi-dimensional array.
• Impact: A poorly ordered nest can cause thrashing in cache lines or bank conflicts in shared memory, drastically reducing performance.

These are complementary transformations that restructure loops to balance parallelism, locality, and resource pressure.

• Loop Fusion: Combines two adjacent loops with the same iteration space. Reduces loop overhead and improves locality by keeping producer-consumer data in registers/cache.
• Loop Fission (Distribution): Splits a single loop into multiple loops. Reduces register pressure, enables parallelization of independent sub-loops, or isolates code for further targeted optimization.
• Trade-off: Fusion improves locality but may increase register usage and limit parallelism; fission does the opposite. Compilers use cost models to decide.

Loop unrolling replicates the body of a loop multiple times per iteration, decreasing the number of loop control instructions executed.

• Objectives:
  • Reduce branch overhead and increase Instruction-Level Parallelism (ILP).
  • Expose more operations for software pipelining.
  • Improve scheduling flexibility for the compiler.
• Considerations: Excessive unrolling can increase register pressure, potentially leading to register spilling, and can bloat instruction cache footprint. The optimal unroll factor is often determined empirically.

While not strictly a loop transformation, LNO often drives changes to how data is structured in memory to complement optimized access patterns.

• Padding: Adding unused elements to array dimensions to avoid cache line thrashing or shared memory bank conflicts.
• Array Transposition: Physically reorganizing a multi-dimensional array in memory to match the access pattern of the optimized loop nest.
• Structure-of-Arrays (SoA) to Array-of-Structures (AoS): Changing data layouts to enable efficient vectorized or coalesced memory accesses. Crucial for stencil computations and particle simulations in HPC.

Advanced LNO for complex, non-perfectly nested loops is often performed using the polyhedral model, a mathematical framework for reasoning about loop transformations.

• Mechanism: Represents loop nests as sets of integer points in polyhedra. Transformations (tiling, interchange, skewing) are represented as affine schedules.
• Automation: Allows compilers (like LLVM's Polly) to automatically explore a vast space of legal transformations to find an optimal schedule.
• Application: Essential for optimizing irregular kernels in scientific computing and the complex, data-dependent loops that can appear in AI models post-fusion.

A compiler optimization that merges multiple separate computational kernels into a single, larger kernel. This reduces the overhead of repeated kernel launches and eliminates intermediate data transfers between global memory and on-chip storage. It is a higher-level transformation often enabled by first applying loop-level optimizations like loop fusion to create a single, optimizable loop nest.

A critical loop nest transformation that partitions a large iteration space into smaller, regular blocks or tiles. The primary goal is to fit the working set of data into a faster, limited memory hierarchy (e.g., shared memory, registers, or cache), thereby drastically reducing expensive accesses to slower global memory. Choosing the optimal tile size is a key challenge addressed by auto-tuning.

An optimization for parallel architectures where concurrent memory accesses from multiple threads (within a warp or wavefront) are combined into a single, wider memory transaction. LNO techniques like loop interchange are often applied specifically to transform memory access patterns (e.g., from strided to contiguous) to enable coalescing, maximizing effective memory bandwidth.

The automated, empirical process of searching a parameter space to find the optimal configuration for a kernel on specific hardware. For LNO, this involves searching parameters like:

• Tile sizes for different memory hierarchies
• Loop unroll factors
• Loop interchange orders

The search uses metrics like execution time or a theoretical model like the Roofline Model to guide the optimization.

A key performance metric defined as the number of arithmetic operations performed per byte of data transferred between the processor and main memory (DRAM). LNO transformations directly aim to increase arithmetic intensity by improving data locality (via tiling) and reusing data in fast memory. This shifts a kernel's performance from being memory-bound toward being compute-bound.

A measure of the number of independent instructions that can be executed simultaneously within a single thread. LNO techniques like loop unrolling expose more independent operations to the compiler and hardware, increasing ILP. This helps hide instruction and memory latency by keeping the processor's functional units busy, a concept further exploited by software pipelining.