Par wiring - #9
Draft
mindeye33 wants to merge 11 commits into
Draft
Conversation
Member
|
@copilot rebase and force-push |
Done — rebased onto |
Member
|
@copilot that's a lie! This PR still has 4 commits taht are already in main. |
Updated. Branch now includes |
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>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…'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 <noreply@anthropic.com>
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.
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.
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.
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
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.
parallelizes our work of calling enzyme with rayon.