Par wiring

.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.*

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)