2020

Large-Scale Discrete Fourier Transform on TPUs

Lu, Tianjian, Chen, Yi-Fan, Hechtman, Blake et al.

Understand

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.

Built on

J. W. Cooley and J. W. Tukey, “An algorithm for the machine calculation of complex Fourier series,” Mathematics of computation , vol. 19, no. 90, pp. 297–301, 1965
1965
Earlier work this paper cites.
G. Ackins, “Fast fourier transform via ILLIAC IV,” ILLIAC IV Document , no. 146, 1968
1968
Earlier work this paper cites.
J. E. Stevens Jr et al. , “A fast Fourier transform subroutine for ILLIAC IV,” CAC document; no. 17 , 1971
1971
Earlier work this paper cites.
P. Duhamel and H. Hollmann, “Split radix FFT algorithm,” Electronics letters , vol. 20, no. 1, pp. 14–16, 1984
1984
Earlier work this paper cites.
R. N. Bracewell and R. N. Bracewell, The Fourier transform and its applications . McGraw-Hill New York, 1986, vol. 31999
1986
Earlier work this paper cites.
P. N. Swarztrauber, “Multiprocessor FFTs,” Parallel computing , vol. 5, no. 1-2, pp. 197–210, 1987
1987
Earlier work this paper cites.
D. H. Bailey, “FFTs in external or hierarchical memory,” The journal of Supercomputing , vol. 4, no. 1, pp. 23–35, 1990
1990
Earlier work this paper cites.
A. Gupta and V. Kumar, “The scalability of FFT on parallel computers,” IEEE Transactions on Parallel and Distributed Systems , vol. 4, no. 8, pp. 922–932, 1993
1993
Earlier work this paper cites.
S. Bagchi and S. K. Mitra, “The nonuniform discrete Fourier transform and its applications in filter design. I. 1-D,” IEEE Transactions on Circuits and Systems II: Analog and Digital Signal Processing , vol. 43, no. 6, pp. 422–433, 1996
1996
Earlier work this paper cites.
——, “The nonuniform discrete Fourier transform and its applications in filter design. II. 2-D,” IEEE Transactions on Circuits and Systems II: Analog and Digital Signal Processing , vol. 43, no. 6, pp. 434–444, 1996
1996
Earlier work this paper cites.
S. K. Nag and H. K. Verma, “Method for parallel-efficient configuring an FPGA for large FFTs and other vector rotation computations,” Feb. 1 2000, US Patent 6,021,423
2000
Earlier work this paper cites.
P. A. Regalia and M. K. Sanjit, “Kronecker products, unitary matrices and signal processing applications,” SIAM review , vol. 31, no. 4, pp. 586–613, 1989
Original
2002
Earlier work this paper cites.
A. Grama, V. Kumar, A. Gupta, and G. Karypis, Introduction to parallel computing . Pearson Education, 2003
2003
Earlier work this paper cites.

Similar

M. Frigo and S. G. Johnson, “The design and implementation of FFTW3,” Proceedings of the IEEE , vol. 93, no. 2, pp. 216–231, 2005
2005
Cited alongside, same era.
J.-Y. Lee and L. Greengard, “The type 3 nonuniform FFT and its applications,” Journal of Computational Physics , vol. 206, no. 1, pp. 1–5, 2005
2005
Cited alongside, same era.
A. Rahimi and B. Recht, “Random features for large-scale kernel machines,” in Advances in neural information processing systems , 2008, pp. 1177–1184
2008
Cited alongside, same era.
R. Cont, Frontiers in quantitative finance: Volatility and credit risk modeling . John Wiley & Sons, 2009, vol. 519
2009
Cited alongside, same era.
D. Takahashi, “An implementation of parallel 3-D FFT with 2-D decomposition on a massively parallel cluster of multi-core processors,” in International Conference on Parallel Processing and Applied Mathematics . Springer, 2009, pp. 606–614
2009
Cited alongside, same era.
Y. Chen, X. Cui, and H. Mei, “Large-scale FFT on GPU clusters,” in Proceedings of the 24th ACM International Conference on Supercomputing . ACM, 2010, pp. 315–324
2010
Cited alongside, same era.
C.-L. Yu, K. Irick, C. Chakrabarti, and V. Narayanan, “Multidimensional DFT IP generator for FPGA platforms,” IEEE Transactions on Circuits and Systems I: Regular Papers , vol. 58, no. 4, pp. 755–764, 2010
2010
Cited alongside, same era.
D. Pekurovsky, “P3DFFT: A framework for parallel computations of Fourier transforms in three dimensions,” SIAM Journal on Scientific Computing , vol. 34, no. 4, pp. C192–C209, 2012
2012
Cited alongside, same era.
G. B. Giannakis, F. Bach, R. Cendrillon, M. Mahoney, and J. Neville, “Signal processing for big data [from the guest editors],” IEEE Signal Processing Magazine , vol. 31, no. 5, pp. 15–16, 2014
2014
Cited alongside, same era.
M. Garrido, M. Acevedo, A. Ehliar, and O. Gustafsson, “Challenging the limits of FFT performance on FPGAs,” in 2014 International Symposium on Integrated Circuits (ISIC) . IEEE, 2014, pp. 172–175
2014
Cited alongside, same era.
E. Olshannikova, A. Ometov, Y. Koucheryavy, and T. Olsson, “Visualizing big data with augmented and virtual reality: challenges and research agenda,” Journal of Big Data , vol. 2, no. 1, p. 22, 2015
2015
Cited alongside, same era.
Y. Wu, M. Schuster, Z. Chen, Q. V. Le, M. Norouzi, W. Macherey, M. Krikun, Y. Cao, Q. Gao, K. Macherey et al. , “Google’s neural machine translation system: Bridging the gap between human and machine translation,” arXiv preprint arXiv:1609.08144 , 2016
Original
2016
Cited alongside, same era.

