Large-scale distributed linear algebra with tensor processing units
Abstract
We demonstrate that Google TPUs, novel computer chips optimized for artificial intelligence, can also efficiently perform scientific computation at massive scale. The concrete examples of matrix multiplication, QR factorization, linear solution, and application of matrix functions are benchmarked upon matrices of linear size up to N =1 ,048 ,576 , achieving throughput around 90% of the single-core value at the largest values of N considered. We believe TPUs and TPU-like architectures will soon play a role in scientific computation analogous to the GPU but at larger scale, and mean for this study to catalyze the process.
- Publication:
-
Proceedings of the National Academy of Science
- Pub Date:
- August 2022
- DOI:
- 10.1073/pnas.2122762119
- arXiv:
- arXiv:2112.09017
- Bibcode:
- 2022PNAS..11922762L
- Keywords:
-
- Physics - Computational Physics;
- Computer Science - Distributed;
- Parallel;
- and Cluster Computing;
- Quantum Physics
- E-Print:
- 12 pages, 8 figures