Soft lauching Terco a `do concurrent` driven, GPU accelerated quantum chemistry library/package

JorgeG94 released terco, an open-source Fortran quantum chemistry library that uses `do concurrent` and the nvfortran 26.9 compiler to compute electron repulsion integrals on Nvidia GPUs, reporting roughly 10x speedup over a prior Fortran standard-parallelism implementation in GAMESS and performance compared against the CUDA C++-based gpu4pyscf. The package collapses the recursive ERI evaluation into fully unrolled, auto-generated routines that can run about 20,000 lines long, with the Boys function base case solved via Chebyshev interpolation in 14 FLOPs, and all GPU optimizations such as shared memory use and launch bounds left to the compiler. The author said the next steps are compiling terco with LFortran and targeting Apple, AMD and Intel GPUs.

If you speak Spanish you know that “terco” means “stubborn”. This package was born out of my delirious dreams of "I am sure we can get close to cuda performance if we use do concurrent and abuse the nvidia Fortran compilers. This is what terco https://github.com/JorgeG94/terco does I’ve already stolen the “powered by Fortran” image, thanks @fxm https://fortran-lang.discourse.group/u/fxm In Molecular Quantum Chemistry everything goes down into evaluating these things called “electron repulsion integrals” or ERIs. They are BLOODY AWFUL. There are O N^4 of them and they are needed for almost everything. They are 6D integrals over a Coulomb operator and they get more difficult with increasing angular momentum, remember your s , p , d from high school chemistry. These integrals have basis functions which have x, y, and z coordinates. They are of the form \mu \nu | \lambda \sigma and they are solved recursively, with the base case of the recursion being usually the evaluation of an integral called the Boys function, which has been shown that the best way to solve it is via Chebyshev interpolation since you can do it in 14 FLOPs. On the GPU, the best way to do this is to collapse these recursions into full unrolled loops, therefore, certain functions can be around ~20k lines long The best performing codes out in the literature all use CUDA C++, some efforts were done previously exploiting OpenMP and Do Concurrent from native Fortran: However, none of these implementations were able to get comparable performance to CUDA written applications at scale - and these papers admit that. So, is CUDA like performance unreachable to standard Fortran for qc applications? No. The main issue facing previous implementations was feeding the GPU with enough work By using some batching schemes and shenanigans, I am able to beat the implementation found in here https://pubs.acs.org/jctcce/article-abstract/19/13/3798/400989/Enabling-Fortran-Standard-Parallelism-in-GAMESS?redirectedFrom=fulltext by a former colleague of mine by around 10x wooo Attached is a figure of the comparison of terco using do concurrent against gpu4pyscf which is a state of the art GPU accelerated code written in CUDA C++ and python as orchestrator. All optimizations done in the Fortran side are algorithmic - shared memory utilization, launch bounds, etc. are all the compiler. Using nvfortran 26.9. Next step is to cause @certik https://fortran-lang.discourse.group/u/certik nightmares and get this code compiled with LFortran, maybe use some Apple GPUs Then getting this to run on AMD and Intel. Nightmares ahead The code is heavily validated but bugs might be present if you do some weird chemistry. Example: do concurrent i=1:nr call pcsi1010 g0 + i - 1 nranks, lo, hi, nseg, sOff, sA, sNB, sOA, sOB, sPA, sPB, sD, & npair, sp i, sp j, sp q, thresh, pcut, jfac, kfac, omega, dsh, nbas, nao, sh l, ao off, & nqp, qp z, qp c, qp px, qp py, qp pz, & sq ax, sq ay, sq az, sq abx, sq aby, sq abz, sq n, ndens, dmat, jmat end do and the function is: the code is auto-generated via a python script pure subroutine pcsi1010 gt, lo, hi, nseg, sOff, sA, sNB, sOA, sOB, sPA, sPB, sD, & npair, sp i, sp j, sp q, thresh, pcut, jfac, kfac, omega, dsh, nbas, nao, sh l, ao off, & nqp, qp z, qp c, qp px, qp py, qp pz, & sq ax, sq ay, sq az, sq abx, sq aby, sq abz, sq n, & ndens, dmat, jmat $acc routine seq integer int64 , intent in :: gt integer, intent in :: lo, hi, nseg, npair, nbas, nao integer int64 , intent in :: sOff nseg + 1 integer, intent in :: sA nseg , sNB nseg , sOA nseg , sOB nseg integer, intent in :: sPA nseg , sPB nseg logical, intent in :: sD nseg integer, intent in :: sp i npair , sp j npair real dp , intent in :: sp q npair , thresh real dp , intent in :: pcut real dp , intent in :: jfac, kfac real dp , intent in :: omega real dp , intent in :: dsh nbas, nbas integer, intent in :: sh l nbas , ao off nbas integer, intent in :: nqp real dp , intent in :: qp z nqp , qp c nqp real dp , intent in :: qp px nqp , qp py nqp , qp pz nqp real dp , intent in :: sq ax npair , sq ay npair , sq az npair real dp , intent in :: sq abx npair , sq aby npair , sq abz npair integer, intent in :: sq n npair integer, intent in :: ndens real dp , intent in :: dmat ndens, nao, nao real dp , intent inout :: jmat ndens, nao, nao integer :: p, q, mid, seg, t, iab, icd, si, sj, sk, sl integer int64 :: nsa, u, kx real dp :: qcut, bnd integer :: slotab, slotcd, nab, ncd integer :: kp, kq, k1, k2 integer :: d, x, cur, ia, ib, ic, id, idx, idens integer :: mu, nu, lam, sig, mui, nuj, lamk, sigl logical :: dij, dkl, dpq real dp :: zeta, eta, zpe, rho, tval, pref, wc real dp :: theta, rho theta real dp :: pqx, pqy, pqz, pax, pay, paz, qcx, qcy, qcz real dp :: wpx, wpy, wpz, wqx, wqy, wqz real dp :: oo2z, oo2e, oo2ze, rz, re, sc, vv real dp :: abx, aby, abz, cdx, cdy, cdz real dp :: rpx, rpy, rpz, rzeta, reta, rzpe real dp :: f 0:BOYS MMAX integer :: bi, bj, bbase real dp :: bx, bx2, b0, b1, b2, btt, bet real dp :: v 16, 0:1 , g1 16 , vbuf 9 real dp :: wq, w, wab, cabmax, ccdmax, qchunk logical :: same ab, same cd, same pair real dp :: jab 3 , jcd 3 , kac 9 real dp :: kad 3 , kbc 3 , kbd 1 real dp :: dab 3 , dcd 3 , dac 9 real dp :: dad 3 , dbc 3 , dbd 1 locate the segment p = lo; q = hi do while p < q mid = p + q + 1 /2 if sOff mid < gt then p = mid else q = mid - 1 end if end do seg = p t = int gt - sOff seg if sD seg then nsa = int sA seg , int64 u = int t - 1, int64 kx = int real 2 nsa + 1, dp - sqrt real 2 nsa + 1, dp 2 - 8.0 dp real u, dp /2.0 dp, int64 if kx < 0 kx = 0 do while kx 0 if kx 2 nsa + 1 - kx kx /2 <= u exit kx = kx - 1 end do do while kx + 1 2 nsa + 1 - kx + 1 kx + 1 /2 <= u kx = kx + 1 end do icd = int kx + 1 iab = icd + int u - kx 2 nsa + 1 - kx kx /2 else icd = t - 1 /sA seg + 1 iab = t - icd - 1 sA seg end if qcut = sp q sOA seg + iab sp q sOB seg + icd if qcut <= thresh return si = sp i sOA seg + iab ; sj = sp j sOA seg + iab sk = sp i sOB seg + icd ; sl = sp j sOB seg + icd if qcut max 4.0 dp dsh si, sj , 4.0 dp dsh sk, sl , & dsh si, sk , dsh si, sl , & dsh sj, sk , dsh sj, sl <= thresh return same ab = si == sj ; same cd = sk == sl same pair = sD seg .and. iab == icd slotab = sOA seg + iab; slotcd = sOB seg + icd nab = sq n slotab ; ncd = sq n slotcd abx = sq abx slotab ; aby = sq aby slotab ; abz = sq abz slotab cdx = sq abx slotcd ; cdy = sq aby slotcd ; cdz = sq abz slotcd do x = 1, 16 g1 x = 0.0 dp end do do k1 = 1, nab kp = sPA seg + k1 - 1 sA seg + iab zeta = qp z kp bnd = TWO PI 2 5 abs qp c kp / zeta sqrt zeta if bnd abs qp c sPB seg + icd /qp z sPB seg + icd <= pcut cycle rpx = qp px kp ; rpy = qp py kp ; rpz = qp pz kp pax = rpx - sq ax slotab pay = rpy - sq ay slotab paz = rpz - sq az slotab rzeta = 1.0 dp/zeta oo2z = 0.5 dp rzeta do k2 = 1, ncd kq = sPB seg + k2 - 1 sNB seg + icd eta = qp z kq zpe = zeta + eta pref = TWO PI 2 5/ zeta eta sqrt zpe qp c kp qp c kq if bnd abs qp c kq /eta <= pcut exit if abs pref <= pcut cycle reta = 1.0 dp/eta rzpe = 1.0 dp/zpe rho = zeta eta rzpe if omega 0.0 dp then theta = omega omega/ rho + omega omega else theta = 1.0 dp end if rho theta = theta rho pqx = rpx - qp px kq pqy = rpy - qp py kq pqz = rpz - qp pz kq qcx = qp px kq - sq ax slotcd qcy = qp py kq - sq ay slotcd qcz = qp pz kq - sq az slotcd wc = zeta rpx + eta qp px kq rzpe wpx = wc - qp px kp ; wqx = wc - qp px kq wc = zeta rpy + eta qp py kq rzpe wpy = wc - qp py kp ; wqy = wc - qp py kq wc = zeta rpz + eta qp pz kq rzpe wpz = wc - qp pz kp ; wqz = wc - qp pz kq wpx = theta wpx; wpy = theta wpy; wpz = theta wpz wqx = theta wqx; wqy = theta wqy; wqz = theta wqz tval = rho theta pqx pqx + pqy pqy + pqz pqz if tval = BOYS TMAX then btt = 1.0 dp/tval f 0 = 0.88622692545275801365 dp sqrt btt bet = 0.0 dp f 1 = 1.0 dp f 0 0.5 dp btt f 2 = 3.0 dp f 1 0.5 dp btt else bi = int tval BOYS DTINV bi = min bi, BOYS NGRID - 1 bx = 2.0 dp tval - real bi, dp BOYS DT BOYS DTINV - 1.0 dp bx2 = 2.0 dp bx bbase = bi BOYS MMAX + 1 BOYS NCHEB + 1 bj = bbase + 2 BOYS NCHEB + 1 b1 = boys table bj + BOYS NCHEB + 1 b2 = 0.0 dp b0 = bx2 b1 - b2 + boys table bj + 5 ; b2 = b1; b1 = b0 b0 = bx2 b1 - b2 + boys table bj + 4 ; b2 = b1; b1 = b0 b0 = bx2 b1 - b2 + boys table bj + 3 ; b2 = b1; b1 = b0 b0 = bx2 b1 - b2 + boys table bj + 2 ; b2 = b1; b1 = b0 f 2 = bx b1 - b2 + boys table bj + 1 bet = exp -tval btt = 2.0 dp tval f 1 = btt f 2 + bet 0.3333333333333333 dp f 0 = btt f 1 + bet 1.0 dp end if pref = pref sqrt theta oo2e = 0.5 dp reta; oo2ze = theta 0.5 dp rzpe rz = rho theta rzeta; re = rho theta reta --- level m = 2 --- v 1,0 = pref f 2 --- level m = 1 --- v 1,1 = pref f 1 v 2,1 = pax v 1,1 + wpx v 1,0 v 3,1 = pay v 1,1 + wpy v 1,0 v 4,1 = paz v 1,1 + wpz v 1,0 --- level m = 0 --- v 1,0 = pref f 0 v 2,0 = pax v 1,0 + wpx v 1,1 v 3,0 = pay v 1,0 + wpy v 1,1 v 4,0 = paz v 1,0 + wpz v 1,1 v 6,0 = qcx v 2,0 + wqx v 2,1 & + 1.0 dp oo2ze v 1,1 v 7,0 = qcx v 3,0 + wqx v 3,1 v 8,0 = qcx v 4,0 + wqx v 4,1 v 10,0 = qcy v 2,0 + wqy v 2,1 v 11,0 = qcy v 3,0 + wqy v 3,1 & + 1.0 dp oo2ze v 1,1 v 12,0 = qcy v 4,0 + wqy v 4,1 v 14,0 = qcz v 2,0 + wqz v 2,1 v 15,0 = qcz v 3,0 + wqz v 3,1 v 16,0 = qcz v 4,0 + wqz v 4,1 & + 1.0 dp oo2ze v 1,1 cur = 0 do x = 1, 16 g1 x = g1 x + v x, cur end do end do end do Single column on every side: the canonical enumeration and the degeneracy weights are the pair-level ones. dij = .not. same ab dkl = .not. same cd dpq = .not. same pair mui = ao off si ; nuj = ao off sj ; lamk = ao off sk ; sigl = ao off sl --- HRR --- vbuf 1 = g1 6 vbuf 2 = g1 7 vbuf 3 = g1 8 vbuf 4 = g1 10 vbuf 5 = g1 11 vbuf 6 = g1 12 vbuf 7 = g1 14 vbuf 8 = g1 15 vbuf 9 = g1 16 wq = 1.0 dp if .not. dij wq = wq 0.5 dp if .not. dkl wq = wq 0.5 dp if .not. dpq wq = wq 0.5 dp if kfac == 0.0 dp then do idens = 1, ndens ib = 0 do ia = 0, 2 dab 1 + ia + 3 ib = dmat idens, mui + ia, nuj + ib jab 1 + ia + 3 ib = 0.0 dp end do id = 0 do ic = 0, 2 dcd 1 + ic + 3 id = dmat idens, lamk + ic, sigl + id jcd 1 + ic + 3 id = 0.0 dp end do id = 0 do ic = 0, 2 ib = 0 do ia = 0, 2 jab 1 + ia + 3 ib = jab 1 + ia + 3 ib & + 4.0 dp jfac wq vbuf 1 + ia + 3 ib + 3 ic + 9 id & dcd 1 + ic + 3 id jcd 1 + ic + 3 id = jcd 1 + ic + 3 id & + 4.0 dp jfac wq vbuf 1 + ia + 3 ib + 3 ic + 9 id & dab 1 + ia + 3 ib end do end do ib = 0 do ia = 0, 2 $acc atomic update jmat idens, mui + ia, nuj + ib = jmat idens, mui + ia, nuj + ib & + jab 1 + ia + 3 ib end do id = 0 do ic = 0, 2 $acc atomic update jmat idens, lamk + ic, sigl + id = jmat idens, lamk + ic, sigl + id & + jcd 1 + ic + 3 id end do end do idens else do idens = 1, ndens do ib = 0, 0 do ia = 0, 2 dab 1 + ia + 3 ib = dmat idens, mui + ia, nuj + ib jab 1 + ia + 3 ib = 0.0 dp end do end do do id = 0, 0 do ic = 0, 2 dcd 1 + ic + 3 id = dmat idens, lamk + ic, sigl + id jcd 1 + ic + 3 id = 0.0 dp end do end do do ic = 0, 2 do ia = 0, 2 dac 1 + ia + 3 ic = dmat idens, mui + ia, lamk + ic kac 1 + ia + 3 ic = 0.0 dp end do end do do id = 0, 0 do ia = 0, 2 dad 1 + ia + 3 id = dmat idens, mui + ia, sigl + id kad 1 + ia + 3 id = 0.0 dp end do end do do ic = 0, 2 do ib = 0, 0 dbc 1 + ib + 1 ic = dmat idens, nuj + ib, lamk + ic kbc 1 + ib + 1 ic = 0.0 dp end do end do do id = 0, 0 do ib = 0, 0 dbd 1 + ib + 1 id = dmat idens, nuj + ib, sigl + id kbd 1 + ib + 1 id = 0.0 dp end do end do idx = 0 do id = 0, 0 do ic = 0, 2 do ib = 0, 0 do ia = 0, 2 idx = idx + 1 sc = wq vbuf idx jab 1 + ia + 3 ib = jab 1 + ia + 3 ib & + 4.0 dp jfac sc dcd 1 + ic + 3 id jcd 1 + ic + 3 id = jcd 1 + ic + 3 id & + 4.0 dp jfac sc dab 1 + ia + 3 ib kac 1 + ia + 3 ic = kac 1 + ia + 3 ic & - kfac sc dbd 1 + ib + 1 id kad 1 + ia + 3 id = kad 1 + ia + 3 id & - kfac sc dbc 1 + ib + 1 ic kbc 1 + ib + 1 ic = kbc 1 + ib + 1 ic & - kfac sc dad 1 + ia + 3 id kbd 1 + ib + 1 id = kbd 1 + ib + 1 id & - kfac sc dac 1 + ia + 3 ic end do end do end do end do do ib = 0, 0 do ia = 0, 2 $acc atomic update jmat idens, mui + ia, nuj + ib = jmat idens, mui + ia, nuj + ib & + jab 1 + ia + 3 ib end do end do do id = 0, 0 do ic = 0, 2 $acc atomic update jmat idens, lamk + ic, sigl + id = jmat idens, lamk + ic, sigl + id & + jcd 1 + ic + 3 id end do end do do ic = 0, 2 do ia = 0, 2 $acc atomic update jmat idens, mui + ia, lamk + ic = jmat idens, mui + ia, lamk + ic & + kac 1 + ia + 3 ic end do end do do id = 0, 0 do ia = 0, 2 $acc atomic update jmat idens, mui + ia, sigl + id = jmat idens, mui + ia, sigl + id & + kad 1 + ia + 3 id end do end do do ic = 0, 2 do ib = 0, 0 $acc atomic update jmat idens, nuj + ib, lamk + ic = jmat idens, nuj + ib, lamk + ic & + kbc 1 + ib + 1 ic end do end do do id = 0, 0 do ib = 0, 0 $acc atomic update jmat idens, nuj + ib, sigl + id = jmat idens, nuj + ib, sigl + id & + kbd 1 + ib + 1 id end do end do end do idens end if end subroutine pcsi1010 My next post will be about the ocean, very different