Benchmarking Quantum Red TEA on CPUs, GPUs, and TPUs

TL;DR

This paper benchmarks quantum many-body simulations across CPUs, GPUs, and TPUs, comparing different linear algebra libraries and optimization strategies to improve performance in tensor network algorithms.

Contribution

It introduces Quantum Red TEA, a library addressing tensor handling across hardware, and demonstrates significant speedups in quantum ground state searches.

Findings

Achieved a 34x speedup on CPU with parameter tuning.

Gained an additional 2.76x speedup on GPU over CPU.

Compared multiple linear algebra backends for quantum simulations.

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.