Optimizing Tensor Contractions in CCSD(T) for Efficient Execution on GPUs

Abstract

Tensor contractions are higher dimensional analogs of matrix multiplications, used in many computational contexts such as high order models in quantum chemistry, deep learning, finite element methods etc. In contrast to the wide availability of high-performance libraries for matrix multiplication on GPUs, the same is not true for tensor contractions. In this paper, we address the optimization of a set of symmetrized tensor contractions that form the computational bottleneck in the CCSD(T) coupled-cluster method in computational chemistry suites like NWChem. Some of the challenges in optimizing tensor contractions that arise in practice from the variety of dimensionalities and shapes for tensors include effective mapping of the high-dimensional iteration space to threads, choice of data buffering in shared-memory and registers, and tile sizes for multi-level tiling. Furthermore, in the case of symmetrized tensor contractions in CCSD(T), it is also a challenge to fuse contractions to reduce data movement cost by exploiting reuse of intermediate tensors. In this paper, we develop an efficient GPU implementation of the tensor contractions in CCSD(T) using shared-memory buffering, register tiling, loop fusion and register transpose. Experimental results demonstrate significant improvement over the current state-of-the-art.

Fast Kronecker Matrix-Matrix Multiplication on GPUs

“…al. [24] improves tensor contractions for coupled cluster methods in quantum chemistry by fusing multiple contractions. However, their approach performs transpose in shared memory and these tensor contractions are different from contractions in Kron-Matmul.…”

From NWChem to NWChemEx: Evolving with the Computational Chemistry Landscape

“…For example, the current V100 GPUs group 32 cores into a warp, and there are 160 warps to a GPU 54,55 . Various existing computational chemistry codes have been adapted 53,[56][57][58][59][60][61][62][63][64][65][66][67] or designed from the outset (e.g., TeraChem) 67,68 to use GPUs.…”

“…High throughput is achieved by operating many warps simultaneously. For example, the current V100 GPUs group 32 cores into a warp, and there are 160 warps to a GPU. , Various existing computational chemistry codes have been adapted ,− or designed from the outset (e.g., TeraChem) , to use GPUs.…”

Athena

“…Sparse tensor contraction. Dense tensor contraction has been studied for decades on diverse hardware platforms [5,19,21,27,28,32,34,42,50,65,72,73], in scientific computing including chemistry, physics, and mechanics. The state-of-the-art studies focus on block-sparse tensor contractions with dense blocks in tensors.…”