DaggerFFT: A Distributed FFT Framework Using Task Scheduling in Julia

TL;DR

DaggerFFT introduces a dynamic, task-based distributed FFT framework implemented in Julia, replacing static domain decompositions with asynchronously scheduled DTasks operating on per-stage distributed arrays. By integrating asynchronous MPI-based redistribution and a locality-aware scheduler with plan caching, it achieves strong performance on both CPU and GPU backends, and outperforms several state-of-the-art libraries in key regimes. The framework is validated through extensive benchmarks and a real-world Oceananigans Poisson solver integration, demonstrating tangible speedups and modularity for large-scale geophysical simulations. This work highlights the practical potential of high-level dynamic runtimes to deliver exascale-ready FFT performance without hardware-specific tuning.

Abstract

The Fast Fourier Transform (FFT) is a fundamental numerical technique with widespread application in a range of scientific problems. As scientific simulations attempt to exploit exascale systems, there has been a growing demand for distributed FFT algorithms that can effectively utilize modern heterogeneous high-performance computing (HPC) systems. Conventional FFT algorithms commonly encounter performance bottlenecks, especially when run on heterogeneous platforms. Most distributed FFT approaches rely on static task distribution and require synchronization barriers, limiting scalability and impacting overall resource utilization. In this paper we present DaggerFFT, a distributed FFT framework, developed in Julia, that treats highly parallel FFT computations as a dynamically scheduled task graph. Each FFT stage operates on a separately defined distributed array. FFT operations are expressed as DTasks operating on pencil or slab partitioned DArrays. Each FFT stage owns its own DArray, and the runtime assigns DTasks across devices using Dagger's dynamic scheduler that uses work stealing. We demonstrate how DaggerFFT's dynamic scheduler can outperform state-of-the-art distributed FFT libraries on both CPU and GPU backends, achieving up to a 2.6x speedup on CPU clusters and up to a 1.35x speedup on GPU clusters. We have integrated DaggerFFT into Oceananigans.jl, a geophysical fluid dynamics framework, demonstrating that high-level, task-based runtimes can deliver both superior performance and modularity in large-scale, real-world simulations.

Figures (15)

• Figure 1: Comparison between standard asynchronous redistribution and the proposed progressive per-chunk pipelining. In prior approaches (left), unpacking and subsequent FFTs begin only after all sends are complete, requiring implicit barriers. In our design (right), receives and unpacks occur progressively as messages arrive, enabling continuous overlap.
• Figure 2: Task-scheduled 3D FFT implementation using pencil decomposition, asynchronous transforms and data movement.
• Figure 3: Complete distributed FFT Framework, from a user's view, for resource assignment. The system dynamically constructs a task graph over stage-specific distributed arrays and maps tasks to computing resources using the scheduler.
• Figure 4: Effect of the scheduling runtime on the first FFT stage ($512^3$ grid, 16 ranks, one per core).
• Figure 5: Effect of work stealing under load imbalance.