Par gradient - #16
Open
ZuseZ4 wants to merge 4 commits into
Open
Conversation
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<f64>, ...) 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Your 2k line PR was too hard to review, so I split it up a little