diff --git a/Cargo.lock b/Cargo.lock index a225da2..047aee4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -420,6 +420,7 @@ version = "0.1.0" dependencies = [ "faer", "libc", + "rayon", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 1b9cf2f..bc5473c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,9 @@ edition = "2021" [dependencies] libc = "0.2.155" faer = "0.19.0" +# already in the tree as a faer dependency; used directly by src/par.rs for +# work-stealing over shell pairs +rayon = "1.11" [profile.dev] #debug = true diff --git a/pyproject.toml b/pyproject.toml index cbb7e0c..ae3f15e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,3 +35,10 @@ librint = ["*.so", "*.dylib", "test_sanity.py"] [tool.pytest.ini_options] testpaths = ["python/tests"] +markers = [ + # Basis sets big enough to cost minutes (def2-tzvp and up). They exist to + # reach code paths the small systems never touch -- f/g shells, general + # contraction -- so skip them for a quick loop, not because they are + # redundant. Deselect with: pytest -m "not slow" + "slow: minutes per case; exercises high-l and general-contraction paths", +] diff --git a/python/librint/_bindings.py b/python/librint/_bindings.py index f506b49..30beb1b 100644 --- a/python/librint/_bindings.py +++ b/python/librint/_bindings.py @@ -170,6 +170,35 @@ def _dylib_suffix(): ) library.denergy_c.restype = ctypes.POINTER(ctypes.c_double) +# Threaded counterparts (src/par.rs). Same arguments as above plus a trailing +# thread count; 0 means rayon's global pool, sized by RAYON_NUM_THREADS. +_PAR_ARGS = ( + ctypes.POINTER(ctypes.c_int), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_int), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_double), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_double), + ctypes.c_size_t, + ctypes.c_size_t, +) + +# Bound defensively: any .so built before src/par.rs existed -- including the +# one committed in this directory -- has none of these symbols, and reaching +# for them eagerly would make `import librint` fail for everyone rather than +# just for the caller who wants a threaded gradient. HAS_PAR lets dscf.py raise +# something a human can act on, and lets the test suite skip instead of error. +HAS_PAR = True +for _name in ("dS_par_c", "dR_par_c", "dHcore_par_c", "danalytical_par_c"): + try: + _fn = getattr(library, _name) + except AttributeError: + HAS_PAR = False + break + _fn.argtypes = _PAR_ARGS + _fn.restype = ctypes.POINTER(ctypes.c_double) + # Releases any buffer returned by the entry points above; len is the element # count that call produced. utils.take() copies then calls this. library.free_c.argtypes = ( diff --git a/python/librint/dscf.py b/python/librint/dscf.py index 0e7e106..86af323 100644 --- a/python/librint/dscf.py +++ b/python/librint/dscf.py @@ -1,6 +1,7 @@ import ctypes import numpy as np +from librint import _bindings from librint import library from librint import utils @@ -105,6 +106,59 @@ def danalyticalf(mol, P: np.ndarray) -> np.ndarray: dR_c = library.danalytical_c(atm_ctypes, len(atm.flatten()), bas_ctypes, len(bas.flatten()), env_ctypes, len(env.flatten()), P_ctypes, len(P.flatten())) return utils.take(dR_c, (s2 - s1,)) +# --------------------------------------------------------------------------- +# Threaded entry points (src/par.rs). +# +# These are separate callables, not a flag on the serial ones: danalyticalf +# stays the finite-difference-validated reference path, byte-for-byte, so +# "parallel == serial" remains a statement about two independent things. +# --------------------------------------------------------------------------- + +def _par(fn, mol, W: np.ndarray, nthreads: int) -> np.ndarray: + if not _bindings.HAS_PAR: + raise RuntimeError( + "this librint.so has no threaded entry points -- it predates " + "src/par.rs. Rebuild (cargo build --release) and point LIBRINT_SO " + "at target/release/librint.so, or use the serial danalyticalf." + ) + atm, bas, env, nelec = utils.prep(mol) + W = np.ascontiguousarray(W, dtype=np.float64) + s1, s2 = utils.split(bas) + + ptr = fn( + atm.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), atm.size, + bas.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), bas.size, + env.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), env.size, + W.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), W.size, + int(nthreads), + ) + return utils.take(ptr, (s2 - s1,)) + + +def dS_par(mol, Q: np.ndarray, nthreads: int = 0) -> np.ndarray: + """Overlap term seeded with the energy-weighted density Q = P F P. + + Unlike dSf this does NOT build F -- pass Q, not P. + """ + return _par(library.dS_par_c, mol, Q, nthreads) + + +def dHcore_par(mol, P: np.ndarray, nthreads: int = 0) -> np.ndarray: + return _par(library.dHcore_par_c, mol, P, nthreads) + + +def dR_par(mol, P: np.ndarray, nthreads: int = 0) -> np.ndarray: + _require_grad_domain(mol) + return _par(library.dR_par_c, mol, P, nthreads) + + +def danalytical_par(mol, P: np.ndarray, nthreads: int = 0) -> np.ndarray: + """Threaded danalyticalf. nthreads=0 uses rayon's global pool, which reads + RAYON_NUM_THREADS; any other value builds a pool of exactly that size.""" + _require_grad_domain(mol) + return _par(library.danalytical_par_c, mol, P, nthreads) + + def denergyf(mol, P: np.ndarray) -> np.ndarray: _require_grad_domain(mol) atm, bas, env, nelec = utils.prep(mol) diff --git a/python/tests/test_gradient_fd.py b/python/tests/test_gradient_fd.py index 8aceec8..948b99e 100644 --- a/python/tests/test_gradient_fd.py +++ b/python/tests/test_gradient_fd.py @@ -68,3 +68,31 @@ def test_gradient_consistency(basis, geo): # Validate against jax benchmark np.testing.assert_allclose(grad_analytical_sorted, grad_fd_sorted, atol=1e-5, rtol=1e-4) + + +@pytest.mark.skipif( + not librint._bindings.HAS_PAR, + reason="librint.so has no threaded entry points; rebuild and set LIBRINT_SO", +) +@pytest.mark.parametrize("basis, geo", MOLECULES) +def test_gradient_consistency_threaded(basis, geo): + """The same finite-difference check, through the threaded path. + + test_par_equiv.py already ties danalytical_par to danalyticalf, and the + test above ties danalyticalf to finite differences, so this is transitively + covered. It is here anyway because the transitive argument breaks silently + if either link is ever weakened, and this one is direct. + """ + molecule = geometries[geo] + atom = '\n'.join([f"{a[0]} {0.529*a[2][0]} {0.529*a[2][1]} {0.529*a[2][2]}" for a in molecule]) + + mol_rpyscf = pyscf.gto.M(atom=atom, basis=basis) + P = librint.scf.density(mol_rpyscf, imax=MAX_ITER) + + grad_fd = calc_fd(mol_rpyscf) + # 0 = rayon's global pool, sized by RAYON_NUM_THREADS; whatever the machine + # running the suite happens to have is a fine width for a correctness check + grad_par = librint.dscf.danalytical_par(mol_rpyscf, P, 0) + + np.testing.assert_allclose(np.sort(grad_par), np.sort(grad_fd), + atol=1e-5, rtol=1e-4) diff --git a/python/tests/test_par_equiv.py b/python/tests/test_par_equiv.py new file mode 100644 index 0000000..3683e83 --- /dev/null +++ b/python/tests/test_par_equiv.py @@ -0,0 +1,160 @@ +"""Does the threaded gradient compute the same thing as the serial one? + +The serial path (danalyticalf) is the reference: test_gradient_fd.py validates +it against central finite differences and test_gradient_pyscfad.py against +pyscfad. The threaded path (danalytical_par) has neither of those on its own, +so it has to be tied to the serial one before any timing of it means anything. + +Exact equality is NOT the criterion and would be the wrong thing to demand: a +work-stealing reduction sums the same terms in a different association order, +so the two agree to round-off, not to the bit. What this asserts is + + 1. every term matches its serial counterpart (dHcore, dR, dS separately, so + a failure localizes), and the assembled gradient matches too; + 2. the agreement does not degrade as threads are added -- a real race shows + up as error growing with thread count, whereas reassociation noise stays + flat. + +The thread sweep is clamped to the cores actually available, so this shrinks to +something meaningful on a laptop instead of failing there. +""" +import os + +import numpy as np +import pytest +import pyscf + +import librint +import librint.dscf +import librint.utils +from librint import _bindings + +from pyscf_comp.geometries import geometries + +# The .so committed in python/librint/ predates src/par.rs. Skipping beats +# failing: nothing here is broken, the library just has no threaded path to +# compare against. Point LIBRINT_SO at a fresh target/release/librint.so. +pytestmark = pytest.mark.skipif( + not _bindings.HAS_PAR, + reason="librint.so has no threaded entry points; rebuild and set LIBRINT_SO", +) + +# Small enough to run every time. sto-3g and def2-svp are s/p only. +FAST = [ + ("H2", "sto-3g"), + ("H2O", "sto-3g"), + ("NH3", "sto-3g"), + ("CH4", "sto-3g"), + ("H2O", "def2-svp"), + ("NH3", "def2-svp"), + ("CH4", "def2-svp"), +] + +# Each of these exists to reach a code path the fast list never touches, so +# they are worth minutes when you want them -- and worth skipping when you do +# not. Run with `-m slow`, or everything with no -m at all. +SLOW = [ + ("H2O", "def2-tzvp"), # f shells (l=3): rys_tab.rs nroots 6-7 Chebyshev + ("NH3", "def2-tzvp"), + ("CH4", "cc-pvdz"), # general contraction (nctr>1): eri_cart_gc path + ("H2O", "cc-pvtz"), + ("H2O", "def2-qzvp"), # g shells (l=4): rys_tab.rs nroots 8-9 Chebyshev +] + +SYSTEMS = ([pytest.param(g, b) for g, b in FAST] + + [pytest.param(g, b, marks=pytest.mark.slow) for g, b in SLOW]) + +THREADS = [1, 2, 4, 8, 16, 32, 64] +RTOL = 1e-9 # relative to max|serial| + + +def build(geo, basis): + atom = "\n".join( + f"{a[0]} {0.529 * a[2][0]} {0.529 * a[2][1]} {0.529 * a[2][2]}" + for a in geometries[geo] + ) + mol = pyscf.gto.M(atom=atom, basis=basis, verbose=0) + mol.cart = True + return mol + + +def rel(got, ref): + scale = max(float(np.abs(ref).max()), 1e-30) + return float(np.abs(got - ref).max()) / scale + + +def thread_counts(): + ncores = len(os.sched_getaffinity(0)) + return [t for t in THREADS if t <= ncores] + + +@pytest.mark.parametrize("geo, basis", SYSTEMS) +def test_par_matches_serial(geo, basis): + mol = build(geo, basis) + mf = pyscf.scf.RHF(mol) + mf.verbose = 0 + mf.conv_tol = 1e-10 + mf.max_cycle = 200 + mf.kernel() + P = mf.make_rdm1() + + # dS_par takes the energy-weighted density directly; dSf builds it + # internally via getF, so construct the same Q here to compare like with + # like. + h = mol.intor("int1e_kin") + mol.intor("int1e_nuc") + eri = mol.intor("int2e") + F = (h + np.einsum("kl,ijkl->ij", P, eri) + - 0.5 * np.einsum("kl,ikjl->ij", P, eri)) + Q = P @ F @ P + + ser = { + "dHcore": np.asarray(librint.dscf.dHcoref(mol, P)), + "dR": np.asarray(librint.dscf.dRf(mol, P)), + "dS": np.asarray(librint.dscf.dSf(mol, P)), + "danalytical": np.asarray(librint.dscf.danalyticalf(mol, P)), + } + + # The SCF above dominates the runtime, so sweep threads inside one test + # rather than parametrizing over them and paying for it once per count. + errs = {} + for T in thread_counts(): + par = { + "dHcore": librint.dscf.dHcore_par(mol, P, T), + "dR": librint.dscf.dR_par(mol, P, T), + "dS": librint.dscf.dS_par(mol, Q, T), + "danalytical": librint.dscf.danalytical_par(mol, P, T), + } + errs[T] = {k: rel(par[k], ser[k]) for k in ser} + + bad = {T: {k: e for k, e in row.items() if not (e < RTOL)} + for T, row in errs.items()} + bad = {T: row for T, row in bad.items() if row} + assert not bad, ( + f"{geo}/{basis}: parallel disagrees with serial beyond {RTOL:.0e}\n" + + "\n".join(f" T={T:3d} " + " ".join(f"{k}={e:.2e}" + for k, e in sorted(row.items())) + for T, row in sorted(errs.items())) + ) + + +@pytest.mark.parametrize("geo, basis", SYSTEMS[:3]) +def test_par_run_to_run(geo, basis): + """Two identical calls must agree to round-off. + + NOT a bitwise check: work stealing decides the fold chunking at run time, + so repeated runs may associate the sum differently. This pins down that the + variation stays at round-off rather than growing into something a caller + would notice. + """ + mol = build(geo, basis) + mf = pyscf.scf.RHF(mol) + mf.verbose = 0 + mf.conv_tol = 1e-10 + mf.max_cycle = 200 + mf.kernel() + P = mf.make_rdm1() + + T = thread_counts()[-1] + a = librint.dscf.danalytical_par(mol, P, T) + b = librint.dscf.danalytical_par(mol, P, T) + assert rel(b, a) < RTOL, f"{geo}/{basis}: run-to-run spread at T={T}" diff --git a/src/lib.rs b/src/lib.rs index 9ec2397..d68c531 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod scf; pub mod dscf; +pub mod par; pub mod p2c; pub mod linalg; diff --git a/src/p2c.rs b/src/p2c.rs index 047288a..158a2bc 100644 --- a/src/p2c.rs +++ b/src/p2c.rs @@ -49,7 +49,7 @@ unsafe fn c2r_arr( // free_c is called on it -- as a boxed slice, so capacity == len and free_c can // reconstruct the Vec exactly. Callers that drop the pointer leak the whole // buffer, which for int2e_c is nao^4 doubles per call. -fn leak_vec(v: Vec) -> *mut f64 { +pub(crate) fn leak_vec(v: Vec) -> *mut f64 { let mut b = v.into_boxed_slice(); let ptr = b.as_mut_ptr(); std::mem::forget(b); diff --git a/src/par.rs b/src/par.rs new file mode 100644 index 0000000..be517f1 --- /dev/null +++ b/src/par.rs @@ -0,0 +1,614 @@ +#![allow(non_snake_case)] +//! Shared-memory parallel gradient accumulation: rayon work-stealing over +//! shell PAIRS. +//! +//! Two design points, both measured rather than assumed (see +//! python/pyscf_comp/bench_par_loops.py): +//! +//! 1. **Granularity: pairs, not the bra index.** Splitting the 2e loop over `i` +//! alone leaves one shell holding 7-11% of the total work, so makespan >= +//! that item and no scheduler helps beyond ~9-14 threads -- block, stride +//! and dynamic all collapse to the same efficiency there. Over `(i, j)` +//! pairs the largest item is 0.4-1.4%, lifting the ceiling to ~70-250. +//! +//! 2. **Dynamic, not static.** At pair granularity on C4H10/def2-svp, T=64: +//! block 0.40, stride 0.71, work-stealing 1.00. A static stride is only +//! worth using where work stealing is unavailable (e.g. across MPI ranks). +//! +//! Per task the closure carries its own scratch and its own clones of +//! atm/bas/env: `#[autodiff_reverse]` demands `&mut` on arguments the loop only +//! reads, and those arrays are kilobytes, so cloning is far cheaper than +//! reworking every signature in the call chain. rayon's `fold` runs the init +//! closure once per work chunk rather than once per pair, so the clones are +//! amortized. The reduction is a sum of `env2.len()` doubles (24-84). +//! +//! ## Pool ownership +//! +//! The `*_in` functions do no pool management -- they run on whatever pool the +//! caller is already inside. The public wrappers add `in_pool`, and +//! `danalytical_par` builds ONE pool for all four loops rather than one each. +//! `nthreads = 0` skips the build entirely and uses rayon's global pool, which +//! honours `RAYON_NUM_THREADS`; that is the preferred path for repeated calls. + +use rayon::prelude::*; + +use crate::cint::CINTOpt; +use crate::cint2e::{cint2e_cart, cint2e_cart_optimizer}; +use crate::cint_bas::CINTcgto_cart; +use crate::dscf::{dkin, dnuc, dovlp, dtwo_ad}; +use crate::linalg::matmult; +use crate::optimizer::CINTdel_optimizer; +use crate::p2c::leak_vec; +use crate::scf::{angl, integral1e, nmol}; +use crate::utils::split; + +/// Per-task state: the mutable copies the autodiff wrappers require, plus the +/// running adjoint accumulator. +struct Task { + atm: Vec, + bas: Vec, + env1: Vec, + env2: Vec, + shls: Vec, + acc: Vec, +} + +fn task_init(atm: &[i32], bas: &[i32], env1: &[f64], env2: &[f64]) -> Task { + Task { + atm: atm.to_vec(), + bas: bas.to_vec(), + env1: env1.to_vec(), + env2: env2.to_vec(), + shls: vec![0i32; 4], + acc: vec![0.0; env2.len()], + } +} + +fn shell_offsets(nbas: usize, bas: &[i32]) -> Vec { + let mut offs = vec![0usize; nbas + 1]; + for s in 0..nbas { + offs[s + 1] = offs[s] + CINTcgto_cart(s, bas) as usize; + } + offs +} + +fn in_pool(nthreads: usize, f: impl FnOnce() -> R + Send) -> R { + if nthreads == 0 { + return f(); // rayon's global pool, sized by RAYON_NUM_THREADS + } + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(nthreads) + .build() + .expect("rayon pool"); + pool.install(f) +} + +fn add_into(acc: &mut [f64], x: &[f64]) { + for t in 0..acc.len() { + acc[t] += x[t]; + } +} + +/// The shape shared by every 1e term: one Enzyme reverse per shell pair over +/// the full `nbas x nbas` loop (no permutational reduction), each seeded with +/// the caller's weight matrix `W`. +/// +/// `rev` is one of the generated reverses `dovlp` / `dkin` / `dnuc`, which have +/// identical signatures. It is taken as a generic rather than a `fn` pointer so +/// each instantiation monomorphizes to a direct call -- an indirect call into +/// an Enzyme-generated body is exactly the kind of thing that type-checks and +/// then misbehaves under fat LTO. +fn pair_1e_in( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + W: &[f64], + rev: F, +) -> Vec +where + F: Fn( + &mut [f64], + &mut [f64], + &mut [i32], + &mut [i32], + &mut [i32], + &mut [f64], + &mut [f64], + &mut [f64], + ) + Sync + + Send, +{ + let (_, nbas) = nmol(atm, bas); + let nshells = angl(bas, 0); + let nparam = env2.len(); + let offs = shell_offsets(nbas, bas); + + let pairs: Vec<(usize, usize)> = + (0..nbas).flat_map(|i| (0..nbas).map(move |j| (i, j))).collect(); + + pairs + .par_iter() + .fold( + || task_init(atm, bas, env1, env2), + |mut st, &(i, j)| { + let (di, dj) = (offs[i + 1] - offs[i], offs[j + 1] - offs[j]); + st.shls[0] = i as i32; + st.shls[1] = j as i32; + + let mut buf = vec![0.0; di * dj]; + let mut dbuf = vec![0.0; di * dj]; + let mut c: usize = 0; + for nuj in offs[j]..offs[j + 1] { + for mui in offs[i]..offs[i + 1] { + dbuf[c] = W[nuj * nshells + mui]; + c += 1; + } + } + + let mut denv = vec![0.0; nparam]; + rev( + &mut buf, &mut dbuf, &mut st.shls, &mut st.atm, &mut st.bas, + &mut st.env1, &mut st.env2, &mut denv, + ); + add_into(&mut st.acc, &denv); + st + }, + ) + .map(|st| st.acc) + .reduce( + || vec![0.0; nparam], + |mut a, b| { + add_into(&mut a, &b); + a + }, + ) +} + +/// Threaded `dscf::dSf`: the overlap term, seeded with the energy-weighted +/// density `Q = P F P`. Note this is `dSf`, not `dSg` -- it does NOT build `F`. +pub fn dS_par( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + Q: &[f64], + nthreads: usize, +) -> Vec { + in_pool(nthreads, || pair_1e_in(atm, bas, env1, env2, Q, dovlp)) +} + +fn dHcore_par_in( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + P: &[f64], +) -> Vec { + // Same term order as dscf::dHcoreg: dT then dV, summed afterwards. + let mut dH = pair_1e_in(atm, bas, env1, env2, P, dkin); + let dV = pair_1e_in(atm, bas, env1, env2, P, dnuc); + add_into(&mut dH, &dV); + dH +} + +/// Threaded `dscf::dHcoreg`. +pub fn dHcore_par( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + P: &[f64], + nthreads: usize, +) -> Vec { + in_pool(nthreads, || dHcore_par_in(atm, bas, env1, env2, P)) +} + +fn dR_par_in( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + P: &[f64], +) -> Vec { + let (_, nbas) = nmol(atm, bas); + let nshells = angl(bas, 0); + let nparam = env2.len(); + let offs = shell_offsets(nbas, bas); + + let w = |a: usize, b: usize, c: usize, d: usize| -> f64 { + 0.5 * (P[a * nshells + b] * P[c * nshells + d] + - 0.5 * P[a * nshells + c] * P[b * nshells + d]) + }; + + // canonical pairs i >= j. Largest-first: rayon splits the range from the + // front, and cost grows with i, so handing out the expensive pairs early + // leaves the cheap ones as filler for whoever finishes first. + let mut pairs: Vec<(usize, usize)> = + (0..nbas).flat_map(|i| (0..=i).map(move |j| (i, j))).collect(); + pairs.reverse(); + + pairs + .par_iter() + .fold( + || task_init(atm, bas, env1, env2), + |mut st, &(i, j)| { + let (di, dj) = (offs[i + 1] - offs[i], offs[j + 1] - offs[j]); + let (mu, nu) = (offs[i], offs[j]); + st.shls[0] = i as i32; + st.shls[1] = j as i32; + + for k in 0..=i { + let dk = offs[k + 1] - offs[k]; + let sig = offs[k]; + st.shls[2] = k as i32; + let lmax = if k == i { j } else { k }; + for l in 0..=lmax { + let dl = offs[l + 1] - offs[l]; + let lam = offs[l]; + st.shls[3] = l as i32; + + let imgs = [ + (i, j, k, l), (j, i, k, l), (i, j, l, k), (j, i, l, k), + (k, l, i, j), (l, k, i, j), (k, l, j, i), (l, k, j, i), + ]; + let mut keep = [true; 8]; + for a in 1..8 { + for b in 0..a { + if imgs[a] == imgs[b] { + keep[a] = false; + break; + } + } + } + + let mut buf = vec![0.0; di * dj * dk * dl]; + let mut dbuf = vec![0.0; di * dj * dk * dl]; + let mut c: usize = 0; + for laml in lam..(lam + dl) { + for sigk in sig..(sig + dk) { + for nuj in nu..(nu + dj) { + for mui in mu..(mu + di) { + let mut ws = 0.0; + if keep[0] { ws += w(mui, nuj, sigk, laml); } + if keep[1] { ws += w(nuj, mui, sigk, laml); } + if keep[2] { ws += w(mui, nuj, laml, sigk); } + if keep[3] { ws += w(nuj, mui, laml, sigk); } + if keep[4] { ws += w(sigk, laml, mui, nuj); } + if keep[5] { ws += w(laml, sigk, mui, nuj); } + if keep[6] { ws += w(sigk, laml, nuj, mui); } + if keep[7] { ws += w(laml, sigk, nuj, mui); } + dbuf[c] = ws; + c += 1; + } + } + } + } + + let lmax_sh = (st.bas[8 * i + 1]) + .max(st.bas[8 * j + 1]) + .max(st.bas[8 * k + 1]) + .max(st.bas[8 * l + 1]); + if lmax_sh as usize > crate::eri::LMAX || st.env1[8] != 0.0 { + panic!( + "2e gradient unsupported for this quartet (max l = \ + {}, omega = {}): the memory-safe reverse (eri.rs) \ + covers only cartesian l <= {}, no range separation.", + lmax_sh, st.env1[8], crate::eri::LMAX, + ); + } + + let mut denv = vec![0.0; nparam]; + dtwo_ad( + &mut buf, &mut dbuf, &mut st.shls, &mut st.atm, + &mut st.bas, &mut st.env1, &mut st.env2, &mut denv, + ); + add_into(&mut st.acc, &denv); + } + } + st + }, + ) + .map(|st| st.acc) + .reduce( + || vec![0.0; nparam], + |mut a, b| { + add_into(&mut a, &b); + a + }, + ) +} + +/// Threaded `dscf::dRf`: the 2e basis-parameter gradient. Keeps the serial +/// 8-fold permutational reduction exactly -- only canonical quartets are +/// differentiated, and each element seed sums the energy weight over the +/// quartet's distinct permutation images. Work is distributed over the +/// canonical `(i, j)` pairs; the `k`/`l` bounds depend only on `i` and `j`, so a +/// pair is a self-contained unit. +pub fn dR_par( + atm: &[i32], + bas: &[i32], + env1: &[f64], + env2: &[f64], + P: &[f64], + nthreads: usize, +) -> Vec { + in_pool(nthreads, || dR_par_in(atm, bas, env1, env2, P)) +} + +/// libcint's shell-quartet optimizer, shared read-only across worker threads. +/// +/// `CINTOpt` is built once and then only read while integrals are evaluated: +/// every `(*opt).` access in the cint2e evaluation path loads into a local, +/// none store back. This is the same contract that lets pyscf hand a single +/// `CINTOpt` to every OpenMP thread. Per-task optimizers would be the +/// conservative alternative, but its tables are O(nbas^2) and would be rebuilt +/// once per work chunk, which at 64 threads costs more memory than the whole +/// rest of the gradient. +struct SharedOpt(*mut CINTOpt); +unsafe impl Send for SharedOpt {} +unsafe impl Sync for SharedOpt {} + +impl SharedOpt { + /// Read the pointer through a method, not the field. Closures capture + /// disjoint fields, so `shared.0` inside one would capture the bare + /// `*mut CINTOpt` -- which is not `Sync` -- and lose this wrapper entirely. + fn get(&self) -> *mut CINTOpt { + self.0 + } +} + +/// Per-task state for the primal Fock build. +struct FockTask { + atm: Vec, + bas: Vec, + env: Vec, + shls: Vec, + G: Vec, +} + +/// Threaded `scf::integral2e_fock` (cartesian only, which is what `getF` uses). +/// +/// Each task accumulates into its own `G`: the Coulomb term writes +/// `G[mui, nuj]`, which is disjoint across `j`, but the exchange term writes +/// `G[mui, laml]` with `l` running over every shell, so two `(i, j)` tasks +/// sharing an `i` collide. Private accumulators plus a sum reduction avoid +/// that without atomics or locking; the cost is one `nao^2` buffer per work +/// chunk. +fn fock2e_par_in( + atm: &[i32], + bas: &[i32], + env: &[f64], + P: &[f64], +) -> Vec { + let (natm, nbas) = nmol(atm, bas); + let n = angl(bas, 0); + let offs = shell_offsets(nbas, bas); + + // The optimizer's tables may reference these, so they outlive the loop. + let mut atm_o = atm.to_vec(); + let mut bas_o = bas.to_vec(); + let mut env_o = env.to_vec(); + let mut opt: *mut CINTOpt = std::ptr::null_mut(); + unsafe { + cint2e_cart_optimizer( + &mut opt, atm_o.as_mut_ptr(), natm as i32, + bas_o.as_mut_ptr(), nbas as i32, env_o.as_mut_ptr(), + ); + } + let shared = SharedOpt(opt); + + // Full nbas^2 pair list (no permutational reduction here, matching the + // serial build), heaviest pairs first so the tail is cheap filler. + let mut pairs: Vec<(usize, usize)> = + (0..nbas).flat_map(|i| (0..nbas).map(move |j| (i, j))).collect(); + pairs.sort_by_key(|&(i, j)| { + std::cmp::Reverse((offs[i + 1] - offs[i]) * (offs[j + 1] - offs[j])) + }); + + let G = pairs + .par_iter() + .fold( + || FockTask { + atm: atm.to_vec(), + bas: bas.to_vec(), + env: env.to_vec(), + shls: vec![0i32; 4], + G: vec![0.0; n * n], + }, + |mut st, &(i, j)| { + let (di, dj) = (offs[i + 1] - offs[i], offs[j + 1] - offs[j]); + let (mu, nu) = (offs[i], offs[j]); + st.shls[0] = i as i32; + st.shls[1] = j as i32; + + for k in 0..nbas { + let dk = offs[k + 1] - offs[k]; + let sig = offs[k]; + st.shls[2] = k as i32; + for l in 0..nbas { + let dl = offs[l + 1] - offs[l]; + let lam = offs[l]; + st.shls[3] = l as i32; + + let mut buf = vec![0.0; di * dj * dk * dl]; + cint2e_cart( + &mut buf, &mut st.shls, &mut st.atm, natm as i32, + &mut st.bas, nbas as i32, &mut st.env, shared.get(), + ); + + let mut c: usize = 0; + for laml in lam..(lam + dl) { + for sigk in sig..(sig + dk) { + for nuj in nu..(nu + dj) { + for mui in mu..(mu + di) { + let v = buf[c]; + c += 1; + st.G[mui * n + nuj] += P[laml * n + sigk] * v; + st.G[mui * n + laml] += + -0.5 * P[nuj * n + sigk] * v; + } + } + } + } + } + } + st + }, + ) + .map(|st| st.G) + .reduce( + || vec![0.0; n * n], + |mut a, b| { + for t in 0..a.len() { + a[t] += b[t]; + } + a + }, + ); + + unsafe { + CINTdel_optimizer(&mut opt); + } + G +} + +/// Threaded `dscf::getF`. The two 1e builds stay serial: they are `nbas^2` +/// primals next to an `nbas^4` one. +fn getF_par_in( + atm: &mut [i32], + bas: &mut [i32], + env: &mut [f64], + P: &[f64], +) -> Vec { + let T = integral1e(atm, bas, env, 0, 1); + let V = integral1e(atm, bas, env, 0, 2); + let G = fock2e_par_in(atm, bas, env, P); + + let mut F = vec![0.0; T.len()]; + for i in 0..F.len() { + F[i] = T[i] + V[i] + G[i]; + } + F +} + +/// Threaded `dscf::danalyticalg`: the whole frozen-P basis-parameter gradient, +/// `dHcore + dR - 0.5 dS`, on one pool. +/// +/// Every `nbas^4` term is threaded, including the primal Fock build that `dSg` +/// hides inside `getF` -- that one is not an Enzyme reverse, but it is the same +/// order of work and leaving it serial caps the whole gradient by Amdahl. See +/// python/pyscf_comp/bench_grad_breakdown.py for the measured split. +pub fn danalytical_par( + atm: &mut [i32], + bas: &mut [i32], + env: &mut [f64], + P: &[f64], + nthreads: usize, +) -> Vec { + let nshells = angl(bas, 0); + + let (s1, s2) = split(bas); + let env1: Vec = env[0..s1].to_vec(); + let env2: Vec = env[s1..s2].to_vec(); + let atm_r: Vec = atm.to_vec(); + let bas_r: Vec = bas.to_vec(); + + let (dH, dR, dS) = in_pool(nthreads, || { + // Q = P F P, the energy-weighted density that seeds the overlap term. + // dscf::dSg builds this internally; hoisted here so F is built once + // and so the Fock build shares this pool instead of running serially. + let F = getF_par_in(atm, bas, env, P); + let pf = matmult(nshells, P, &F); + let Q = matmult(nshells, &pf, P); + + let dH = dHcore_par_in(&atm_r, &bas_r, &env1, &env2, P); + let dR = dR_par_in(&atm_r, &bas_r, &env1, &env2, P); + let dS = pair_1e_in(&atm_r, &bas_r, &env1, &env2, &Q, dovlp); + (dH, dR, dS) + }); + + let mut dtotal = vec![0.0; dH.len()]; + for i in 0..dtotal.len() { + dtotal[i] = dH[i] + dR[i] - 0.5 * dS[i]; + } + dtotal +} + +fn c_args( + atm_p: *mut i32, atm_l: usize, + bas_p: *mut i32, bas_l: usize, + env_p: *mut f64, env_l: usize, + W_p: *mut f64, W_l: usize, +) -> (Vec, Vec, Vec, Vec, Vec) { + let atm: Vec = unsafe { std::slice::from_raw_parts(atm_p, atm_l) }.to_vec(); + let bas: Vec = unsafe { std::slice::from_raw_parts(bas_p, bas_l) }.to_vec(); + let env: Vec = unsafe { std::slice::from_raw_parts(env_p, env_l) }.to_vec(); + let W: Vec = unsafe { std::slice::from_raw_parts(W_p, W_l) }.to_vec(); + let (s1, s2) = split(&bas); + let env1 = env[0..s1].to_vec(); + let env2 = env[s1..s2].to_vec(); + (atm, bas, env1, env2, W) +} + +/// `dS_par` over the C ABI. W is the nshells x nshells adjoint seed (Q = P F P +/// for the overlap term -- `dscf::dSg` builds it internally, this takes it). +/// nthreads = 0 uses rayon's global pool. +#[no_mangle] +pub extern "C" fn dS_par_c( + atm_p: *mut i32, atm_l: usize, + bas_p: *mut i32, bas_l: usize, + env_p: *mut f64, env_l: usize, + W_p: *mut f64, W_l: usize, + nthreads: usize, +) -> *mut f64 { + let (atm, bas, env1, env2, W) = + c_args(atm_p, atm_l, bas_p, bas_l, env_p, env_l, W_p, W_l); + leak_vec(dS_par(&atm, &bas, &env1, &env2, &W, nthreads)) +} + +/// `dR_par` over the C ABI. W is the density matrix P. +#[no_mangle] +pub extern "C" fn dR_par_c( + atm_p: *mut i32, atm_l: usize, + bas_p: *mut i32, bas_l: usize, + env_p: *mut f64, env_l: usize, + W_p: *mut f64, W_l: usize, + nthreads: usize, +) -> *mut f64 { + let (atm, bas, env1, env2, W) = + c_args(atm_p, atm_l, bas_p, bas_l, env_p, env_l, W_p, W_l); + leak_vec(dR_par(&atm, &bas, &env1, &env2, &W, nthreads)) +} + +/// `dHcore_par` over the C ABI. W is the density matrix P. +#[no_mangle] +pub extern "C" fn dHcore_par_c( + atm_p: *mut i32, atm_l: usize, + bas_p: *mut i32, bas_l: usize, + env_p: *mut f64, env_l: usize, + W_p: *mut f64, W_l: usize, + nthreads: usize, +) -> *mut f64 { + let (atm, bas, env1, env2, W) = + c_args(atm_p, atm_l, bas_p, bas_l, env_p, env_l, W_p, W_l); + leak_vec(dHcore_par(&atm, &bas, &env1, &env2, &W, nthreads)) +} + +/// `danalytical_par` over the C ABI: the threaded counterpart of +/// `danalytical_c`. W is the density matrix P; nthreads = 0 uses rayon's +/// global pool. +#[no_mangle] +pub extern "C" fn danalytical_par_c( + atm_p: *mut i32, atm_l: usize, + bas_p: *mut i32, bas_l: usize, + env_p: *mut f64, env_l: usize, + W_p: *mut f64, W_l: usize, + nthreads: usize, +) -> *mut f64 { + let mut atm: Vec = unsafe { std::slice::from_raw_parts(atm_p, atm_l) }.to_vec(); + let mut bas: Vec = unsafe { std::slice::from_raw_parts(bas_p, bas_l) }.to_vec(); + let mut env: Vec = unsafe { std::slice::from_raw_parts(env_p, env_l) }.to_vec(); + let P: Vec = unsafe { std::slice::from_raw_parts(W_p, W_l) }.to_vec(); + leak_vec(danalytical_par(&mut atm, &mut bas, &mut env, &P, nthreads)) +}