Automatic Generation of Mappings for Distributed Fourier Operations

The Fourier transform is an ubiquitous mathematical operation used in a multitude of scientific applications. Most distributed Fourier transform libraries provide rigid implementations that force developers of high performance applications to mold their code around the Fourier computation, omitting opportunities for minimizing communication across the Fourier transforms and the surrounding computation. In this work, we introduce a new automatic approach to generate distributed mappings for multi-dimensional Fourier operations, offering a solution to this problem. Our approach decides how to decompose, map, and schedule the computation as smaller and lower-dimensional parallel operations. We design and implement a novel non-linear iterative formulation that optimizes across Fourier and linear algebra operations. Our scheme leverages the Z3 SMT solver to minimize the number of communication steps across key MPI collectives, while selecting the grid shape. We evaluate the effectiveness of our new scheme and demonstrate 2 × -31 × speedups over coupled heFFTe and COSMA solutions.