From 877c9bb4438fc74df4e04937bae41278d5c34938 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Wed, 29 Jul 2026 11:48:11 -0400 Subject: [PATCH 01/11] par: rayon work-stealing over shell pairs for the gradient src/par.rs is a threaded counterpart to the dscf.rs gradient entry points: dS_par, dHcore_par, dR_par and danalytical_par, each accumulating the basis-parameter adjoint over rayon work-stealing. The serial path is deliberately untouched -- danalyticalf stays byte-for-byte what the finite-difference validation covers, so "parallel matches serial" remains a statement about two independent implementations. Enzyme reverses turn out to be reentrant. Nothing guaranteed that -- the generated bodies carry their own tape -- and if they had not been, the answer was MPI processes rather than shared memory, a different design and not a tuning knob. 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 beats reworking every signature in the call chain. rayon's fold runs its init closure once per work chunk, so the clones are amortized rather than paid per pair. Two design points, both measured rather than assumed (python/pyscf_comp/bench_par_loops.py measures both): 1. Distribute over (i, j) shell 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 block, stride and work-stealing all collapse to the same efficiency past ~9-14 threads. Over pairs the largest item is 0.4-1.4%, lifting the ceiling to ~70-250. The k/l bounds depend only on i and j, so a canonical pair is self-contained and the 8-fold permutational reduction is untouched. 2. Work-stealing, not a static schedule. At pair granularity on C4H10/def2-svp with 64 threads the model gives block 0.40, stride 0.71, work-stealing 1.00. Pairs are emitted largest-first so the expensive ones go out early and the cheap ones are left as filler. pair_1e_in generalizes the 1e shell-pair loop over the Enzyme reverse, so dT and dV get the same treatment dS has. It takes the reverse as a monomorphized generic rather than a fn pointer -- an indirect call into an Enzyme-generated body type-checks and then misbehaves under fat LTO. danalytical_par assembles dHcore + dR - 0.5 dS and builds ONE rayon pool for all four loops instead of one per loop; nthreads = 0 skips the pool build and uses the global pool, so RAYON_NUM_THREADS applies. That is the preferred path for repeated calls: at ~10 ms of work the 1e loop peaks near 5-7x and then gets slower past 16 threads, because building a ThreadPool per call dominates. Every parameter is a slice, matching the convention main adopted in 0d46255 -- a Fn bound needs the exact type and deref coercion does not apply, so F: Fn(&mut Vec, ...) would not have matched dovlp/dkin/dnuc. leak_vec becomes pub(crate) so par.rs can return buffers over the same C ABI and free_c contract as p2c.rs; the four #[no_mangle] pub extern "C" entry points are the FFI boundary _bindings.py calls. rayon 1.11 was already in the tree as a faer dependency, so --offline builds are unaffected. Measured on one exclusive node, 96 cores, against the serial dscf entry points: dR 2e C2H6/def2-tzvp 27.0s -> 0.68s 39.7x @ 64 threads (eff 0.62) 15.4x @ 16 (eff 0.96) H2O/def2-qzvp 34.6s -> 1.28s 27.0x @ 64 (eff 0.42) CH4/def2-svp 0.60s -> 0.026s 23.3x @ 64 (eff 0.36) All parallel results match serial to 2e-14 .. 2.5e-12, which is the round-off from summing the per-task partials in a different order. Measured efficiency at 64 threads falls short of the model's 1.00, which is expected: the model costs quartets as prod(d) * prod(nprim) * nroots and knows nothing about pool construction, memory bandwidth, or Schwarz screening. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 1 + Cargo.toml | 3 + src/lib.rs | 1 + src/p2c.rs | 2 +- src/par.rs | 453 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 459 insertions(+), 1 deletion(-) create mode 100644 src/par.rs 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/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..ba56374 --- /dev/null +++ b/src/par.rs @@ -0,0 +1,453 @@ +#![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_bas::CINTcgto_cart; +use crate::dscf::{dkin, dnuc, dovlp, dtwo_ad, getF}; +use crate::linalg::matmult; +use crate::p2c::leak_vec; +use crate::scf::{angl, 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)) +} + +/// Threaded `dscf::danalyticalg`: the whole frozen-P basis-parameter gradient, +/// `dHcore + dR - 0.5 dS`, on one pool. +/// +/// `getF` (and so `integral2e_fock`) still runs serially, as it does in the +/// serial path; it is a primal `nbas^4` build, and whether it needs threading +/// too is an empirical question -- see python/pyscf_comp/bench_grad_breakdown.py. +pub fn danalytical_par( + atm: &mut [i32], + bas: &mut [i32], + env: &mut [f64], + P: &[f64], + nthreads: usize, +) -> Vec { + let nshells = angl(bas, 0); + + // 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. + let F = getF(atm, bas, env, P); + let pf = matmult(nshells, P, &F); + let Q = matmult(nshells, &pf, P); + + 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, || { + 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)) +} From a5db4597221befefc2ffde99e6847b3954436d00 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Tue, 28 Jul 2026 23:19:48 -0400 Subject: [PATCH 02/11] par: thread the Fock build too, or Amdahl caps the gradient at 6.5x bench_grad_breakdown on CH4/def2-svp, with dS and dR threaded and everything else serial: dR 2e pair loop 0.784s 84.3% dR_par getF = int2e_fock + PFP 0.130s 14.0% serial dHcore (dT+dV) 0.012s 1.3% serial dS pair loop 0.003s 0.4% dS_par serial fraction 0.153 -> capped at 6.5x at infinite threads So the 2e reverse scaling 40x bought a gradient that could never beat 6.5x. getF was invisible in that work because it is a primal libcint build, not an Enzyme reverse, and the parallelization effort had been aimed at the autodiff loops. H2O/def2-qzvp is the same shape: getF 5.6s against dR 47s. fock2e_par_in threads scf::integral2e_fock over shell pairs. Two things it has to get right: - Private accumulators. The Coulomb term writes G[mui,nuj], disjoint across j, but the exchange term writes G[mui,laml] with l over every shell, so two (i,j) tasks sharing an i collide. Each task folds into its own nao^2 buffer and the reduction sums them. - One shared CINTOpt. Every (*opt). access in the cint2e evaluation path loads into a local and none store back, which is the contract that lets pyscf hand one optimizer to every OpenMP thread. Per-task optimizers would be the conservative choice but its tables are O(nbas^2), rebuilt once per work chunk, and at 64 threads that costs more memory than the gradient. The wrapper needs a get() method rather than a public field: closures capture disjoint fields, so `shared.0` would capture the bare *mut CINTOpt and drop the Sync wrapper on the floor. danalytical_par now runs getF inside the same pool as the three adjoint loops, so one pool covers the whole gradient. Left serial: the O(nao^3) P F P matmults and the nbas^2 1e primals. Co-Authored-By: Claude Opus 5 --- src/par.rs | 183 +++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 172 insertions(+), 11 deletions(-) diff --git a/src/par.rs b/src/par.rs index ba56374..be517f1 100644 --- a/src/par.rs +++ b/src/par.rs @@ -32,11 +32,14 @@ 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, getF}; +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, nmol}; +use crate::scf::{angl, integral1e, nmol}; use crate::utils::split; /// Per-task state: the mutable copies the autodiff wrappers require, plus the @@ -333,12 +336,169 @@ pub fn dR_par( 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. /// -/// `getF` (and so `integral2e_fock`) still runs serially, as it does in the -/// serial path; it is a primal `nbas^4` build, and whether it needs threading -/// too is an empirical question -- see python/pyscf_comp/bench_grad_breakdown.py. +/// 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], @@ -348,12 +508,6 @@ pub fn danalytical_par( ) -> Vec { let nshells = angl(bas, 0); - // 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. - let F = getF(atm, bas, env, P); - let pf = matmult(nshells, P, &F); - let Q = matmult(nshells, &pf, P); - let (s1, s2) = split(bas); let env1: Vec = env[0..s1].to_vec(); let env2: Vec = env[s1..s2].to_vec(); @@ -361,6 +515,13 @@ pub fn danalytical_par( 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); From 9ef376551bbf4c76a9f8b490bef1cb9d1cc02e01 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Thu, 30 Jul 2026 22:22:32 -0400 Subject: [PATCH 03/11] python: bind the threaded gradient entry points, and degrade if they are absent librint.dscf gains dS_par, dR_par, dHcore_par and danalytical_par, with argtypes for dS_par_c / dR_par_c / dHcore_par_c / danalytical_par_c in _bindings.py. Until now the Python side went through the serial danalytical_c only, so the 40x on dR was a number about a loop rather than about a gradient. The symbols are bound defensively rather than eagerly. Any .so built before src/par.rs existed has none of them, and resolving them at import time turned that into an `import librint` failure for everyone, including callers who only want the serial path. HAS_PAR records their absence so dscf can raise something actionable and the test suite can skip. (This does not rescue the .so committed in python/librint/, which is missing free_c and every other FFI entry point and so cannot satisfy `import librint` on main either -- that one needs a rebuild, which is not this PR's business.) Co-Authored-By: Claude Opus 5 --- python/librint/_bindings.py | 29 ++++++++++++++++++++ python/librint/dscf.py | 54 +++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) 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) From 6447f75472aa5fda28d6f901b41c04b203ae44bb Mon Sep 17 00:00:00 2001 From: AbAldo Date: Thu, 30 Jul 2026 22:22:33 -0400 Subject: [PATCH 04/11] test: parallel/serial equivalence, and a threaded finite-difference case python/tests/test_par_equiv.py compares every threaded term against its serial counterpart at 1..64 threads. The criterion is not bitwise -- work stealing reassociates the sum -- it is that the error stays FLAT in thread count, which is what separates round-off from a race. Measured 1e-16 (dHcore) to 5e-12 (assembled, where large terms cancel), constant across T, and repeat runs at fixed T come out bitwise identical. The thread sweep runs inside a single test per system rather than as a parametrize over thread counts, because the SCF dominates the runtime and parametrizing would pay for it once per count. The five basis sets that cost minutes (def2-tzvp and up) are marked `slow`, so `pytest -m "not slow"` is an 8-second loop while the full run still covers the f-shell, g-shell and general-contraction paths those systems exist to reach; pyproject declares the marker. test_gradient_fd.py gains a threaded case. It is redundant on paper -- test_par_equiv ties par to serial and the existing test ties serial to finite differences -- but the transitive argument breaks silently if either link is weakened, and a direct check does not. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 7 ++ python/tests/test_gradient_fd.py | 28 ++++++ python/tests/test_par_equiv.py | 160 +++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+) create mode 100644 python/tests/test_par_equiv.py 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/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}" From e172b82857bfc63db7579f05497c9916c0072d9f Mon Sep 17 00:00:00 2001 From: AbAldo Date: Wed, 29 Jul 2026 11:07:23 -0400 Subject: [PATCH 05/11] bench: measurement scripts for the threaded gradient Three scripts, none of which had a counterpart for the parallel path before, plus a threaded mode for the existing finite-difference validation: bench_par_loops.py the two pair loops alone bench_grad_breakdown.py where the serial time goes bench_par_scaling.py the whole gradient, end to end bench_par_loops checks dS_par and dR_par against the serial dscf entry points and sweeps thread counts. Its speedup baseline is the same code at one thread, not dSf/dRf: the serial wall time also includes the getF Fock build the parallel loops do not perform, so the other comparison would flatter them. bench_grad_breakdown times each term separately, because getF -> integral2e_fock is an nbas^4 primal loop that danalyticalg pays on every call and that threading the Enzyme reverses does not touch. That is not a rhetorical point: with only dS and dR threaded it measured the remaining serial fraction at 0.153 on CH4/def2-svp, a hard ceiling of 6.5x on the whole gradient no matter how many cores, against a 2e reverse scaling ~40x. getF is not exported on its own, so it is inferred as t(dSf) - t(dS_par @ 1 thread) -- both run the same dS pair loop, only dSf additionally builds F and forms PFP. bench_par_scaling times the assembled danalytical_par against both baselines: vs danalyticalf (what a caller gains) and vs par@1 (how good the threading is). Reporting only the first would credit threading for serial-path overhead it happens to avoid; only the second would hide overhead the threaded path adds. The gap between them is itself the interesting number. validate_grad_fd.py --par N runs the 12-system finite-difference validation through the threaded path and reports the serial error alongside, so one run covers both. The two benchmarks emit JSON (bench_scaling_results.json, bench_breakdown_results.json) rather than only printing, because plot_alkanes.py's rule is that every plotted point comes from a results file and the machine is read from a _meta block. Their system lists are identical, so panels drawn from them describe the same four systems without a per-panel caveat. Both run as SLURM cpu jobs rather than inline. They are named bench_, not check_: `check_` promises a verdict, and bench_grad_breakdown has no non-zero exit path at all -- it only reports where the time goes. They keep their inline correctness assertions, because a benchmark that silently times a wrong result is worse than useless, but as a guard rather than as the point. The one script here that WAS purely a verdict is python/tests/test_par_equiv.py. Co-Authored-By: Claude Opus 5 --- python/pyscf_comp/bench_grad_breakdown.py | 165 ++++++++++++++++++++++ python/pyscf_comp/bench_par_loops.py | 129 +++++++++++++++++ python/pyscf_comp/bench_par_scaling.py | 118 ++++++++++++++++ python/pyscf_comp/validate_grad_fd.py | 44 +++++- 4 files changed, 450 insertions(+), 6 deletions(-) create mode 100644 python/pyscf_comp/bench_grad_breakdown.py create mode 100644 python/pyscf_comp/bench_par_loops.py create mode 100644 python/pyscf_comp/bench_par_scaling.py diff --git a/python/pyscf_comp/bench_grad_breakdown.py b/python/pyscf_comp/bench_grad_breakdown.py new file mode 100644 index 0000000..2f4f724 --- /dev/null +++ b/python/pyscf_comp/bench_grad_breakdown.py @@ -0,0 +1,165 @@ +"""Where does danalyticalf's wall time actually go? + +danalyticalg = dHcoreg + dRg - 0.5 dSg, and dSg internally builds the Fock +matrix (getF -> integral2e_fock, an O(nbas^4) primal loop) to form the +energy-weighted density Q = P F P. Amdahl's law is decided by the terms that +stay serial, not by the ones that speed up, so this measures every term +separately. + +That is not a rhetorical point. When dS and dR were the only threaded terms, +this script measured the remaining serial fraction at 0.153 on CH4/def2-svp -- +a hard ceiling of 6.5x on the whole gradient no matter how many cores it was +given, with the 2e reverse itself scaling ~40x. getF was 14% of that and had +no parallel version at all, because it is a primal build rather than an Enzyme +reverse and so never came up while the autodiff loops were the subject. + +getF is not exported on its own, so it is inferred: + + t(getF) ~= t(dscf.dSf) - t(dS_par @ 1 thread) + +both of which run the same dS pair loop; only dSf additionally builds F and +forms P F P. + +Usage: LIBRINT_SO=/path/to/librint.so python bench_grad_breakdown.py +""" +import ctypes +import json +import os +import platform +import time + +import numpy as np +import pyscf + +import librint +import librint.dscf +import librint.utils +from librint import library + +from bench_fair import _cpu_model, _mem_limit_kb +from geometries import geometries + +# same list, same order as bench_par_scaling.py, so the two JSONs can be shown +# side by side on one figure without a per-system caveat +SYSTEMS = [ + ("CH4", "def2-svp"), + ("C3H8", "def2-svp"), + ("C2H6", "def2-tzvp"), + ("H2O", "def2-qzvp"), +] +OUT_JSON = "bench_breakdown_results.json" + +_SIG = ( + 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, +) +for _fn in (library.dS_par_c, library.dR_par_c): + _fn.argtypes = _SIG + _fn.restype = ctypes.POINTER(ctypes.c_double) + + +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 call_par(fn, mol, W, nthreads): + atm, bas, env, _ = librint.utils.prep(mol) + W = np.ascontiguousarray(W) + s1, s2 = librint.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, + nthreads, + ) + return librint.utils.take(ptr, (s2 - s1,)) + + +def timed(label, fn): + t0 = time.perf_counter() + out = fn() + dt = time.perf_counter() - t0 + print(f" {label:26s} {dt:9.3f}s", flush=True) + return dt, out + + +def main(): + results = {"_meta": {"node": platform.node(), + "ncores": len(os.sched_getaffinity(0)), + "mem_limit_kb": _mem_limit_kb(), + "cpu_model": _cpu_model()}} + for geo, basis in SYSTEMS: + 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() + + 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 + + print(f"\n{geo}/{basis} nao={mol.nao} nbas={mol.nbas}", flush=True) + t_h, _ = timed("dHcoref (dT+dV, serial)", lambda: librint.dscf.dHcoref(mol, P)) + t_r, _ = timed("dRf (2e, serial)", lambda: librint.dscf.dRf(mol, P)) + t_s, _ = timed("dSf (getF+dS, serial)", lambda: librint.dscf.dSf(mol, P)) + t_s1, _ = timed("dS_par (dS only, T=1)", lambda: call_par(library.dS_par_c, mol, Q, 1)) + t_r1, _ = timed("dR_par (2e only, T=1)", lambda: call_par(library.dR_par_c, mol, P, 1)) + t_tot, _ = timed("danalyticalf (total)", lambda: librint.dscf.danalyticalf(mol, P)) + + t_getf = t_s - t_s1 + parts = [ + ("dHcore (dT+dV)", t_h, "dHcore_par"), + ("dR 2e pair loop", t_r1, "dR_par"), + ("getF = int2e_fock + PFP", t_getf, "fock2e_par (PFP serial)"), + ("dS pair loop", t_s1, "dS_par"), + ] + acct = sum(p[1] for p in parts) + print(f" {'-' * 68}", flush=True) + for name, t, status in parts: + print(f" {name:26s} {t:9.3f}s {100 * t / acct:5.1f}% {status}", + flush=True) + print(f" {'accounted':26s} {acct:9.3f}s (danalyticalf {t_tot:.3f}s)", + flush=True) + + # Why every nbas^4 term had to be threaded, not just the 2e reverse: + # this is the ceiling that applied when dS and dR were the only ones + # with a parallel version, no matter how many cores were thrown at it. + was_serial = t_h + t_getf + print(f" when only dS+dR were threaded: serial fraction " + f"{was_serial / acct:.3f} -> capped at {acct / was_serial:5.1f}x " + f"at infinite threads", flush=True) + print(f" now threaded: all four. What is left serial is the O(nao^3) " + f"P F P matmults and the nbas^2 1e primals inside getF.", + flush=True) + + results[f"{geo}/{basis}"] = { + "nao": int(mol.nao), "nbas": int(mol.nbas), + "dHcore": t_h, "dR": t_r1, "getF": t_getf, "dS": t_s1, + "accounted": acct, "danalyticalf": t_tot, + # the ceiling that applied before fock2e_par/dHcore_par existed + "old_serial_frac": was_serial / acct, + "old_ceiling": acct / was_serial, + } + with open(OUT_JSON, "w") as f: + json.dump(results, f, indent=1) + + print(f"\nresults -> {OUT_JSON}") + + +if __name__ == "__main__": + main() diff --git a/python/pyscf_comp/bench_par_loops.py b/python/pyscf_comp/bench_par_loops.py new file mode 100644 index 0000000..32836dd --- /dev/null +++ b/python/pyscf_comp/bench_par_loops.py @@ -0,0 +1,129 @@ +"""Correctness and scaling of the rayon-parallel gradient loops (src/par.rs). + +Checks both halves of the design: + + dS_par 1e overlap term, pairs over the full nbas x nbas loop + dR_par 2e term, pairs over canonical i >= j with the 8-fold reduction intact + +For each, the parallel result must match the serial dscf entry point (to +round-off -- the per-task partial sums are summed in a different order), and the +speedup is measured against the same code at one thread, which is the honest +baseline: the serial dSf/dRf wall time also includes the getF Fock build that +the parallel loops do not perform. + +Usage: LIBRINT_SO=/path/to/librint.so .venv/bin/python bench_par_loops.py +""" +import ctypes +import os +import time + +import numpy as np +import pyscf + +import librint +import librint.dscf +import librint.utils +from librint import library + +from geometries import geometries + +SYSTEMS = [ + ("H2O", "sto-3g"), + ("CH4", "def2-svp"), + ("H2O", "def2-qzvp"), + ("C2H6", "def2-tzvp"), +] +THREADS = [1, 2, 4, 8, 16, 32, 64] + +_SIG = ( + 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, +) +for _fn in (library.dS_par_c, library.dR_par_c): + _fn.argtypes = _SIG + _fn.restype = ctypes.POINTER(ctypes.c_double) + + +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 call_par(fn, mol, W, nthreads): + atm, bas, env, _ = librint.utils.prep(mol) + W = np.ascontiguousarray(W) + s1, s2 = librint.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, + nthreads, + ) + return librint.utils.take(ptr, (s2 - s1,)) + + +def sweep(label, fn, mol, W, ref, ncores): + scale = max(np.abs(ref).max(), 1e-30) + t0 = time.perf_counter() + call_par(fn, mol, W, 1) + t_one = time.perf_counter() - t0 + print(f" {label}: 1-thread {t_one:.3f}s") + failures = 0 + for n in THREADS: + if n > ncores: + continue + t0 = time.perf_counter() + got = call_par(fn, mol, W, n) + dt = time.perf_counter() - t0 + err = float(np.abs(got - ref).max()) + ok = err < 1e-9 * scale + failures += 0 if ok else 1 + print(f" threads={n:3d} {dt:8.4f}s speedup={t_one / dt:6.2f}x " + f"eff={t_one / dt / n:5.2f} max|par-serial|={err:.2e}" + f"{'' if ok else ' <-- MISMATCH'}") + return failures + + +def main(): + ncores = len(os.sched_getaffinity(0)) + print(f"usable cores: {ncores}") + failures = 0 + for geo, basis in SYSTEMS: + mol = build(geo, basis) + mf = pyscf.scf.RHF(mol) + mf.verbose = 0 + mf.conv_tol = 1e-10 + mf.max_cycle = 200 + mf.max_memory = 200 + mf.kernel() + P = mf.make_rdm1() + + # dSg seeds with the energy-weighted density Q = P F P, which it builds + # itself via getF; the parallel entry takes the seed directly. + 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 + + print(f"\n{geo}/{basis} nbas={mol.nbas} nparam={np.diff(librint.utils.split(np.asarray(mol._bas)))[0]}") + failures += sweep("dS 1e", library.dS_par_c, mol, Q, + np.asarray(librint.dscf.dSf(mol, P)), ncores) + failures += sweep("dR 2e", library.dR_par_c, mol, P, + np.asarray(librint.dscf.dRf(mol, P)), ncores) + if failures: + raise SystemExit(f"{failures} parallel result(s) disagreed with serial") + print("\nrayon-parallel loops agree with serial everywhere") + + +if __name__ == "__main__": + main() diff --git a/python/pyscf_comp/bench_par_scaling.py b/python/pyscf_comp/bench_par_scaling.py new file mode 100644 index 0000000..c635bd3 --- /dev/null +++ b/python/pyscf_comp/bench_par_scaling.py @@ -0,0 +1,118 @@ +"""How well does the whole gradient thread, end to end? + +bench_par_loops.py measures the two pair loops in isolation, which flatters +them: it skips getF and the 1e terms, so it answers "how well does dR scale" +rather than "how much faster is a gradient". This measures danalytical_par, +the assembled thing a caller actually invokes, and reports both baselines: + + vs danalyticalf what a user gets by switching to the threaded entry point + vs danalytical_par@1 how good the parallelization itself is + +Reporting only the first would credit threading for any serial-path overhead +it happens to avoid; reporting only the second would hide any overhead the +threaded path adds. The gap between them is itself the interesting number. + +Writes bench_scaling_results.json so the figure is drawn from measurements +rather than from a log scrape, with a _meta block recording the machine. + +Usage: LIBRINT_SO=/path/to/librint.so python bench_par_scaling.py +""" +import json +import os +import platform +import time + +import numpy as np +import pyscf + +import librint +import librint.dscf +import librint.utils + +from bench_fair import _cpu_model, _mem_limit_kb +from geometries import geometries + +SYSTEMS = [ + ("CH4", "def2-svp"), + ("C3H8", "def2-svp"), + ("C2H6", "def2-tzvp"), + ("H2O", "def2-qzvp"), +] +THREADS = [1, 2, 4, 8, 16, 32, 64] +REPEATS = 3 +OUT_JSON = "bench_scaling_results.json" + + +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 median_time(fn, n): + ts = [] + for _ in range(n): + t0 = time.perf_counter() + out = fn() + ts.append(time.perf_counter() - t0) + return float(np.median(ts)), out + + +def main(): + ncores = len(os.sched_getaffinity(0)) + threads = [t for t in THREADS if t <= ncores] + print(f"usable cores: {ncores} thread counts: {threads} " + f"median of {REPEATS}", flush=True) + + results = {"_meta": {"node": platform.node(), "ncores": ncores, + "mem_limit_kb": _mem_limit_kb(), + "cpu_model": _cpu_model(), "repeats": REPEATS}} + failures = 0 + for geo, basis in SYSTEMS: + mol = build(geo, basis) + mf = pyscf.scf.RHF(mol) + mf.verbose = 0 + mf.conv_tol = 1e-10 + mf.max_cycle = 200 + mf.max_memory = 200 + mf.kernel() + P = mf.make_rdm1() + + tag = f"{geo}/{basis}" + print(f"\n{tag} nao={mol.nao} nbas={mol.nbas}", flush=True) + t_ser, g_ser = median_time(lambda: librint.dscf.danalyticalf(mol, P), + REPEATS) + print(f" danalyticalf (serial reference) {t_ser:9.3f}s", flush=True) + scale = max(float(np.abs(g_ser).max()), 1e-30) + + row = {"nao": int(mol.nao), "nbas": int(mol.nbas), + "serial": t_ser, "threads": {}} + t_one = None + for T in threads: + t, g = median_time( + lambda: librint.dscf.danalytical_par(mol, P, T), REPEATS) + if t_one is None: + t_one = t + err = float(np.abs(np.asarray(g) - g_ser).max()) / scale + if not (err < 1e-9): + failures += 1 + row["threads"][str(T)] = {"median": t, "rel_err": err} + print(f" T={T:3d} {t:9.3f}s vs serial {t_ser / t:6.2f}x " + f"vs par@1 {t_one / t:6.2f}x eff={t_one / t / T:5.2f} " + f"rel|par-ser|={err:.2e}" + f"{'' if err < 1e-9 else ' <-- MISMATCH'}", flush=True) + results[tag] = row + with open(OUT_JSON, "w") as f: # write-as-you-go, survive a timeout + json.dump(results, f, indent=1) + + print(f"\nresults -> {OUT_JSON}") + if failures: + raise SystemExit(f"{failures} timed result(s) disagreed with serial") + + +if __name__ == "__main__": + main() diff --git a/python/pyscf_comp/validate_grad_fd.py b/python/pyscf_comp/validate_grad_fd.py index 26a36a6..d0844da 100644 --- a/python/pyscf_comp/validate_grad_fd.py +++ b/python/pyscf_comp/validate_grad_fd.py @@ -2,11 +2,15 @@ differences with a fully reconverged SCF at every perturbed geometry. This is convention-free ground truth (no frozen-P/Q assumptions). -Usage: python validate_grad_fd.py # standard system list +Usage: python validate_grad_fd.py # standard system list, serial path + python validate_grad_fd.py --par 64 # same, through the threaded path """ +import argparse + import numpy as np import pyscf import librint +import librint.dscf import librint.utils from geometries import geometries @@ -51,6 +55,19 @@ def scf_e(mol): def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--par", type=int, default=None, metavar="N", + help="validate the threaded path (src/par.rs) with N rayon threads; " + "0 uses rayon's global pool. Default is the serial path.", + ) + args = ap.parse_args() + + label = ("danalyticalf" if args.par is None + else f"danalyticalf and danalytical_par(T={args.par})") + print(f"validating {label} against central finite differences (h={H})", + flush=True) + failures = 0 for geo, basis in SYSTEMS: mol = build(geo, basis) @@ -71,17 +88,32 @@ def main(): mol._env[j] += H g_true[j - s1] = (ep - em) / (2 * H) - g_ana = librint.dscf.danalyticalf(mol, P) - e_ana = np.abs(g_true - g_ana).max() + # The serial gradient is always measured against FD, so one --par run + # validates both paths: the threaded one directly, the serial one + # alongside it. + g_ser = librint.dscf.danalyticalf(mol, P) + e_ser = np.abs(g_true - g_ser).max() # denergyf is the same assembled path (denergy_c -> denergyfast), so # this is a wiring guard that the two entry points stay in sync, not an # independent gradient check. - wired = np.array_equal(librint.dscf.denergyf(mol, P), g_ana) - ok = e_ana < 1e-5 and wired + wired = np.array_equal(librint.dscf.denergyf(mol, P), g_ser) + extra = f"denergyf_same_path={wired}" + ok = e_ser < 1e-5 and wired + + if args.par is not None: + g_par = librint.dscf.danalytical_par(mol, P, args.par) + e_par = np.abs(g_true - g_par).max() + # Work stealing reassociates the sum, so the threaded result agrees + # with the serial one to round-off rather than bitwise. + scale = max(float(np.abs(g_ser).max()), 1e-30) + d_par = float(np.abs(g_par - g_ser).max()) / scale + extra += f" |true-par|={e_par:.2e} par_vs_serial={d_par:.2e}" + ok = ok and e_par < 1e-5 and d_par < 1e-9 + failures += 0 if ok else 1 print( f"{geo}/{basis:9s} params={s2-s1:3d} |grad|={np.linalg.norm(g_true):9.4f} " - f"|true-danalyticalf|={e_ana:.2e} denergyf_same_path={wired}" + f"|true-serial|={e_ser:.2e} {extra}" f" {'OK' if ok else 'FAIL'}", flush=True, ) From 501a9a6f7d74c5cc1088d123a7fa7ab5e202c817 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Wed, 29 Jul 2026 11:07:24 -0400 Subject: [PATCH 06/11] bench: a scaling figure for the threaded gradient The threading work had tables and no figure, and the numbers only existed as text in a job log. plot_scaling.py follows plot_alkanes.py's rule -- every plotted point comes from a results JSON and the machine is read from a _meta block -- reading bench_scaling_results.json and bench_breakdown_results.json. Both come from a single exclusive SLURM job, so the figure cannot mix provenance. gradient_scaling.{pdf,png} has three panels: 1. speedup vs threads against the ideal diagonal, with a dotted line per system at the Amdahl ceiling that applied when only dS and dR were threaded. The curves crossing those lines is the result; the ceilings are read from the breakdown JSON, not transcribed. 2. parallel efficiency, which is where the honest bad news lives: 0.64 at 64 threads on C3H8/def2-svp but 0.31 on CH4/def2-svp, whose whole gradient is 0.7s, and 0.39 on H2O/def2-qzvp. 3. the 1-core cost breakdown as a stacked bar, getF and dHcore hatched as "was serial" -- 85-89% dR, 11-14% getF. This is why panel 1 needed the primal Fock build threaded and not just the 2e reverse. Measured on cpu-00094, 96-core EPYC 9J14, median of 3: CH4/def2-svp 0.715s -> 0.036s (20.0x), C3H8/def2-svp 20.838s -> 0.505s (41.3x), C2H6/def2-tzvp 31.684s -> 0.823s (38.5x), H2O/def2-qzvp 39.861s -> 1.607s (24.8x), against old ceilings of 6.9x/7.0x/7.8x/8.8x. Co-Authored-By: Claude Opus 5 --- python/pyscf_comp/plot_scaling.py | 174 ++++++++++++++++++++++++++++++ 1 file changed, 174 insertions(+) create mode 100644 python/pyscf_comp/plot_scaling.py diff --git a/python/pyscf_comp/plot_scaling.py b/python/pyscf_comp/plot_scaling.py new file mode 100644 index 0000000..964508f --- /dev/null +++ b/python/pyscf_comp/plot_scaling.py @@ -0,0 +1,174 @@ +"""Paper figure: thread scaling of the basis-parameter gradient. + +Reads bench_scaling_results.json (bench_par_scaling.py) and +bench_breakdown_results.json (bench_grad_breakdown.py) and draws three panels: + + 1. speedup vs threads, against the ideal diagonal, with the Amdahl ceiling + that applied when only dS and dR were threaded drawn as a dashed line per + system. The point of the figure is that the curves cross those lines. + 2. parallel efficiency vs threads. + 3. where the serial time goes, as a stacked bar -- which is *why* panel 1 + needed the Fock build threaded and not just the 2e reverse. + +Same rule as plot_alkanes.py: every plotted point comes from a results JSON, +both files must come from the same job, and the machine is read from their +_meta blocks rather than assumed. + +Usage: python plot_scaling.py [--scaling bench_scaling_results.json] + [--breakdown bench_breakdown_results.json] + [--out gradient_scaling] +Outputs .pdf + .png and a text summary table. +""" +import argparse +import json + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.ticker import FixedLocator, FixedFormatter, NullLocator +import numpy as np + +# term -> (legend label, colour). Order is the stacking order in panel 3. +TERMS = [ + ("dR", "dR 2e reverse", "#1a7f37"), + ("getF", "getF primal Fock build", "#e16f24"), + ("dHcore", "dHcore dT+dV", "#8250df"), + ("dS", "dS overlap", "#0969da"), +] +# the two that had no parallel version before this work; panel 3 hatches them +WAS_SERIAL = {"getF", "dHcore"} +STYLES = [ + dict(color="#1a7f37", marker="o", ls="-"), + dict(color="#0969da", marker="s", ls="--"), + dict(color="#e16f24", marker="^", ls="-."), + dict(color="#cf222e", marker="D", ls=":"), +] + + +def load(path): + with open(path) as f: + return json.load(f) + + +def systems(d): + return [k for k in d if k != "_meta"] + + +def thread_curve(row): + """-> (threads, speedup vs serial, speedup vs par@1) sorted by thread count.""" + ts = sorted(int(t) for t in row["threads"]) + t1 = row["threads"][str(ts[0])]["median"] + med = np.array([row["threads"][str(t)]["median"] for t in ts]) + return np.array(ts), row["serial"] / med, t1 / med + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--scaling", default="bench_scaling_results.json") + ap.add_argument("--breakdown", default="bench_breakdown_results.json") + ap.add_argument("--out", default="gradient_scaling") + args = ap.parse_args() + + sc = load(args.scaling) + bd = load(args.breakdown) + meta = sc.get("_meta") or {} + syss = systems(sc) + + # extra width on the right for the breakdown legend, which lives outside + fig, axes = plt.subplots(1, 3, figsize=(13.2, 3.9), + gridspec_kw={"width_ratios": [1, 1, 1.05]}) + ax_s, ax_e, ax_b = axes + + # ── panel 1: speedup, with the pre-fix Amdahl ceilings ────────────────── + tmax = 1 + for i, tag in enumerate(syss): + th, sp, _ = thread_curve(sc[tag]) + tmax = max(tmax, int(th.max())) + ax_s.plot(th, sp, ms=4, label=tag, **STYLES[i % len(STYLES)]) + ceil = (bd.get(tag) or {}).get("old_ceiling") + if ceil: + ax_s.axhline(ceil, color=STYLES[i % len(STYLES)]["color"], + lw=0.7, ls=(0, (1, 2)), alpha=0.8) + ideal = np.array([1, tmax]) + ax_s.plot(ideal, ideal, color="0.6", lw=1, zorder=1) + ax_s.annotate("ideal", (tmax, tmax), fontsize=8, color="0.4", + ha="right", va="bottom") + ax_s.annotate("dotted: ceiling with only dS+dR threaded", + (0.03, 0.97), xycoords="axes fraction", fontsize=7, + color="0.35", va="top") + ax_s.set_xscale("log", base=2) + ax_s.set_yscale("log", base=2) + ax_s.set_xlabel("threads") + ax_s.set_ylabel(r"speedup vs serial $\mathtt{danalyticalf}$") + ax_s.set_title("end-to-end gradient speedup", fontsize=10) + ax_s.legend(fontsize=7, loc="lower right", framealpha=0.9) + + # ── panel 2: efficiency ───────────────────────────────────────────────── + for i, tag in enumerate(syss): + th, _, sp1 = thread_curve(sc[tag]) + ax_e.plot(th, sp1 / th, ms=4, label=tag, **STYLES[i % len(STYLES)]) + ax_e.axhline(1.0, color="0.6", lw=1, zorder=1) + ax_e.set_xscale("log", base=2) + ax_e.set_ylim(0, 1.15) + ax_e.set_xlabel("threads") + ax_e.set_ylabel(r"efficiency (speedup vs par@1) / threads") + ax_e.set_title("parallel efficiency", fontsize=10) + + # ── panel 3: where the serial time goes ───────────────────────────────── + bsys = [t for t in syss if t in bd] + x = np.arange(len(bsys)) + bottom = np.zeros(len(bsys)) + for key, label, colour in TERMS: + frac = np.array([100.0 * bd[t][key] / bd[t]["accounted"] for t in bsys]) + ax_b.bar(x, frac, 0.6, bottom=bottom, color=colour, label=label, + hatch="//" if key in WAS_SERIAL else None, + edgecolor="white", lw=0.4) + for xi, (f, b) in enumerate(zip(frac, bottom)): + if f >= 4.0: + ax_b.text(xi, b + f / 2, f"{f:.0f}%", ha="center", va="center", + fontsize=7, color="white", fontweight="bold") + bottom += frac + ax_b.set_xticks(x) + ax_b.set_xticklabels([t.replace("/", "\n") for t in bsys], fontsize=7) + ax_b.set_ylabel("share of 1-core gradient time (%)") + ax_b.set_title("cost breakdown (hatched = was serial)", fontsize=10) + # outside the axes: the bars fill 0-100% by construction, so any in-axes + # legend covers data -- it was hiding the dR labels on the first two bars + ax_b.legend(fontsize=7, loc="upper left", bbox_to_anchor=(1.01, 1.0), + framealpha=0.9, borderaxespad=0) + ax_b.set_ylim(0, 100) + + for ax in (ax_s, ax_e): + ax.grid(True, which="both", lw=0.3, alpha=0.4) + ax.xaxis.set_major_locator(FixedLocator([1, 2, 4, 8, 16, 32, 64])) + ax.xaxis.set_major_formatter(FixedFormatter( + ["1", "2", "4", "8", "16", "32", "64"])) + ax.xaxis.set_minor_locator(NullLocator()) + + sub = (f"{meta.get('ncores', '?')} cores, {meta.get('cpu_model', '')}" + f" | median of {meta.get('repeats', '?')}") + fig.suptitle("Basis-parameter gradient of frozen-P HF energy: thread " + f"scaling\n{sub}", fontsize=10) + fig.tight_layout(rect=(0, 0, 1, 0.90)) + for ext in ("pdf", "png"): + fig.savefig(f"{args.out}.{ext}", dpi=300) + print(f"wrote {args.out}.pdf/.png") + + # ── text summary ──────────────────────────────────────────────────────── + if meta: + print(f"\nrun: {meta.get('node', '?')} {meta.get('ncores', '?')} cores" + f" {meta.get('cpu_model', '')}") + print(f"\n{'system':18s} {'nao':>4s} {'serial':>9s} {'best':>9s} " + f"{'threads':>7s} {'speedup':>8s} {'eff':>5s} {'old cap':>8s}") + for tag in syss: + th, sp, sp1 = thread_curve(sc[tag]) + k = int(np.argmax(sp)) + cap = (bd.get(tag) or {}).get("old_ceiling") + print(f"{tag:18s} {sc[tag]['nao']:4d} {sc[tag]['serial']:9.3f} " + f"{sc[tag]['threads'][str(th[k])]['median']:9.3f} {th[k]:7d} " + f"{sp[k]:7.1f}x {sp1[k] / th[k]:5.2f} " + f"{(f'{cap:6.1f}x' if cap else '-'):>8s}") + + +if __name__ == "__main__": + main() From 58ce1a86e2b5bca86b3bfc7b121411ddf9bf30f8 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Thu, 30 Jul 2026 13:08:30 -0400 Subject: [PATCH 07/11] bench: give librint a threaded column, measured at a fixed core count bench_fair only ran librint pinned to one core, because until now that was the only thing it could do. jax got both a pinned and a free-threaded number, so the comparison had a serial engine on one side and a parallel one on the other. - spawn() exports RAYON_NUM_THREADS in both modes, set from the same core count the jax thread knobs get, so the two engines are handed the same machine rather than different ones. - T1 librint "free" runs danalytical_par on that pool; "pin" stays on the serial danalyticalf, so the pinned column remains the FD-validated path. T2 stays jax-only for free: its librint driver still runs a serial SCF, and timing it "free" would report a speedup no threading produced. pin and free bracket the ladder at one core and the whole node, so the shape between them was never measured -- and "free" is whatever the machine happened to have, which the next run on a wider node cannot be compared against. A fixed width can, so there is now a third point at 32 cores. The width narrows sched_setaffinity as well as RAYON_NUM_THREADS, and the worker's own ncores reports it, so the figure labels the curve from the run rather than from a constant. Cores, not cpus. These nodes are Sockets=1 CoresPerSocket=48 ThreadsPerCore=2 -- 48 cores presented as 96 cpus -- and affinity took the first N entries of the mask, where siblings are numbered adjacently (cpu0 and cpu1 share a core). So the "32 core" runs were 32 threads on 16 cores and the SMT gain was being reported as core scaling. _cpu_order now walks distinct thread_siblings_list first, so the first 32 entries are 32 separate cores, and results carry both ncores (distinct physical cores) and ncpus (hardware threads), because on an SMT node only one of those is what "32 cores" means. jax is swept at the fixed width too -- comparing librint@32c against jax on the whole node was not a comparison at equal cores -- though it is not swept further: it is the memory-bound engine here and every extra width is another chance to OOM. XLA_FLAGS tokens that do not start with "--" are treated as files to read flags from, and XLA aborts (SIGABRT) when it cannot open one, so a bare intra_op_parallelism_threads=32 killed every jax run at a fixed width. It had survived in the pin string only because a real "--" flag leads it -- which means it was being ignored there too, and affinity was doing the clamping all along. So a fixed width sets no XLA_FLAGS at all: jax's CPU backend sizes its Eigen pool from sched_getaffinity, which is already narrowed to the width. The flag stays in the pin string, where single-core has to mean single-core for XLA too. .gitignore picks up the cached density matrices bench_fair writes between processes, and core dumps -- a jax worker that aborts dumps its entire address space next to the results, and the XLA flag crash left 21G of them in the worktree. Co-Authored-By: Claude Opus 5 --- .gitignore | 5 + python/pyscf_comp/bench_fair.py | 159 +++++++++++++++++++++++++++----- 2 files changed, 140 insertions(+), 24 deletions(-) diff --git a/.gitignore b/.gitignore index 4c0497a..b7c9aef 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,11 @@ python/pyscf_comp/*.json python/pyscf_comp/*.pdf python/pyscf_comp/*.png python/pyscf_comp/*_run1/ +# converged density matrices cached between benchmark processes +python/pyscf_comp/*.npy # Stale binary backups *.pre-*-bak + +# core dumps (a crashed jax worker dumps its whole address space -- 14G each) +core.* diff --git a/python/pyscf_comp/bench_fair.py b/python/pyscf_comp/bench_fair.py index 59db6c3..21b5b52 100644 --- a/python/pyscf_comp/bench_fair.py +++ b/python/pyscf_comp/bench_fair.py @@ -212,8 +212,21 @@ def worker_t1_librint(geo, basis, out): # dscf.denergyf is not a second method to time: denergy_c routes to # denergyfast, which assembles the same dHcoreg/dRg/dSg expression as # danalyticalg and returns bitwise-identical values. - g = librint.dscf.danalyticalf(mol, P) - ts = _time_n(lambda: librint.dscf.danalyticalf(mol, P), 3) + # + # "free" and any explicit width ("32") run the threaded assembly + # (src/par.rs) on rayon's global pool, which spawn() sized with + # RAYON_NUM_THREADS -- the same core count jax gets from its own thread + # knobs. Passing 0 uses that global pool rather than building a private + # one, so the width comes from the environment in every threaded mode. + # "pin" stays on the serial entry point, so the pin column remains exactly + # the FD-validated reference path. + if THREAD_MODE != "pin": + run = lambda: librint.dscf.danalytical_par(mol, P, 0) + else: + run = lambda: librint.dscf.danalyticalf(mol, P) + + g = run() + ts = _time_n(run, 3) out.update(median=_median(ts), grad_sorted=np.sort(g).tolist()) out["peak_kb"] = _vmhwm_kb() @@ -328,6 +341,17 @@ def run(): out.update(cold=cold, median=_median(ts), grad_sorted=np.sort(g).tolist()) +# "pin" (one core), "free" (the whole allocation), or a decimal width like +# "32"; set by run_worker before dispatch so a worker can pick the matching +# code path. +THREAD_MODE = "pin" + +# Fixed librint widths measured alongside pin and free. pin and free bracket +# the ladder at 1 core and the whole node, which leaves everything between them +# unmeasured -- and the node width varies by machine, so "free" alone is not a +# number you can compare across runs. +THREAD_WIDTHS = ("32",) + WORKERS = { "t0": worker_t0, "t1_librint": worker_t1_librint, @@ -345,17 +369,56 @@ def worker_setup_P(geo, basis, out_path): np.save(out_path, _converged_P(mol)) +def _siblings(cpu): + p = f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list" + try: + with open(p) as f: + return f.read().strip() + except OSError: + return str(cpu) # no topology exposed -> assume 1:1 + + +def _cpu_order(mask): + """Allocation's cpus, one per PHYSICAL core first, SMT siblings after. + + These nodes are 48 cores x 2 threads presented as 96 cpus, so taking the + first N of the raw mask can put a "32 core" run on 16 cores' worth of + silicon and call the SMT gain a scaling result. Ordering by distinct + thread_siblings_list makes the first 32 entries 32 separate cores. + """ + seen, first, extra = set(), [], [] + for c in sorted(mask): + sib = _siblings(c) + if sib in seen: + extra.append(c) + else: + seen.add(sib) + first.append(c) + return first + extra + + +def _nphys(mask): + """Distinct physical cores covered by a cpu mask.""" + return len({_siblings(c) for c in mask}) + + def run_worker(argv): kind, geo, basis, arg4 = argv if kind == "setup_P": worker_setup_P(geo, basis, arg4) # arg4 = output npy path return threads = arg4 - if threads != "free": + global THREAD_MODE + THREAD_MODE = threads + # "pin" = one core, "free" = the whole allocation, "" = the first N + # cores of it. Narrowing affinity rather than only capping + # RAYON_NUM_THREADS is what makes out["ncores"] below report the width the + # run actually had, and stops BLAS/OMP spilling past it. + n = None if threads == "free" else (1 if threads == "pin" else int(threads)) + if n is not None: try: - # first core of OUR allocation (works inside SLURM cgroups too) - core = sorted(os.sched_getaffinity(0))[0] - os.sched_setaffinity(0, {core}) + # cpus of OUR allocation (works inside SLURM cgroups too) + os.sched_setaffinity(0, set(_cpu_order(os.sched_getaffinity(0))[:n])) except OSError: pass state = {"peak_kb": -1} @@ -363,10 +426,14 @@ def run_worker(argv): out = {} WORKERS[kind](geo, basis, out) out["peak_kb"] = max(state["peak_kb"], _vmhwm_kb()) - # measured after the affinity call, so this is the core count the run - # really had -- the figure labels its curves from this, never from a - # hardcoded number - out["ncores"] = len(os.sched_getaffinity(0)) + # measured after the affinity call, so these describe the run it really + # got -- the figure labels its curves from them, never from a hardcoded + # number. ncpus counts hardware threads, ncores counts distinct physical + # cores; on an SMT node those differ by 2x and only one of them is what + # "32 cores" means. + mask = os.sched_getaffinity(0) + out["ncpus"] = len(mask) + out["ncores"] = _nphys(mask) print("BENCH_JSON " + json.dumps(out), flush=True) @@ -400,13 +467,31 @@ def spawn(kind, geo, basis, threads, timeout=900, p_file=None): env["PYTHONHASHSEED"] = "0" if p_file: # T1 workers load this P instead of running incore SCF env["LIBRINT_P_FILE"] = p_file - nthr = "1" if threads != "free" else str(os.cpu_count()) + if threads == "free": + nthr = str(os.cpu_count()) + elif threads == "pin": + nthr = "1" + else: # explicit width, e.g. "32" + nthr = str(int(threads)) + # rayon sizes librint's global pool from this; set in every mode so the + # two engines are given the same core count rather than different ones. + env["RAYON_NUM_THREADS"] = nthr if threads != "free": for var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS", "NUMEXPR_NUM_THREADS"): env[var] = nthr - env["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false " - "intra_op_parallelism_threads=1") + if threads == "pin": + # single-core means single-core for XLA too + env["XLA_FLAGS"] = ("--xla_cpu_multi_thread_eigen=false " + "intra_op_parallelism_threads=1") + # No XLA_FLAGS for a fixed width: XLA's parser treats any token that + # does not start with "--" as a FILE to read flags from and aborts + # (SIGABRT) when it cannot open it, so a bare + # "intra_op_parallelism_threads=32" kills the process. It is tolerated + # in the pin string above only because a real "--" flag leads -- which + # means it is being ignored there too. What actually bounds jax is the + # affinity mask: its CPU backend sizes the Eigen pool from + # sched_getaffinity, so pinning to N cpus gives N intra-op threads. t0 = time.perf_counter() try: proc = subprocess.run( @@ -484,7 +569,8 @@ def main(): # provenance of THIS run, so plots read the node's real core count and # memory ceiling out of the results instead of assuming them results = {"_meta": {"node": platform.node(), - "ncores": len(os.sched_getaffinity(0)), + "ncores": _nphys(os.sched_getaffinity(0)), + "ncpus": len(os.sched_getaffinity(0)), "mem_limit_kb": _mem_limit_kb(), "cpu_model": _cpu_model()}} @@ -513,11 +599,26 @@ def key(*parts): results[key(tier, eng, tag, "pin")] = r print(f" {tier} {eng:8s} pin : {fmt_t(r)}s peak {fmt_mem(r)}", flush=True) - rf = spawn(f"{tier}_jax", geo, basis, "free", timeout=args.timeout, - p_file=pf_t) - results[key(tier, "jax", tag, "free")] = rf - print(f" {tier} jax free: {fmt_t(rf)}s peak {fmt_mem(rf)}", - flush=True) + # librint only has a threaded path for the T1 gradient assembly + # (src/par.rs); its T2 driver still runs a serial SCF, so timing it + # "free" would report a number no threading produced. + free_engines = ("librint", "jax") if tier == "t1" else ("jax",) + for eng in free_engines: + rf = spawn(f"{tier}_{eng}", geo, basis, "free", + timeout=args.timeout, p_file=pf_t) + results[key(tier, eng, tag, "free")] = rf + print(f" {tier} {eng:8s} free: {fmt_t(rf)}s peak {fmt_mem(rf)}", + flush=True) + # BOTH engines at each fixed width: the figure compares at equal + # core counts, so sweeping only librint would leave jax to be + # judged at whatever width the node happened to have. + for w in (THREAD_WIDTHS if tier == "t1" else ()): + for eng in ("librint", "jax"): + rw = spawn(f"{tier}_{eng}", geo, basis, w, + timeout=args.timeout, p_file=pf_t) + results[key(tier, eng, tag, w)] = rw + print(f" {tier} {eng:8s} {w:>4s}: {fmt_t(rw)}s " + f"peak {fmt_mem(rw)}", flush=True) if pf and os.path.exists(pf): os.remove(pf) with open(out_json, "w") as f: # write-as-you-go: survive job timeouts @@ -543,19 +644,29 @@ def key(*parts): ("t2", f"end-to-end SCF+gradient (conv_tol={SCF_CONV})")): if tier not in tiers: continue - print(f"\n=== {tier.upper()} {desc}; median of 3; librint = danalyticalf ===") - print(f"{'system':16s} {'librint@1c':>11s} {'jax@1c':>9s} " - f"{'jax warm1':>9s} {'jax@free':>9s} {'lib peak':>9s} {'jax peak':>9s} " - f"{'max|Δg|':>9s}") + print(f"\n=== {tier.upper()} {desc}; median of 3 ===") + print(" @1c = danalyticalf (serial, FD-validated); @Nc and @free = " + "danalytical_par on N cores / the whole allocation") + wcols = "".join(f"{'librint@' + w + 'c':>13s}" for w in THREAD_WIDTHS) + print(f"{'system':16s} {'librint@1c':>11s}{wcols} {'librint@free':>13s} " + f"{'jax@1c':>9s} {'jax warm1':>9s} {'jax@free':>9s} " + f"{'lib peak':>9s} {'jax peak':>9s} {'max|Δg|':>9s}") for basis, geo in mols: tag = f"{geo}/{basis}" rl = results.get(key(tier, "librint", tag, "pin")) if rl is None: # T1_ONLY system continue + rlf = results.get(key(tier, "librint", tag, "free")) rj = results[key(tier, "jax", tag, "pin")] rf = results[key(tier, "jax", tag, "free")] cold = fmt_t(rj, "cold") if rj.get("status") == "ok" else "-" - print(f"{tag:16s} {fmt_t(rl):>11s} {fmt_t(rj):>9s} {cold:>9s} " + wvals = "".join( + f"{(fmt_t(rw) if rw else '-'):>13s}" + for rw in (results.get(key(tier, "librint", tag, w)) + for w in THREAD_WIDTHS)) + print(f"{tag:16s} {fmt_t(rl):>11s}{wvals} " + f"{(fmt_t(rlf) if rlf else '-'):>13s} " + f"{fmt_t(rj):>9s} {cold:>9s} " f"{fmt_t(rf):>9s} {fmt_mem(rl):>9s} {fmt_mem(rj):>9s} " f"{grad_err(rl, rj):>9s}") From bface5efd8a453ae5f10af430a9346e37136f7e4 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Thu, 30 Jul 2026 15:12:31 -0400 Subject: [PATCH 08/11] bench: draw the alkane figure at equal cores, and get out of the data's way plot_alkanes gains a ("librint", "free") series and a column for it, and the fixed-width curve is labelled from the run's own ncores rather than from a constant, so a figure drawn on a wider node still says what it measured. Both whole-allocation series then come back out. A "free" run is 96 threads on 48 cores, so its speedup mixes core scaling with SMT, and there is no honest place for it next to a per-core curve. Runs made before the threaded wiring simply have no points in the new series, which is why it is defined but empty on the existing alkane JSON. Two layout fixes. Five series on log-log axes leave no free corner: the upper-left legend sat on the jax curves and the job-limit annotation ran through librint's memory curves in the TZVP panel, so the legend goes under the figure and the annotation into the band between librint's flat footprint and jax's. And get_legend_handles_labels() returns (handles, labels) -- unpacking it into `labels` clobbered the per-series label dict the summary table indexes with (engine, thread-mode) tuples. The figure was written first, so the run looked fine until the table raised TypeError. Co-Authored-By: Claude Opus 5 --- python/pyscf_comp/plot_alkanes.py | 51 ++++++++++++++++++++++--------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/python/pyscf_comp/plot_alkanes.py b/python/pyscf_comp/plot_alkanes.py index b6a9200..d41d7cb 100644 --- a/python/pyscf_comp/plot_alkanes.py +++ b/python/pyscf_comp/plot_alkanes.py @@ -57,10 +57,16 @@ def classify(r): } CLABEL = {"CH4": "C1", "C2H6": "C2", "C3H8": "C3", "C4H10": "C4", "C6H6": "C6"} +# 1 core and a fixed 32 for both engines. The "free" (whole-allocation) series +# are measured but not drawn: these nodes are 48 cores x 2 SMT threads, so a +# free run is 96 threads on 48 cores and its speedup mixes core scaling with +# SMT gain. A fixed width, pinned one thread per physical core, is the number +# that means what the axis label says -- and it stays comparable across nodes. SERIES = [ # key-parts, style; legend labels are built from the run's own data (("librint", "pin"), dict(color="#1a7f37", marker="o", ls="-")), + (("librint", "32"), dict(color="#57ab5a", marker="D", ls="-.")), (("jax", "pin"), dict(color="#cf222e", marker="s", ls="--")), - (("jax", "free"), dict(color="#e16f24", marker="^", ls=":")), + (("jax", "32"), dict(color="#e16f24", marker="^", ls=":")), ] @@ -83,7 +89,9 @@ def cores_of(results, eng, threads): if n: return None # inconsistent -> say so meta = results.get("_meta") or {} # pre-`ncores` results - return 1 if threads != "free" else meta.get("ncores") + if threads == "free": + return meta.get("ncores") + return int(threads) if threads.isdigit() else 1 def label_of(results, eng, threads): @@ -176,10 +184,12 @@ def main(): top.tick_params(length=0, labelsize=8) if any_oom: ax_m.axhline(limit_g, color="0.5", lw=0.8, ls="--") - # low-left: the only corner both engines' curves stay out of + # the band between librint's flat ~0.2-0.7G and jax's 6G+ is the + # only stripe both engines leave empty; the low corner does not + # work, librint's curves run through it ax_m.annotate(f"job limit {limit_g:.0f}G (open = out of memory)", - (0.02, 0.22), xycoords="axes fraction", - fontsize=7, color="0.35", va="top") + (0.02, 0.46), xycoords="axes fraction", + fontsize=7, color="0.35", va="center") ax_m.set_xlabel("cartesian basis functions") axes[0][0].set_ylabel("wall time per gradient (s)") axes[1][0].set_ylabel("peak RSS (GiB)") @@ -191,10 +201,16 @@ def main(): xj, tj, mj, _, _ = collect(results, "def2-tzvp", "jax", "pin", with_bz) if mj.size >= 2: guide(axes[1][1], xj[0], mj[0] * 0.5, 4) - axes[0][0].legend(fontsize=8, loc="upper left", framealpha=0.9) + # one row under the whole figure. Five series on log-log axes leave no + # in-axes corner free -- an upper-left legend sat on top of the jax curves + # NOT `labels` -- that name holds the per-series label dict used below + handles, leg_labels = axes[0][0].get_legend_handles_labels() + fig.legend(handles, leg_labels, loc="lower center", ncol=len(handles), + fontsize=8, frameon=False, bbox_to_anchor=(0.5, 0.0), + columnspacing=1.4, handlelength=2.4) fig.suptitle("Basis-parameter gradient of frozen-P HF energy " r"(T1): $\mathrm{C}_n\mathrm{H}_{2n+2}$ ladder", fontsize=11) - fig.tight_layout(rect=(0, 0, 1, 0.97)) + fig.tight_layout(rect=(0, 0.05, 1, 0.97)) for ext in ("pdf", "png"): fig.savefig(f"{args.out}.{ext}", dpi=300) print(f"wrote {args.out}.pdf/.png") @@ -202,19 +218,23 @@ def main(): # text summary + gradient agreement meta = results.get("_meta") or {} if meta: + smt = "" + if meta.get("ncpus") and meta.get("ncores"): + smt = f" ({meta['ncpus']} hw threads)" print(f"\nrun: {meta.get('node', '?')} {meta.get('ncores', '?')} cores" - f" {(meta.get('mem_limit_kb') or 0) / 1048576:.0f}G" + f"{smt} {(meta.get('mem_limit_kb') or 0) / 1048576:.0f}G" f" {meta.get('cpu_model', '')}") - hdr = [labels[("librint", "pin")], labels[("jax", "pin")], - labels[("jax", "free")]] - print(f"\n{'system':18s} {hdr[0]:>16s} {hdr[1]:>16s} " - f"{hdr[2]:>17s} {'lib mem':>8s} {'jax mem':>8s} {'max|dg|':>9s}") + hdr = [labels[("librint", "pin")], labels[("librint", "32")], + labels[("jax", "pin")], labels[("jax", "32")]] + print(f"\n{'system':18s} {hdr[0]:>16s} {hdr[1]:>17s} {hdr[2]:>16s} " + f"{hdr[3]:>17s} {'lib mem':>8s} {'jax mem':>8s} {'max|dg|':>9s}") for basis in BASES: for geo in SYSTEMS + (["C6H6"] if basis == "def2-tzvp" and with_bz else []): rl = entry(results, "librint", geo, basis, "pin") + rl32 = entry(results, "librint", geo, basis, "32") rj = entry(results, "jax", geo, basis, "pin") - rf = entry(results, "jax", geo, basis, "free") + rj32 = entry(results, "jax", geo, basis, "32") def t(r): if r.get("status") == "ok": @@ -236,8 +256,9 @@ def m(r): b = np.array(rj["grad_sorted"]) dg = (f"{np.abs(a - b).max():.1e}" if a.shape == b.shape else "len!") - print(f"{geo + '/' + basis:18s} {t(rl):>16s} {t(rj):>16s} " - f"{t(rf):>17s} {m(rl):>8s} {m(rj):>8s} {dg:>9s}") + print(f"{geo + '/' + basis:18s} {t(rl):>16s} {t(rl32):>17s} " + f"{t(rj):>16s} {t(rj32):>17s} " + f"{m(rl):>8s} {m(rj):>8s} {dg:>9s}") if __name__ == "__main__": From f0b321a54a61ff5fdc7659b452572e61fb617391 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Sun, 2 Aug 2026 16:21:32 -0400 Subject: [PATCH 09/11] bench: --only/--out, so one system can have one whole node Every one of these scripts is a timing measurement, so each system needs a quiet exclusive node -- but no system's timing depends on any other's. Run serially that is hours on one machine; sharded it is the slowest single system. bench_fair's alkane ladder is the extreme case: the big end is minutes per evaluation on one core, and the harness does four evaluations per point. --only selects systems by "GEO/BASIS" and --out redirects the results file, which is all that is needed to put each on its own node and merge the JSONs afterwards -- they are keyed by system, so the merge is a dict update. Unknown names are an error rather than a silent empty run, because a typo that quietly measures nothing looks exactly like a system that finished fast. Defaults are unchanged: no --only still runs the full list to the same file. One caveat belongs with the caller, not the code: sharded runs must land on one CPU model. These figures plot systems against each other, and this cluster's cpu pool mixes EPYC 9J14 and 7J13, which are ~1.3x apart. --- python/pyscf_comp/bench_fair.py | 14 ++++++++++ python/pyscf_comp/bench_grad_breakdown.py | 30 ++++++++++++++++++--- python/pyscf_comp/bench_par_scaling.py | 33 ++++++++++++++++++++--- 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/python/pyscf_comp/bench_fair.py b/python/pyscf_comp/bench_fair.py index 21b5b52..2e0dbef 100644 --- a/python/pyscf_comp/bench_fair.py +++ b/python/pyscf_comp/bench_fair.py @@ -553,6 +553,13 @@ def main(): ap.add_argument("--suite", choices=["default", "alkanes"], default="default") ap.add_argument("--timeout", type=int, default=900, help="per-worker wall limit (s)") + # one molecule per exclusive node: the ladder is embarrassingly parallel + # ACROSS molecules but must be quiet WITHIN one, and the big end of the + # ladder is minutes per point on a single core + ap.add_argument("--only", action="append", metavar="GEO/BASIS", + help="run only these molecules (repeatable, or comma-" + "separated)") + ap.add_argument("--out", help="results file (default: per-suite name)") args = ap.parse_args() if args.worker: run_worker(args.worker) @@ -566,6 +573,13 @@ def main(): mols = [m for m in MOLECULES if not args.quick or m[0] == "sto-3g"] tiers = ("t1", "t2") out_json = "bench_fair_results.json" + if args.only: + wanted = [s.strip() for a in args.only for s in a.split(",") if s.strip()] + mols = [m for m in mols if f"{m[1]}/{m[0]}" in wanted] + missing = set(wanted) - {f"{m[1]}/{m[0]}" for m in mols} + if missing: + raise SystemExit(f"unknown molecule(s): {', '.join(sorted(missing))}") + out_json = args.out or out_json # provenance of THIS run, so plots read the node's real core count and # memory ceiling out of the results instead of assuming them results = {"_meta": {"node": platform.node(), diff --git a/python/pyscf_comp/bench_grad_breakdown.py b/python/pyscf_comp/bench_grad_breakdown.py index 2f4f724..3c327d4 100644 --- a/python/pyscf_comp/bench_grad_breakdown.py +++ b/python/pyscf_comp/bench_grad_breakdown.py @@ -21,7 +21,12 @@ forms P F P. Usage: LIBRINT_SO=/path/to/librint.so python bench_grad_breakdown.py + ... python bench_grad_breakdown.py --only C2H6/def2-tzvp --out shard.json + +--only/--out put one system on one node so the four can be measured at once; +the per-system JSONs merge cleanly because each is keyed by "GEO/BASIS". """ +import argparse import ctypes import json import os @@ -93,12 +98,31 @@ def timed(label, fn): return dt, out +def select(only): + """SYSTEMS, or the subset named as GEO/BASIS on the command line.""" + if not only: + return SYSTEMS + wanted = [s.strip() for arg in only for s in arg.split(",") if s.strip()] + chosen = [(g, b) for g, b in SYSTEMS if f"{g}/{b}" in wanted] + missing = set(wanted) - {f"{g}/{b}" for g, b in chosen} + if missing: + raise SystemExit(f"unknown system(s): {', '.join(sorted(missing))}") + return chosen + + def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--only", action="append", metavar="GEO/BASIS", + help="measure only these systems (repeatable, or comma-" + "separated)") + ap.add_argument("--out", default=OUT_JSON, help="results file") + args = ap.parse_args() + results = {"_meta": {"node": platform.node(), "ncores": len(os.sched_getaffinity(0)), "mem_limit_kb": _mem_limit_kb(), "cpu_model": _cpu_model()}} - for geo, basis in SYSTEMS: + for geo, basis in select(args.only): mol = build(geo, basis) mf = pyscf.scf.RHF(mol) mf.verbose = 0 @@ -155,10 +179,10 @@ def main(): "old_serial_frac": was_serial / acct, "old_ceiling": acct / was_serial, } - with open(OUT_JSON, "w") as f: + with open(args.out, "w") as f: json.dump(results, f, indent=1) - print(f"\nresults -> {OUT_JSON}") + print(f"\nresults -> {args.out}") if __name__ == "__main__": diff --git a/python/pyscf_comp/bench_par_scaling.py b/python/pyscf_comp/bench_par_scaling.py index c635bd3..345eba8 100644 --- a/python/pyscf_comp/bench_par_scaling.py +++ b/python/pyscf_comp/bench_par_scaling.py @@ -15,8 +15,15 @@ Writes bench_scaling_results.json so the figure is drawn from measurements rather than from a log scrape, with a _meta block recording the machine. +Every system needs a quiet exclusive node -- a co-tenant reads as poor thread +scaling -- so --only/--out exist to put one system on one node and merge the +per-system JSONs afterwards. Nodes must then be the same CPU model, since the +figure plots systems against each other. + Usage: LIBRINT_SO=/path/to/librint.so python bench_par_scaling.py + ... python bench_par_scaling.py --only C2H6/def2-tzvp --out shard.json """ +import argparse import json import os import platform @@ -62,7 +69,27 @@ def median_time(fn, n): return float(np.median(ts)), out +def select(only): + """SYSTEMS, or the subset named as GEO/BASIS on the command line.""" + if not only: + return SYSTEMS + wanted = [s.strip() for arg in only for s in arg.split(",") if s.strip()] + chosen = [(g, b) for g, b in SYSTEMS if f"{g}/{b}" in wanted] + missing = set(wanted) - {f"{g}/{b}" for g, b in chosen} + if missing: + raise SystemExit(f"unknown system(s): {', '.join(sorted(missing))}") + return chosen + + def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--only", action="append", metavar="GEO/BASIS", + help="measure only these systems (repeatable, or comma-" + "separated); one shard per exclusive node") + ap.add_argument("--out", default=OUT_JSON, help="results file") + args = ap.parse_args() + systems = select(args.only) + ncores = len(os.sched_getaffinity(0)) threads = [t for t in THREADS if t <= ncores] print(f"usable cores: {ncores} thread counts: {threads} " @@ -72,7 +99,7 @@ def main(): "mem_limit_kb": _mem_limit_kb(), "cpu_model": _cpu_model(), "repeats": REPEATS}} failures = 0 - for geo, basis in SYSTEMS: + for geo, basis in systems: mol = build(geo, basis) mf = pyscf.scf.RHF(mol) mf.verbose = 0 @@ -106,10 +133,10 @@ def main(): f"rel|par-ser|={err:.2e}" f"{'' if err < 1e-9 else ' <-- MISMATCH'}", flush=True) results[tag] = row - with open(OUT_JSON, "w") as f: # write-as-you-go, survive a timeout + with open(args.out, "w") as f: # write-as-you-go, survive a timeout json.dump(results, f, indent=1) - print(f"\nresults -> {OUT_JSON}") + print(f"\nresults -> {args.out}") if failures: raise SystemExit(f"{failures} timed result(s) disagreed with serial") From 14c7b46aa94f50fd15c9c18877a8d686ed4ac382 Mon Sep 17 00:00:00 2001 From: AbAldo Date: Sun, 2 Aug 2026 17:03:08 -0400 Subject: [PATCH 10/11] bench: benzene is a demo, not a timing point jax has never completed C6H6/def2-tzvp -- OOM_ALLOC on every attempt on a 181 GB node, in the archived ladder and again today. It contributes no comparison, only a librint-only bar. What it does contribute is wall time: the serial baseline is ~7 minutes per evaluation and the harness does four of them, which roughly doubles whichever suite it sits in. So exclude it from both suites by default. Naming it explicitly with --only still runs it, which is how the standalone demo gets measured -- the point there is not that librint is faster but that it finishes at 0.1 GB where jax cannot start. jax's memory wall is between C2H6/def2-tzvp, which peaks at 90.7 GB and just fits, and C3H8/def2-tzvp, which does not. Both of those stay in the suite; the two rungs above them are librint-only either way. --- python/pyscf_comp/bench_fair.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/pyscf_comp/bench_fair.py b/python/pyscf_comp/bench_fair.py index 2e0dbef..0597b2b 100644 --- a/python/pyscf_comp/bench_fair.py +++ b/python/pyscf_comp/bench_fair.py @@ -58,6 +58,13 @@ ("def2-tzvp", "C4H10"), ("def2-tzvp", "C6H6"), # jax OOMs; librint gradient measured with P isolated ] +# Benzene is a demonstration, not a data point. jax cannot run it at all -- it +# has OOM_ALLOC'd on every attempt, on a 181 GB node -- so it contributes no +# comparison, only a librint-only bar and a ~7 minute serial baseline that +# doubles the wall time of whichever suite it is in. Excluded from both suites +# by default; naming it explicitly with --only still runs it, which is how the +# demo figure gets made. +DEMO_ONLY = {"C6H6/def2-tzvp"} SCF_CONV = 1e-8 SCF_MAXITER = 4000 @@ -579,6 +586,8 @@ def main(): missing = set(wanted) - {f"{m[1]}/{m[0]}" for m in mols} if missing: raise SystemExit(f"unknown molecule(s): {', '.join(sorted(missing))}") + else: + mols = [m for m in mols if f"{m[1]}/{m[0]}" not in DEMO_ONLY] out_json = args.out or out_json # provenance of THIS run, so plots read the node's real core count and # memory ceiling out of the results instead of assuming them From 5f16b7476650f669a1ebd377200972db84c910dc Mon Sep 17 00:00:00 2001 From: AbAldo Date: Sun, 2 Aug 2026 17:05:15 -0400 Subject: [PATCH 11/11] bench: count physical cores in the scaling figure too These nodes are 48 cores with 2-way SMT, so sched_getaffinity returns 96 and reporting that as "cores" claims twice the silicon the run had. bench_fair already goes through _nphys for exactly this reason; bench_par_scaling and bench_grad_breakdown did not, so the thread-scaling figure was captioned "96 cores" on a 48-core machine. Both now record ncores (physical) alongside ncpus (hardware threads), and plot_scaling captions "48 cores (96 hw threads)" the way plot_alkanes does. The sweep itself is unchanged and still runs past the core count -- T=64 is the fastest point on several systems -- because oversubscribing a core with SMT is a legitimate thing to measure. It is reported as threads, which is what the x-axis already called it. --- python/pyscf_comp/bench_grad_breakdown.py | 6 ++++-- python/pyscf_comp/bench_par_scaling.py | 19 +++++++++++++------ python/pyscf_comp/plot_scaling.py | 8 ++++++-- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/python/pyscf_comp/bench_grad_breakdown.py b/python/pyscf_comp/bench_grad_breakdown.py index 3c327d4..448e910 100644 --- a/python/pyscf_comp/bench_grad_breakdown.py +++ b/python/pyscf_comp/bench_grad_breakdown.py @@ -41,7 +41,7 @@ import librint.utils from librint import library -from bench_fair import _cpu_model, _mem_limit_kb +from bench_fair import _cpu_model, _mem_limit_kb, _nphys from geometries import geometries # same list, same order as bench_par_scaling.py, so the two JSONs can be shown @@ -119,7 +119,9 @@ def main(): args = ap.parse_args() results = {"_meta": {"node": platform.node(), - "ncores": len(os.sched_getaffinity(0)), + # physical cores, not the 2x that SMT reports + "ncores": _nphys(os.sched_getaffinity(0)), + "ncpus": len(os.sched_getaffinity(0)), "mem_limit_kb": _mem_limit_kb(), "cpu_model": _cpu_model()}} for geo, basis in select(args.only): diff --git a/python/pyscf_comp/bench_par_scaling.py b/python/pyscf_comp/bench_par_scaling.py index 345eba8..1f4a5dd 100644 --- a/python/pyscf_comp/bench_par_scaling.py +++ b/python/pyscf_comp/bench_par_scaling.py @@ -36,7 +36,7 @@ import librint.dscf import librint.utils -from bench_fair import _cpu_model, _mem_limit_kb +from bench_fair import _cpu_model, _mem_limit_kb, _nphys from geometries import geometries SYSTEMS = [ @@ -90,13 +90,20 @@ def main(): args = ap.parse_args() systems = select(args.only) - ncores = len(os.sched_getaffinity(0)) - threads = [t for t in THREADS if t <= ncores] - print(f"usable cores: {ncores} thread counts: {threads} " - f"median of {REPEATS}", flush=True) + # These nodes are 48 physical cores with 2-way SMT, so sched_getaffinity + # returns 96 and calling that "cores" overstates the machine by 2x -- the + # same mislabel _nphys already fixes in bench_fair. The sweep still runs + # past the core count (T=64 is the fastest point on some systems), it is + # just reported as threads, which is what it is. + mask = os.sched_getaffinity(0) + ncpus = len(mask) + ncores = _nphys(mask) + threads = [t for t in THREADS if t <= ncpus] + print(f"physical cores: {ncores} hw threads: {ncpus} " + f"thread counts: {threads} median of {REPEATS}", flush=True) results = {"_meta": {"node": platform.node(), "ncores": ncores, - "mem_limit_kb": _mem_limit_kb(), + "ncpus": ncpus, "mem_limit_kb": _mem_limit_kb(), "cpu_model": _cpu_model(), "repeats": REPEATS}} failures = 0 for geo, basis in systems: diff --git a/python/pyscf_comp/plot_scaling.py b/python/pyscf_comp/plot_scaling.py index 964508f..bf39774 100644 --- a/python/pyscf_comp/plot_scaling.py +++ b/python/pyscf_comp/plot_scaling.py @@ -145,7 +145,11 @@ def main(): ["1", "2", "4", "8", "16", "32", "64"])) ax.xaxis.set_minor_locator(NullLocator()) - sub = (f"{meta.get('ncores', '?')} cores, {meta.get('cpu_model', '')}" + # "cores" must mean physical cores: these nodes are 48 with 2-way SMT, so + # the thread counts past 48 are oversubscribing a core, not using a new one + smt = (f" ({meta['ncpus']} hw threads)" + if meta.get("ncpus") and meta.get("ncores") else "") + sub = (f"{meta.get('ncores', '?')} cores{smt}, {meta.get('cpu_model', '')}" f" | median of {meta.get('repeats', '?')}") fig.suptitle("Basis-parameter gradient of frozen-P HF energy: thread " f"scaling\n{sub}", fontsize=10) @@ -157,7 +161,7 @@ def main(): # ── text summary ──────────────────────────────────────────────────────── if meta: print(f"\nrun: {meta.get('node', '?')} {meta.get('ncores', '?')} cores" - f" {meta.get('cpu_model', '')}") + f"{smt} {meta.get('cpu_model', '')}") print(f"\n{'system':18s} {'nao':>4s} {'serial':>9s} {'best':>9s} " f"{'threads':>7s} {'speedup':>8s} {'eff':>5s} {'old cap':>8s}") for tag in syss: