Benchmarking Quantum Red TEA on CPUs, GPUs, and TPUs

Abstract: We benchmark simulations of many-body quantum systems on heterogeneous hardware platforms using CPUs, GPUs, and TPUs. We compare different linear algebra backends, e.g., NumPy versus the PyTorch, JAX, or TensorFlow libraries, as well as a mixed-precision-inspired approach and optimizations for the target hardware. Quantum Red TEA out of the Quantum TEA library specifically addresses handling tensors with different libraries or hardware, where the tensors are the building blocks of tensor network algorithms. The benchmark problem is a variational search of a ground state in an interacting model. This is a ubiquitous problem in quantum many-body physics, which we solve using tensor network methods. This approximate state-of-the-art method compresses quantum correlations which is key to overcoming the exponential growth of the Hilbert space as a function of the number of particles. We present a way to obtain speedups of a factor of 34 when tuning parameters on the CPU, and an additional factor of 2.76 on top of the best CPU setup when migrating to GPUs.