Fusion, layout, scheduling, and memory planning

Algebraic simplification and common expressions

Continue with the running graph `y = GELU(x @ w + bias)`. Suppose shape code also computes `hidden = 2048 * 2`, the graph contains `transpose(transpose(w))`, and an earlier diagnostic branch produces a tensor that no output reads. Before scheduling a GPU kernel, the compiler can simplify these concrete pieces while leaving the matrix multiplication, bias, and GELU behavior unchanged.

The inputs `2048` and `2` are known while the program is compiled, so the compiler evaluates the multiplication once and replaces it with `4096`. This rewrite is constant folding. The runtime no longer needs an operation to calculate the hidden dimension, and the known value can make a fixed tile or vector width legal.

Applying a transpose twice with the same two-dimensional permutation returns the original logical layout. Replacing `transpose(transpose(w))` with `w` is an algebraic simplification. Removing the pair can also prevent a later physical layout conversion. The rewrite is legal only when the transpose operations have the assumed semantics and no intervening mutation changes the value.

If two branches independently compute the same pure expression, such as `shape_of(x)[1]`, the compiler can retain one result and direct both readers to it. This rewrite is common-subexpression elimination. The compiler cannot make that substitution merely because two calls have the same printed arguments when a call reads mutable state, produces randomness, or has another side effect.

The unused diagnostic tensor has no reader and does not affect an output, mutation, exception, or other observable behavior, so dead-code elimination removes its producing operations. Constant folding and removal of the transpose pair change later opportunities: the constant `4096` can expose a static tile, and the original weight layout can now remain available for fusion. Floating-point reassociation is more restricted because changing the order of additions can change rounding; the declared numerical mode determines whether such a rewrite is permitted.

Fusion and layout propagation

Producer-consumer fusion computes a value near the operation that consumes it, which can remove an HBM intermediate. Sibling fusion lets two outputs share a common input traversal. A reduction, transpose, or multiple-consumer graph can impose incompatible iteration spaces. Resource estimates must account for all live values in the fused schedule.

Layout propagation selects physical representations that serve a chain of operations. A layout suited to a matrix instruction may differ from one suited to a following reduction. Inserting a conversion has a cost, but forcing one layout through the complete region can make several operations inefficient. The compiler must compare these alternatives with the full consumer graph.

Schedule selection

A schedule chooses tile sizes, loop order, parallel dimensions, vector widths, pipeline stages, and memory placement. The schedule parameters determine reuse, coalescing, instruction eligibility, register count, shared-memory allocation, and available parallel blocks. Two programs with identical arithmetic can differ substantially because of their schedules.

The schedule space is too large for exhaustive search on every input. Heuristics remove unlikely candidates. Analytical cost models rank candidates from estimated work, bytes, and resources. Autotuning measures a bounded set. A model trained or tuned on square matrices can rank narrow decode shapes poorly, so workload coverage and cache keys are part of schedule-system correctness.

Memory planning and rematerialization

Buffer planning uses value lifetimes and alias rules to assign storage. Values whose lifetimes do not overlap can reuse the same buffer. In-place updates can remove allocations only when no later operation needs the original value and aliasing permits mutation. Reuse can reduce peak memory but can also add dependencies that limit execution overlap.

Rematerialization discards a value and computes the value again near a later use. Rematerialization increases arithmetic but can reduce saved-tensor storage, high-bandwidth-memory traffic, or live range. The rewrite is profitable when the recomputed operation is inexpensive relative to the storage and movement avoided. FlashAttention backward is a concrete large-scale example: the backward pass recomputes score blocks instead of saving the complete probability matrix.

Specialization and compilation cost

Known shapes and alignments permit unrolling, fixed tiles, vector operations, and removal of bounds checks. A dynamic implementation covers more inputs but may retain predicates and conservative schedules. Guarded variants combine these approaches by specializing important cases and using a general fallback.

Every variant consumes compilation time, tuning time, code-cache memory, and maintenance attention. Include these costs when shapes are rare or processes are short lived. Profile-guided frequency can concentrate specialization on shapes that contribute most time or latency. Do not assess an optimization pipeline only from warmed execution of one compiled graph.