Quantum Chemistry and First Principles Molecular Dynamics on GPUs Todd J. Martinez Dept of Chemistry and PULSE Institute, Stanford University SLAC National Accelerator Laboratory TeraChem Team: Ivan Ufimtsev (Core SCF and Gradient Routines) Alexey Titov (meta-programming strategies) Nathan Luehr (DFT Quadrature and Dynamic Precision) Jeffrey Gour (TDDFT Gradients and Nonadiabatic Couplings) Christine Isborn (TDDFT and CIS) Heather Kulik (Ab Initio Protein Structure Optimization) Funding: NSF, AFOSR-STTR, DoD, DOE

Quantum Chemistry Where are the electrons? Molecular Dynamics Where are the atoms? or Potential Energy Surface Two Pillars of Chemical Simulation

Hartree-Fock Approximation Ψ electronic is antisymmetrized product of N 1-electron orbitals ψ Expand ψ over predefined basis set φ Often taken to be atom-centered Gaussian functions in molecular (non-periodic) case Density Functional Theory – Include K xc [ρ(ψ)] in mean-field Hamiltonian Adds need for numerical quadrature but this is one of the most straightforward tasks to implement efficiently on GPU (1-2 days) and it is not a major bottleneck

Hartree-Fock equations All matrices are of size NxN (N ~ 1,000 … 10,000) N 3 operations to solve HF equations (diagonalization of F) N 4 operations to get F (formally) Fock matrix construction is usually the bottleneck

What are GPUs good for? “Stream processing” Massively parallel problems with millions threads running concurrently (more than 10,000 threads can run in parallel) Simple kernels – to launch as many threads as possible to hide memory latency High computational load (big FLOPS/MOPS ratio) Little or no communication between threads, except for small blocks of threads (up to 512-thread blocks) that can share data very efficiently

6 Every atomic orbital is a fixed contraction of Gaussians Molecular orbitals are orthogonal contractions of AOs Antisymmetrized products of MOs Total electronic wfn is contraction of APs Large amount of data management and linear algebra Quantum chemistry is “stream-friendly!” Large amount of data management and linear algebra Quantum chemistry is “stream-friendly!” Quantum Chemistry - The Never-Ending Contraction

Computation of ERIs on GPU CPUGPU “Pair quantities” for bra and ket – N 2 [bra|ket] integrals – N 4 1 thread for each primitive integral!

Arrangement of Two-Electron Integrals Different code for each integral type, e.g. [ss|ss] vs [sp|ss] Rearrange computation to optimize kernel execution

Further Arrangement of Integrals leaves only N 2 out of N 4 integrals SIMD warp Most negligibly small integrals will be calculated SIMD warp Only significant integrals will be calculated [ij| |kl] ssspppsssppp x x x10 -9 ss sp pp ss sp pp Note: “Pre”-Computation to put problem in form suitable for GPU – this is essential and always (in our experience) the key to good performance Likely also better for modern CPUs and multi-core (but CPUs are more forgiving)

J-matrix implementation Next critical step is ordering threads into blocks and specifying how to traverse the problem

How Well Does this Work in Practice? RHF/3-21G

Riding the Wave of Faster GPUs… GT200  Fermi No code modification

TeraChem / GAMESS Benchmarks The speedups are bigger for larger molecules: GAMESS, 32 Intel Clovertown cores: s TeraChem, 4 GeForce 295 GTX GPUs: 407 s 127X GPU vs CPU core speedup 1 TeraStation ($10K) = 128 Dual Quad-Core Intel Nodes ($512K-$1M) ≈ BPTI Protein

Numerical Precision on GPU Take advantage of faster single precision arithmetic Can never use only single precision, but can use mixed precision Diagonalize in double precision Accumulate contributions to Fock matrix in double precision (MP1) Use DP for largest integrals, SP for smallest (MP2)

Ascorbic AcidLactoseCyano Toxin GAMESS TeraChem DP TeraChem MP Neurokinin A5x6 Nanotube GAMESS TeraChem DP TeraChem MP Numerical Precision on GPU

Beyond Mixed Precision Double PrecisionDynamic Precision Precision Error Conv Thres Final EnergyIterFinal EnergyIter Ascorbic Acid (1000K) E E Lactose (Minimum) E E Cyano Toxin (2000K) E E Neurokinin A (Minimum) E E Nanotube (Minimum) E E Nanotube (2000K) E E Crambin E E Ubiquitin E E T-Cadherin EC E E Ribonuclease A E E Change precision dynamically… 2-4x speedup over pure double precision (2x for GTX480, 4x for GTX295)

Summary Some Key issues: Compute upfront on the CPU to optimize computation on GPU Need many lightweight threads – look for the data which scales worst in order to determine parallelization direction Often want to waste compute cycles to maximize throughput e.g., do not take advantage of integral index symmetries (ij|kl)=(kl|ij) in J-matrix build Consider writing codes (or using a combination of Maple/Python/C) to write codes – forerunner to an ATLAS-like strategy for more complex tasks than linear algebra Many of these ideas are also useful on CPUs – the primary difference is that bad GPU code is 100x slower than good GPU code while bad CPU code is 5x slower than good CPU code…