Abstract:In this work, we present two parallel algorithms for the large-scale discrete Fourier transform (DFT) on Tensor Processing Unit (TPU) clusters. The two parallel algorithms are associated with two formulations of DFT: one is based on the Kronecker product, to be specific, dense matrix multiplications between the input data and the Vandermonde matrix, denoted as KDFT in this work; the other is based on the famous Cooley-Tukey algorithm and phase adjustment, denoted as FFT in this work. Both KDFT and FFT formulations take full advantage of TPU's strength in matrix multiplications. The KDFT formulation allows direct use of nonuniform inputs without additional step. In the two parallel algorithms, the same strategy of data decomposition is applied to the input data. Through the data decomposition, the dense matrix multiplications in KDFT and FFT are kept local within TPU cores, which can be performed completely in parallel. The communication among TPU cores is achieved through the one-shuffle scheme in both parallel algorithms, with which sending and receiving data takes place simultaneously between two neighboring cores and along the same direction on the interconnect network. The one-shuffle scheme is designed for the interconnect topology of TPU clusters, minimizing the time required by the communication among TPU cores. Both KDFT and FFT are implemented in TensorFlow. The three-dimensional complex DFT is performed on an example of dimension $8192 \times 8192 \times 8192$ with a full TPU Pod: the run time of KDFT is 12.66 seconds and that of FFT is 8.3 seconds. Scaling analysis is provided to demonstrate the high parallel efficiency of the two DFT implementations on TPUs.

An in-place truncated Fourier transform and applications to polynomial multiplication

Frobenius Additive Fast Fourier Transform

FourierPIM: High-Throughput In-Memory Fast Fourier Transform and Polynomial Multiplication

Fast Forth Power and Its Application in Inversion Computation for a Special Class of Trinomials

Some New Parallel Fast Fourier Transform Algorithms

XFT: An Extension of the Discrete Fractional Fourier Transform

Quantum Fourier Transform Revisited

Fourier optimization, the least quadratic non-residue, and the least prime in an arithmetic progression

Fast Fourier transforms for the evaluation of convolution products: CPU versus GPU implementation

Faster truncated integer multiplication

Faster Integer Multiplication Using Preprocessing

Efficient Algorithm for 2-D Arithmetic Fourier Transform

Fast Approximate Polynomial Multipoint Evaluation and Applications

A Split Fast Fourier Transform Algorithm for Block Toeplitz Matrix-Vector Multiplication

Did Fourier Really Meet Möbius? Fast Subset Convolution via FFT

Non-Abelian Fourier Transforms and Normalized Intertwining Operators for General Parabolics over Finite Fields, and the Kloosterman Fourier Transform for Quadric Cones

Practical Implementation of the Arithmetic Fourier Transform

Large-Scale Discrete Fourier Transform on TPUs

A polynomial time algorithm for calculating Fourier-Dedekind sums