Shortest-Path FFT: Optimal SIMD Instruction Scheduling via Graph Search

Abstract: An $N$-point FFT admits many valid implementations that differ in radix choice, stage ordering, and register-blocking strategy. These alternatives use different SIMD instruction mixes with different latencies, yet produce the same mathematical result. We show that finding the fastest implementation is a shortest-path problem on a directed acyclic graph. We formalize two variants of this graph. In the \emph{context-free} model, nodes represent computation stages and edge weights are independently measured instruction costs. In the \emph{context-aware} model, nodes are expanded to encode the \emph{predecessor edge type}, so that edge weights capture inter-operation correlations such as cache warming -- the cost of operation~B depends on which operation~A preceded it. This addresses a limitation identified but deliberately bypassed by FFTW \citep{FrigoJohnson1998}: that optimal-substructure assumptions break down ``because of the different states of the cache.'' Applied to Apple M1 NEON, the context-free Dijkstra finds an arrangement at 22.1~GFLOPS (74\% of optimal). The context-aware Dijkstra discovers $\text{R4} \to \text{R2} \to \text{R4} \to \text{R4} \to \text{Fused-8}$ at 29.8~GFLOPS -- a $5.2\times$ improvement over pure radix-2 and 34\% faster than the context-free result. This arrangement includes a radix-2 pass \emph{sandwiched between} radix-4 passes, exploiting cache residuals that only exist in context. No context-free search can discover this.

• The paper introduces a context-aware graph search approach that models first-order dependencies to optimize SIMD instruction scheduling for FFTs.
• It demonstrates that a context-aware DAG yields a 34% performance improvement over context-free models and a 5.2× speedup over pure radix-2 scheduling.
• The framework is hardware-sensitive, enabling architecture-specific scheduling optimizations that can extend to other multi-stage computational pipelines.

Optimal SIMD Instruction Scheduling for FFT via Context-Aware Graph Search

Problem Formulation and Motivation

The Fast Fourier Transform (FFT) is a central computational primitive whose efficient implementation on modern hardware demands close attention to low-level instruction scheduling, cache locality, and register resources. Despite all valid decompositions of the Cooley-Tukey FFT algorithm producing identical results, their instruction-level implementations differ drastically in terms of achievable throughput due to variations in algorithm stages (radix-2, radix-4, radix-8), register blocking, and cache behavior.

Historically, tools like FFTW and SPIRAL have approached this challenge using heuristic or empirical search strategies, often relying on dynamic programming with the optimal substructure assumption—that is, claiming that the fastest way to transform a partial result is independent of its computational context. However, these assumptions break down due to real-world effects such as cache state transitions and register pressure, which create first-order dependencies between scheduling decisions.

This work systematically addresses and resolves the limitations of the optimal substructure assumption by recasting the problem as a shortest-path search in a specifically constructed directed acyclic graph (DAG). Critically, the graph is augmented to encode context, such that edge weights capture not only the cost of an operation but also how that cost depends on the operation's preceding context—primarily through cache and register state.

Context-Free vs. Context-Aware Graph Models

The baseline, context-free DAG formulation represents each computation stage as a node and each valid FFT scheduling alternative (e.g., a radix-2 pass, a fused register block) as an edge, with empirically measured cost as the edge weight. Solution paths from initial to final nodes represent specific FFT decomposition choices; the shortest path corresponds to the fastest scheduling.

The context-aware variant expands the node-space to explicitly encode the predecessor edge (operation type), so that each node represents not just the current computation stage but the kind of operation that preceded it. This allows edge weights to be conditioned on the operation history, thus capturing inter-operation effects such as cache warming and register spill/fill costs. Measurement protocols are adapted: for each potential transition, the previous operation is executed (untimed) before the candidate operation is timed, providing a realistic cost measure that reflects memory and cache residencies.

The context-aware model introduces a state-space size increase by a multiplicative factor of the number of edge types but remains computationally tractable (e.g., 77 nodes for $N=1024$) and, crucially, models performance-affecting correlations that context-free abstractions miss.

Implementation Details on Apple M1 NEON

The graph-based framework is evaluated on the Apple M1 NEON SIMD engine. The implementation comprises highly optimized butterfly cores (in split-complex storage), and several fused register-blocked FFT kernels (newly including an FFT-32 block, uniquely suited to NEON's 32-register architecture).

Empirical measurement demonstrates that maximizing block size is not universally optimal; register pressure penalizes large fused blocks (FFT-32 underperforms FFT-16 and FFT-8 due to insufficient registers for twiddle storage). The search framework automatically detects these tradeoffs due to direct measurement of each planning alternative in realistic context.

Experimental Results

The context-aware graph search demonstrates strong quantitative improvements:

• Context-aware search achieves 29.8 GFLOPS on M1 NEON for $N=1024$ split-complex FFT.
• This is a $5.2\times$ speedup over pure radix-2 scheduling and 34% faster than the best context-free search result.
• The identified optimal sequence is R4 $\to$ R2 $\to$ R4 $\to$ R4 $\to$ Fused-8—a pattern that would not have been proposed by any context-free or maximize-radix search, due to the non-obvious inclusion of a radix-2 pass whose cost is beneficial only because of the specific cache state left by the prior R4.
• Classical “maximize the radix at every step” strategies fail: pure radix-8 variants achieve only 25% of optimal throughput.
• The framework is architecture-sensitive: the optimal schedule for AVX2 (Intel Haswell) diverges drastically from M1 NEON, confirming that hardware-specific costs must be empirically measured and integrated into decomposition search.

The approach’s computational burden for search and measurement is negligible compared to even established libraries’ planning overheads (e.g., FFTW’s codelet evaluation).

Theoretical and Practical Implications

By explicitly modeling first-order Markov dependencies in stage scheduling, the framework formalizes, quantifies, and operationalizes an effect recognized but previously sidestepped in FFT research (especially in FFTW and SPIRAL). The empirical performance difference underscores the inadequacy of context-free decompositions and highlights the methodological advantage of incorporating context into instruction scheduling.

This approach generalizes beyond FFT to any staged signal processing or scientific computing pipeline where low-level operations have non-independent performance characteristics and where prior state (cache, register, pipeline) is a nontrivial factor. For instance, matrix factorizations, multi-stage filters, and neural network layer fusion can benefit from similar context-aware shortest-path search formulations.

The model also enables straightforward porting to new hardware: edge weights are re-measured, and the shortest-path search is rerun, yielding a new schedule matched to the actual microarchitectural details of the new target.

Conclusion

Context-aware shortest-path graph search offers a principled, formal solution for optimal SIMD instruction scheduling of FFTs, capturing cache and register correlations neglected by previous approaches. The demonstrated 34% gain over context-free models validates the method’s necessity on modern CPUs. Architecture-specific, non-obvious decompositions can be discovered automatically, obviating the need for human-crafted heuristics or oversimplified dynamic programming. The method is of broad applicability wherever stage-local costs are context-sensitive.

Source code is available at: https://github.com/aminems/fft.

Reference: "Shortest-Path FFT: Optimal SIMD Instruction Scheduling via Graph Search" (2604.04311).