Large-Scale Discrete Fourier Transform on TPUs

02/09/2020
by   Tianjian Lu, et al.
0

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 × 8192 × 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.

READ FULL TEXT

page 1

page 2

research
02/12/2020

CROFT: A scalable three-dimensional parallel Fast Fourier Transform (FFT) implementation for High Performance Clusters

The FFT of three-dimensional (3D) input data is an important computation...
research
06/24/2020

Accelerating MRI Reconstruction on TPUs

The advanced magnetic resonance (MR) image reconstructions such as the c...
research
06/07/2019

Lightweight Parallel Foundations: a model-compliant communication layer

We present the Lightweight Parallel Foundations (LPF), an interoperable ...
research
08/02/2018

Parallelization of the FFT on SO(3)

In this paper, a work-optimal parallelization of Kostelec and Rockmore's...
research
08/16/2018

Novel Model-based Methods for Performance Optimization of Multithreaded 2D Discrete Fourier Transform on Multicore Processors

In this paper, we use multithreaded fast Fourier transforms provided in ...
research
07/20/2018

Interactive Supercomputing on 40,000 Cores for Machine Learning and Data Analysis

Interactive massively parallel computations are critical for machine lea...
research
03/22/2022

Minimizing communication in the multidimensional FFT

We present a parallel algorithm for the fast Fourier transform (FFT) in ...

Please sign up or login with your details

Forgot password? Click here to reset