Abstract:Tensor Processing Units (TPUs) are specialized hardware accelerators developed by Google to support large-scale machine-learning tasks, but they can also be leveraged to accelerate and scale other linear-algebra-intensive computations. In this paper we demonstrate the usage of TPUs for massively parallel, classical simulations of quantum many-body dynamics on long timescales. We apply our methods to study the phenomenon of Floquet prethermalization, i.e., exponentially slow heating in quantum spin chains subject to high-frequency periodic driving. We simulate the dynamics of L=34 qubits for over $10^5$ Floquet periods, corresponding to circuits with millions of two-qubit gates. The circuits simulated have no additional symmetries and represent a pure-state evolution in the full $2^L$-dimensional Hilbert space. This is achieved by distributing the computation over 128 TPU cores. On that size TPU cluster, we find speedups in wall-clock runtime of 230x and 15x when compared to reference CPU and single-GPU simulations, respectively, for shorter 30-qubit simulations that can be handled by all three platforms. We study the computational cost of the simulations, as a function of both the number of qubits and the number of TPU cores used, up to our maximum capacity of L=40 qubits, which requires a ``full pod" of 2048 TPU cores with tens of terabytes of memory in total. For these simulations, an 8-TPU-core machine is comparable to a single A100 GPU, and thus the full TPU pod is comparable to a machine with hundreds of GPUs. However, the TPU pod is more energy and cost efficient, and readily accessible (via Google Cloud), unlike such large many-GPU configurations. We also study the accumulation of numerical error as a function of circuit depth in very deep circuits. Our work demonstrates that TPUs can offer significant advantages for state-of-the-art simulations of quantum many-body dynamics.

Proposal for a High Precision Tensor Processing Unit

Implementation of High Performance Single-Precision Divider

High-resolution imaging on TPUs

A Carbon-Nanotube-based Tensor Processing Unit

GPTPU: Accelerating Applications using Edge Tensor Processing Units

Flex-TPU: A Flexible TPU with Runtime Reconfigurable Dataflow Architecture

Exploration of TPUs for AI Applications

A fine-grained mixed precision DNN accelerator using a two-stage big-little core RISC-V MCU.

Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution

Hardware-Efficient Mixed-Precision CP Tensor Decomposition

TPU as Cryptographic Accelerator

Mixed precision in Graphics Processing Unit

Heterogeneous Integration of In-Memory Analog Computing Architectures with Tensor Processing Units

TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

Simulation of quantum many-body dynamics with Tensor Processing Units: Floquet prethermalization

Query Processing on Tensor Computation Runtimes

Photonic tensor cores for machine learning

TATAA: Programmable Mixed-Precision Transformer Acceleration with a Transformable Arithmetic Architecture

Photonic tensor processing unit with single dataflow and programmable high-precision weighting control

High-Performance Tensor-Train Primitives Using GPU Tensor Cores

Accelerating MRI Reconstruction on TPUs