Automatic code generation for many-body electronic structure methods: the tensor contraction engine‡‡
Alexander A. Auer, Gerald Baumgartner, David E. Bernholdt, Alina Bibireata, Venkatesh Choppella, Daniel Cociorva, Xiaoyang Gao, Robert J. Harrison, Sriram Krishnamoorthy, Sandhya Krishnan, Chi‐Chung Lam, Qingda Lu, Marcel Nooijen, Russell M. Pitzer, Jagannathan Ramanujam, Ponnuswamy Sadayappan, Alexander Sibiryakov
As both electronic structure methods and the computers on which they are run become increasingly complex, the task of producing robust, reliable, high-performance implementations of methods at a rapid pace becomes increasingly daunting. In this paper we present an overview of the Tensor Contraction Engine (TCE), a unique effort to address issues of both productivity and performance through automatic code generation. The TCE is designed to take equations for many-body methods in a convenient high-level form and acts like an optimizing compiler, producing an implementation tuned to the target computer system and even to the specific chemical problem of interest. We provide examples to illustrate the TCE approach, including the ability to target different parallel programming models, and the effects of particular optimizations.
As both electronic structure methods and the computers on which they are run become increasingly complex, the task of producing robust, reliable, high-performance implementations of methods at a rapid pace becomes increasingly daunting. In this paper we present an overview of the Tensor Contraction Engine (TCE), a unique effort to address issues of both productivity and performance through automatic code generation. The TCE is designed to take equations for many-body methods in a convenient high-level form and acts like an optimizing compiler, producing an implementation tuned to the target computer system and even to the specific chemical problem of interest. We provide examples to illustrate the TCE approach, including the ability to target different parallel programming models, and the effects of particular optimizations.
We have developed a symbolic manipulation program and program generator (tensor contraction engine or TCE) that abstracts and automates the time-consuming, error-prone processes of deriving the working equations of a well-defined model of second-quantized many-electron theories and synthesizing efficient parallel computer programs on the basis of these equations. Provided an ansatz of a many-electron theory model, TCE performs valid contractions of creation and annihilation operators according to Wick's theorem, consolidates identical terms, and reduces the expressions into the form of multiple tensor contractions acted upon by permutation operators. It subsequently determines the binary contraction order for each multiple tensor contraction with the minimal operation and memory cost, factorizes common binary contractions (defines intermediate tensors), and identifies reusable intermediates. The resulting ordered list of binary tensor contractions, additions, and index permutations is translated into an optimized program that is combined with the NWCHEM and UTCHEM computational chemistry software packages. The programs synthesized by TCE take advantage of spin symmetry (within the spin−orbital formalisms), real Abelian point-group symmetry, and index permutation symmetry at every stage of the calculations to minimize the number of arithmetic operations and storage requirement, adjust the peak local memory usage by index-range tiling, and support parallel I/O interfaces and dynamic load balancing for parallel executions. We demonstrate the utility of TCE through automatic derivation and implementation of parallel programs for a range of predictive computational methods configuration-interaction theory (CISD, CISDT, CISDTQ), generalized many-body perturbation theory [MBPT(2), MBPT(3), MBPT(4)], and coupled-cluster theory (LCCD, CCD, LCCSD, CCSD, QCISD, CCSDT, and CCSDTQ), some for the first time and discuss the performance of the implemented programs.
The ultimate goal of electronic structure theory is solving the electronic Schr¨odinger Equation. However, even accurate approximations of solving Schr¨odinger Equation, such as high order coupled cluster theories, require computational efforts that are too demanding to be applied on large chemical systems. This thesis tackles the problem of curse of dimensionality: how to reduce the time complexity of high-accuracy coupled cluster methods in order to accelerate computations of molecular energy. On one hand, we believe that low-rank approximation (i.e. Tensor HyperContraction) of high-order tensors appearing in coupled cluster theory is a promising way to achieve rank-reduced coupled cluster theory. On the other hand, we think that the work of deriving, optimizing and implementing low-rank approximated coupled cluster theory is both tedious and repetitious, and shall be performed by an automatic code engine other than chemists. In this thesis, we introduce our chemistry equation compiler, Autom: a code engine that is designed for automatically deriving, optimizing and implementing low-rank approximated coupled cluster theories. With Autom, the development of electronic theories is greatly simplified. Autom can easily generate runnable codes in minutes for tensor-algebra based chemistry equations with customized low-rank approximation rules; conventionally, such procedure may take an expert weeks or even months to manually program. This thesis roughly consists of two parts. In the first part, we aim to cover the lowrank approximation techniques for high-accuracy coupled cluster theories, emphasizing on two new factorization techniques we developed especially for several coupled cluster theories including triple contributions. In the second part, we demonstrate the algorithms of our automatic code engine. We also summarize recent advances of domain specific compilers in Machine Learning community and discuss future directions for compute graph optimization. We believe that this new software infrastructure holds great promise for future chemistry package designs.
The continuous demand for higher accuracy in computational chemistry necessitates the development of advanced many-body electronic structure methods. However, the derivation and efficient implementation of these theories constitute a significant bottleneck. As the rank of the associated tensors increases, the governing equations explode in complexity, rendering manual implementation labor-intensive, error-prone, and difficult to optimize for modern hardware. To address this challenge, this dissertation presents SeQuant, a comprehensive framework for the automated derivation and parallel implementation of many-body quantum chemistry methods. Built upon a robust symbolic algebra engine, SeQuant allows for the expression of theories in the natural language of second quantization. It automates the transformation of high-level theoretical ansatzes into explicit tensor contraction expressions and subsequently generates optimized, high-performance C++ code. A central innovation of this work is the extension of automated implementation to reduced-scaling methods, which exploit the sparsity inherent in electronic correlation. We introduce a novel "tensor-of-tensors" data structure designed to manage the irregular sparsity patterns of Pair Natural Orbital (PNO) formulations. This development enables the first fully automated implementation of PNO-Coupled Cluster (PNO-CC) methods, bridging the gap between symbolic abstraction and the runtime requirements of sparse tensor algebra. The results demonstrate that SeQuant not only reproduces established dense methods (such as CCSD and CCSDT) with high fidelity but also effectively handles the complexity of sparse, local correlation approaches. By decoupling the complexity of the physics from the details of the implementation, this framework establishes a new paradigm for method development, dramatically accelerating the translation of theoretical insights into computational reality.
Computational many-body theories in quantum chemistry, condensed matter, and nuclear physics aim to provide sufficiently accurate description of and insights into the motion of many interacting particles. Due to their intrinsic complexity, the development of such theories generally involves very complex, tedious, and error-prone symbolic manipulations. Here a complete solution to automate the symbolics in many-body theory development is attempted. General data structures based on an existing computer algebra system are designed to specifically address the symbolic problems for which there is currently no satisfactory handling. Based on the data structures, algorithms are given to accomplish common symbolic manipulations and simplifications. Noncommutative algebraic systems, tensors with symmetry, and symbolic summations can all enjoy deep simplifications efficient enough for theories of very complex form. After the symbolic derivation, novel algorithms for automatic optimization of tensor contractions and their sums are devised, which can be used together with automatic code generation tools. In this way, the burden of symbolic tasks in theory development can be vastly reduced, with the potential to spare scientists more time and energy for the actual art and science of many-body theories.
Large-scale tensor network simulations are crucial for developing robust complexity-theoretic bounds on classical quantum simulation, enabling circuit cutting approaches, and optimizing circuit compilation, all of which aid efficient quantum computation on limited quantum resources. Modern exascale high-performance computing platforms offer significant potential for advancing tensor network quantum circuit simulation capabilities. We implement TN-Sim, a tensor network simulator backend within the NWQ-Sim software package that utilizes the Tensor Algebra for Many-body Methods (TAMM) framework to support both distributed HPC-scale computations and local simulations with ITensor. To optimize the scale up in computation across multiple nodes we implement a task based parallelization scheme to demonstrate parallelized gate contraction for wide quantum circuits with many gates per layer. Through the integration of the TAMM framework with Matrix Product State (MPS) tensor network approaches, we deliver a simulation environment that can scale from local systems to HPC clusters. We demonstrate an MPS tensor network simulator running on the state-of-the-art Perlmutter (NVIDIA) supercomputer and discuss the potential portability of this software to HPC clusters such as Frontier (AMD) and Aurora (Intel). We also discuss future improvements including support for different tensor network topologies and enhanced computational efficiency.
Calculation of high-rank tensor contractions plays a central role in computational physics, quantum chemistry, and computer science. The ability to perform a tensor contraction within a given computational budget poses a challenge to feasibility and restricts the types of systems that can be investigated computationally. We present the development and implementation of the Stratified Stochastic Tensor Contraction (SSTC) method as an alternative to sequential evaluation of tensor contraction. The SSTC approach partitions the tensor indices into stratified segments and employs composite index mappings to guide efficient sampling. By leveraging intrinsic structure within tensors, SSTC achieves rapid convergence to exact results with controllable sampling errors. We demonstrate the method on representative cases, two-point Coulomb integrals, four-point kernels, and correlation energy calculations using 2nd order Møller–Plesset perturbation theory, all of which are relevant to explicitly correlated electronic structure theory. In all examples, the SSTC method provides an accurate estimate of the tensor contraction with systematically reducible errors. We present the theory of exact tensor contractions as a foundation for the stochastic approach, followed by the mathematical derivation of the SSTC framework. Key components include the construction of a composite index to enable loop compression, and the application of functional minimization techniques for variance reduction. A discussion comparing sequential and stochastic contraction strategies, analysis of the distribution of sampling points through cumulative distribution functions and correlations between sampling assignments and kernel values are presented. A description of the SSTC method’s impact on computational clock times and dependence on segment number is included. Additional considerations include sampling with vs without replacement and the computational advantages of index mapping. These results highlight the effectiveness of the SSTC method as a versatile and scalable alternative to the conventional sequential tensor contraction approach, with broad applicability to quantum chemistry, many-body physics, stochastic quantum mechanics, and tensor-based machine learning.