Par gradient

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

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",
+]

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 = (

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)

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)

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}"

src/lib.rs

@@ -3,6 +3,7 @@
 pub mod scf;
 
 pub mod dscf;
+pub mod par;
 pub mod p2c;
 
 pub mod linalg;

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<f64>) -> *mut f64 {
+pub(crate) fn leak_vec(v: Vec<f64>) -> *mut f64 {
     let mut b = v.into_boxed_slice();
     let ptr = b.as_mut_ptr();
     std::mem::forget(b);