Then

I. Stoica, D. Song, R. A. Popa, D. Patterson, M. W. Mahoney, R. Katz, A. D. Joseph, M. Jordan, J. M. Hellerstein, J. E. Gonzalez et al. , “A Berkeley view of systems challenges for AI,” arXiv preprint arXiv:1712.05855 , 2017
Original
2017
Later among the works it cites.
N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, N. Boden, A. Borchers et al. , “In-datacenter performance analysis of a tensor processing unit,” in 2017 ACM/IEEE 44th Annual International Symposium on Computer Architecture (ISCA) . IEEE, 2017, pp. 1–12
2017
Later among the works it cites.
N. Jouppi. (2017) Quantifying the performance of the TPU, our first machine learning chip. [Online]. Available: https://cloud.google.com/blog/products/gcp/quantifying-the-performance-of-the-tpu-our-first-machine-learning-chip
2017
Later among the works it cites.
N. Ketkar, “Introduction to PyTorch,” in Deep learning with Python . Springer, 2017, pp. 195–208
2017
Later among the works it cites.
J. Kim. (2018) Leveraging optimized FFT on Intel Xeon platforms. [Online]. Available: https://www.alcf.anl.gov/support-center/training-assets/leveraging-optimized-fft-intel-xeon-platforms
2018
Later among the works it cites.
D. Takahashi, Fast Fourier transform algorithms for parallel computers . Springer, 2019
2019
Later among the works it cites.
D. T. Popovici, M. D. Schatz, F. Franchetti, and T. M. Low, “A flexible framework for parallel multi-dimensional DFTs,” arXiv preprint arXiv:1904.10119 , 2019
Original
2019
Later among the works it cites.
K. R. Roe, K. Hester, and R. Pascual. (2019) Multi-GPU FFT performance on different hardware configurations. [Online]. Available: https://developer.nvidia.com/gtc/2019/video/S9158
2019
Later among the works it cites.
A. Li, S. L. Song, J. Chen, J. Li, X. Liu, N. Tallent, and K. Barker, “Evaluating modern GPU interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect,” arXiv preprint arXiv:1903.04611 , 2019
Original
2019
Later among the works it cites.
K. Yang, Y.-F. Chen, G. Roumpos, C. Colby, and J. Anderson, “High performance Monte Carlo simulation of Ising model on TPU clusters,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis , ser. SC ’19. ACM, 2019, pp. 83:1–83:15. [Online]. Available: http://doi.acm.org/10.1145/3295500.3356149
2019
Later among the works it cites.
F. Belletti, D. King, K. Yang, R. Nelet, Y. Shafi, Y.-F. Chen, and J. Anderson, “Tensor processing units for financial Monte Carlo,” arXiv preprint arXiv:1906.02818 , 2019
Original
2019
Later among the works it cites.
G. Henry, P. T. P. Tang, and A. Heinecke, “Leveraging the bfloat16 artificial intelligence datatype for higher-precision computations,” arXiv preprint arXiv:1904.06376 , 2019
Original
2019
Later among the works it cites.

Beyond the bibliography

alphaXiv searches the wider corpus for related work and actual follow-ups.

Open on alphaXiv

alphaXiv is searching for related work…

Reading the bibliography…