Add OrbitalEvaluator and Gaussian cube writer utility - #200
Open
ConradJohnston wants to merge 31 commits into
Open
Add OrbitalEvaluator and Gaussian cube writer utility#200ConradJohnston wants to merge 31 commits into
ConradJohnston wants to merge 31 commits into
Conversation
ConradJohnston
marked this pull request as draft
April 29, 2026 17:40
ConradJohnston
force-pushed
the
feature/cojoh/orbital-evaluator-and-cube
branch
2 times, most recently
from
April 29, 2026 18:06
d7dc190 to
efae7ca
Compare
Adds a public, host-parallel point-set evaluator for molecular orbitals
and densities, plus a self-contained Gaussian cube file writer. The two
are deliberately decoupled: OrbitalEvaluator returns plain numerical
arrays, and write_cube depends only on Molecule + raw field data. Either
can be used independently.
OrbitalEvaluator (include/gauxc/orbital_evaluator.hpp,
src/orbital_evaluator.cxx)
- Constructed from a BasisSet; caches the LocalHostWorkDriver internally
so it can be reused across many evaluations for the same molecule.
- eval_orbital / eval_orbitals: chi_j(r) = sum_mu C[mu, j] * phi_mu(r),
one or many MOs at a time on an arbitrary AoS point set.
- eval_density: rho(r) = sum_{mu,nu} D[mu, nu] * phi_mu(r) * phi_nu(r),
for a symmetric AO density matrix.
- Internally batches points with an adaptive batch size (~16 MB AO
scratch per thread) and parallelises the batch loop with OpenMP. Each
thread owns its own scratch.
- Currently Host-only; the public API takes ExecutionSpace so device
backends can be slotted in later without an ABI break.
write_cube + CubeGrid (include/gauxc/external/cube.hpp,
src/external/cube.cxx)
- Standard Gaussian cube text format, byte-compatible with the output of
PySCF / Gaussian / NWChem cubegen utilities.
- CubeGrid::from_molecule builds a default axis-aligned grid that
encloses the molecule with a configurable margin (default 3.0 Bohr,
matching PySCF cubegen).
- The data block is formatted in parallel: each (ix, iy) row is
serialised independently into a pre-sized byte buffer and committed
with a single fwrite at the end. The per-value formatter is a
hand-rolled %13.5E that matches glibc snprintf bit-for-bit on the
fast path (1-2 digit exponent) and falls back to snprintf for the
rare 3-digit-exponent case. This is the same writer that produced the
measured 5x end-to-end speedup over PySCF cubegen on the QDK-Chemistry
workloads that motivated this PR.
- Self-contained: no third-party dependencies, no GauXC kernel coupling.
Tests (tests/orbital_evaluator_test.cxx, [orbital_evaluator] / [cube]):
- Water cc-pVDZ kernel checks against direct eval_collocation reference:
one-hot MO, multi-MO contraction, identity density, rank-1 density.
- CubeGrid layout: first/last point coordinates, iz-fastest ordering.
- write_cube round-trip: parses back the written file and verifies
header (atoms, origin, axes) and field values to %13.5E precision.
- write_cube formatter: byte-exact comparison of the data block against
glibc snprintf("%13.5E", v) for a battery of edge-case values
(signed zero, near-power-of-ten boundaries, 3-digit exponents).
- All 1570 assertions pass; existing [collocation] suite unaffected.
Use gen_compressed_submat_map + Shell::cutoff_radius() to skip shells whose cutoff sphere doesn't overlap any grid point in the batch. Reduces collocation work for sparse systems (e.g. water at 80^3/16t: 10.4x -> 11.2x vs PySCF).
Add eval_orbital/eval_orbitals/eval_density overloads that accept a CubeGrid directly and generate per-batch point coordinates on-the-fly, avoiding the 3*N*8-byte temporary coordinate array (~12 MB at 80^3, ~98 MB at 160^3). Also add CubeGrid::points_into(double*) for callers that manage their own buffer. Benchmark improvement (vs shell, 16 threads): water 160^3 orb: 11.9x -> 12.5x benzene 160^3 orb: 2.9x -> 3.2x benzene 160^3 den: 3.1x -> 3.3x
Replace the omp-for loop in the CubeGrid eval_orbitals overload with a hybrid dispatch: when n_batches < 2*nthreads, use an OMP task pipeline that overlaps batch B+1's collocation with batch B's GEMM via dependent tasks on a ring of pre-allocated buffer slots. Otherwise fall back to the plain omp-for loop which has lower overhead at saturation. Benchmark (orbital speedup vs PySCF, benzene cc-pVDZ): 80^3/16t: 3.0x -> 3.6x (+20%) 80^3/64t: 4.7x -> 7.0x (+49%) 160^3/4t: 1.3x -> 1.9x (+46%) 160^3/16t: 3.2x -> 3.6x (+13%) 160^3/64t: 8.6x -> 7.6x (within noise, no regression) 160^3/208t: 11.3x -> 11.9x (+5%)
ConradJohnston
force-pushed
the
feature/cojoh/orbital-evaluator-and-cube
branch
from
April 29, 2026 19:43
dccd4d0 to
2b9f197
Compare
Move the nbf*nbf eval_xmat scratch buffer from per-omp-parallel-region allocation to a pre-allocated per-thread pool in the Impl. Saves nthreads * nbf^2 * 8 bytes of allocation churn per eval_density call when the evaluator is reused across multiple density evaluations (e.g. dumping 30 MO densities). No measurable impact in single-call benchmarks (allocation was already per-thread, not per-batch), but eliminates a malloc hot-spot in the multi-call pattern.
ConradJohnston
force-pushed
the
feature/cojoh/orbital-evaluator-and-cube
branch
from
April 29, 2026 19:56
2b9f197 to
5dc6806
Compare
Add write_cube_hdf5() alongside write_cube() for writing cube field data to HDF5 files. Guarded behind GAUXC_HAS_HDF5 and uses the existing HighFive dependency. The HDF5 layout stores: /field (nx, ny, nz) float64 — the scalar field /grid/origin (3,) float64 /grid/spacing (3,) float64 /grid/shape (3,) int64 /atoms/Z (natom,) int64 /atoms/coords (natom, 3) float64 /comment string attribute on /field This is useful for downstream analysis tools (h5py, HDFView) that can read the field directly as a numpy array without parsing text. Includes a round-trip test verifying field values, grid metadata, atomic geometry, and comment attribute.
- Guard #include <omp.h> with #ifdef _OPENMP and provide serial fallbacks for omp_get_max_threads / omp_get_thread_num, so the code compiles without OpenMP enabled. - Remove leftover internal optimization labels from comments. - Replace POSIX-only unistd.h / getpid() with a portable temp-path helper to fix Windows builds. - Add test case for CubeGrid eval_orbital / eval_density overloads, verifying they produce identical results to the pointer-based API.
ConradJohnston
marked this pull request as ready for review
April 29, 2026 22:23
The new orbital_evaluator_test cases that round-trip cube/HDF5 files
all wrote to /tmp paths derived from a per-process static counter.
Under GAUXC_MPI_TEST (mpiexec -n 2) both ranks share /tmp and produced
identical paths, racing on fopen("w")/fwrite/ifstream/std::remove
and causing assertion failures.
Match the pattern used in moltypes_test / collocation / runtime:
gate the I/O TEST_CASEs on rank 0 via a small is_root_rank() helper
that uses MPI_Initialized to stay safe when MPI is disabled. Also
include the PID in make_temp_path() as a defensive measure for any
future I/O test that might forget the guard.
The pure-compute TEST_CASEs are left to run on all ranks (they have
no shared state).
Blockers - Remove the per-evaluator xmat scratch pool. It was sized from omp_get_max_threads() at construction but indexed by omp_get_thread_num() at call time, so raising the thread count after construction indexed out of bounds, and nested or concurrent evaluation aliased a single buffer across teams. Scratch is now per-thread and per-call inside the parallel region, and the eval_xmat buffer grows to the largest nbe actually seen rather than nbf^2. This also drops the eager multi-GiB allocation that callers using only eval_orbitals were paying for. - Add coverage for the screened paths, none of which any test reached: nbe == 0, 0 < nbe < nbf against an unscreened reference, multi-batch grids, and thread-count invariance. Structure - Collapse four copies of the screen/collocate/gather/contract loop into a single batched_eval parameterised over a point source and a contractor. - Delete the OMP task pipeline. It engaged only when n_batches was below 2*nthreads and could not add parallelism, since every phase-B task depends serially on its phase-A task. - Size batches from both cache footprint (~1 MiB of AO scratch per thread) and load balance (npts/(4*nthreads)), instead of clamping to a constant 8192 for every nbf <= 256. Cube writer - Emit NAN/INF in the case glibc %E uses, and honour signbit for NaN. - Place rows at exact offsets (nz*13 + ceil(nz/6)) and stage them in bounded chunks, removing the serial compaction pass and the ~1.7 GiB staging buffer at 512^3. - Wrap the FILE* so the large allocations cannot leak it. Conventions - Move the pimpl to detail::OrbitalEvaluatorImpl, add OrbitalEvaluatorFactory to own the ExecutionSpace switch, use size_t for point counts and leading dimensions, and move CubeGrid out from behind the cube-format header. Also rejects ldo values that would truncate in the BLAS int parameter, computes grid batch bounding boxes analytically instead of rescanning the coordinates just generated, and drops dead members, unused includes and the POSIX-only test helpers. Note that batch size now depends on the thread count, so screened results can differ between thread counts by up to the shell tolerance. This is documented on the class and pinned by a test.
…ng tolerance Screening quality is set by the bounding box of each batch, and a contiguous index range makes a poor one: with iz varying fastest, a batch that crosses a row boundary spans the entire z extent, so shells anywhere along that column survive. Walking the grid in spatially compact tiles instead cuts the mean surviving basis count on a 110-atom system from 166 to 123, worth 1.4x on orbitals and 1.65x on density. Tiles are used only where they can help. A tighter box can remove nothing if the whole z extent already lies within a shell cutoff radius, and the scatter it costs is then pure overhead -- measured at 0.82x on benzene, whose grid spans 0.88 cutoff radii along z against taxol's 3.08. Grids below that threshold keep the contiguous path. A batch is now described by its runs of output indices, so the two traversals share one kernel: contiguous batches are the single-run case and still write straight into the output. Batch size was derived from a 1 MiB target for the AO block alone, but each thread holds three such blocks -- the AO block, the collocation kernel's own transpose staging, and, for density, the D*AO product. Sharing the target across all of them rather than granting it to one is worth ~1.5x on density for small systems and leaves large ones unchanged. The AO and D*AO blocks are also sized to the screened count actually seen rather than the nbf worst case, cutting per-thread scratch 2.5x on the 110-atom case at no cost in time. Screening tolerance becomes a constructor argument applied to the evaluator's own copy of the basis. It was previously reachable only by mutating the caller's basis in place, which would have silently loosened XC screening for anyone sharing that object with an SCF setup. A basis inherited from such a setup may sit at machine epsilon, which costs ~6x here for precision a five-significant-digit cube file cannot represent. Orbital error tracks the tolerance directly and density error goes as its square, so density tolerates a looser setting; both are pinned by a test. Two test assertions claimed bitwise agreement between batch decompositions that tiling makes genuinely distinct -- grid against pointer traversal, and one thread count against another. Both now assert agreement to within the screening tolerance, which is what the code guarantees. The thread-count case is renamed to what it actually guards: that scratch follows the thread count in force at evaluation time. Validated against an unscreened PySCF reference on the 110-atom case: density agrees to 5e-16 L2 relative, orbitals to 1.2e-10 absolute.
Spatial tiling is selected only for grids spanning well beyond a shell cutoff radius, and every grid in the suite sat below that threshold, so the tiled path ran in none of them. Instrumenting the predicate reported "linear" for all seven grid evaluations. Adds a case whose 32 Bohr z extent is several times the largest cc-pVDZ cutoff radius, with the two centres far enough apart that screening is active within it. Checks the tiled result against both a dense unscreened eval_collocation reference and the contiguous traversal reached through the pointer overload, so a fault in either the tile geometry or the output scatter is caught.
Coverage measurement put src/external/cube.cxx at 83% of lines, the
weakest file in this series, with the shortfall inside format_e13_5
rather than on error paths.
Two branches were unreachable:
- The snprintf padding path. "%13.5E" sets a *minimum* field width of
13, so snprintf cannot return fewer than 13 characters and the pad
loop was dead.
- The downward mantissa renormalisation. Over the range where
pow(10,-exp10) is accurate to a few ulp the scaled mantissa cannot
fall below 1, so the branch only fired for subnormal inputs, whose
exponent then sent them to the snprintf deferral anyway, discarding
the work. Testing the exponent range before scaling makes this
explicit and removes the branch.
Verified byte-identical to the previous implementation over 11.17M
values covering random bit patterns, a 4M-point subnormal sweep, every
power of ten and its neighbours, and 6th-significant-digit boundaries.
Also extends the snprintf byte-equality battery with subnormals, both
signs of 3-digit exponent, and a mantissa carrying from E+99 into a
3-digit exponent, which is the one deferral the old tests missed.
The doc comment and the neighbouring test claimed the formatter differs
from glibc only at exact half-way ties. That is too narrow: 9.99999...e-98
sits just below the tie and still differs, because scaling perturbs the
mantissa across the boundary. Both remain within one unit of the last
printed digit, which is what the test pins. Comments corrected.
cube.cxx line coverage 83.3% -> 92.8%; every line now exercised except
GAUXC_GENERIC_EXCEPTION throws.
bbox_for_index_range has three cases: a batch spanning whole grid rows, one spanning whole z columns, and one contained in a single row, which gets an exact z sub-range. Only the first two were reached; the third had no coverage. Exercising it needs three conditions at once. The row must outlast a batch, so nz has to exceed the batch size. The grid must stay off the tiled path, so the z extent has to sit under twice the median cutoff radius, which the test derives from the basis rather than hard-coding a spacing that a different tolerance would invalidate. Evaluation is pinned to one thread so the batch is npts/4 instead of thread-count dependent, which puts three quarters of a row in a single batch on any machine. The geometry matters as much as the sizes. A first attempt placed the molecule at the near end of z; it executed the branch but a mutant that collapsed the box to a point survived, because the face nearest the shells was unchanged and screening did not move. Putting the molecule at the far end makes the sub-range the deciding face: the same mutant now fails 10425 assertions with errors of order 0.1, against a 1e-9 margin. Also cross-checks against the pointer overload, which boxes the points it is handed rather than deriving a box from grid indices, so the index arithmetic is validated against an independent decomposition. orbital_evaluator.cxx line coverage 95.9% -> 96.5%; every remaining uncovered line is a throw, a move operation, or the basis() accessor.
Review item 17 suggests skipping the coefficient gather when every shell survives screening, or when the surviving AOs form one contiguous range. Implemented and measured, it is not worth taking. The two conditions are close to mutually exclusive with the cost they target. The gather only exists because shells were screened out; the bypass only fires when they were not. Counting batches with 20 orbitals: water 80^3 nbe==nbf 100% contiguous 100% gathered 0.0 MB benzene 80^3 nbe==nbf 4.6% contiguous 5.7% gathered 21.0 MB benzene 160^3 nbe==nbf 2.5% contiguous 3.0% gathered 161.7 MB taxol 40^3 nbe==nbf 0.0% contiguous 0.8% gathered 29.1 MB Water hits the bypass on every batch and saves nothing, because with 24 basis functions there was no gather to speak of. Everywhere else the bypass fires on a few percent of batches and removes under 2% of the copied volume. End to end on benzene 160^3 the difference was unmeasurable: 13 of 25 interleaved pairs favoured the bypass, a sign test p of 0.50, against 51% thermal drift on this machine. It also has a cost. Handing C to BLAS in place puts ldc into an int parameter for the first time, which needs a range check that rejects leading dimensions the current code accepts. The relative weight of the gather did rise when the batch size was corrected, from the 0.24% the review measured at 8192 points to roughly 2% at present sizes, so the concern is real. This bypass is not the remedy; avoiding gau2grid's transpose (review item 14) would be.
The Contractor::init size hint was never read by either contractor. The comment on eval_collocation's AO packing flagged a discrepancy in the non-gau2grid fallback without saying that the top-level CMakeLists makes gau2grid a hard dependency, so that path cannot be built.
… buffer
C_compressed was the one scratch buffer sized up front rather than grown on
demand, which is why the Contractor concept carried an init hook. Every
other buffer here (dm_ao, xmat_scr, staging) already uses the same
grow-if-short idiom, so switching this one over lets init go entirely: it
was empty in DensityContractor and a single resize in OrbitalContractor.
This is a uniformity change, not an optimisation. The high-water argument
that motivates growing dm_ao on demand does not transfer: dm_ao is nbe by
np and scales with the batch point count, whereas C_compressed is nbe by
nmo and comes to roughly 175 KB per thread at nbf 1099. Measured on taxol
80^3, 20 orbitals, 16 threads:
peak RSS 109 MB both ways
runtime 17 of 40 interleaved pairs favoured on-demand growth, sign
test p 0.74, against 39% drift
so it neither saves memory nor costs time.
The header branches on GAUXC_HAS_HDF5 but only saw it through a seven-deep transitive chain (cube_grid.hpp -> molecule.hpp -> ... -> gauxc_config.hpp). Had that chain changed, write_cube_hdf5 would have quietly stopped being declared while still being compiled into the library.
A tiled batch is scattered into the output column by column, but every test that used the tiled path asked for one orbital and every test that asked for several used the pointer overload, which batches contiguously and hands the GEMM the output directly. So the scatter had only ever run with nmo == 1 and ldo == npts, leaving both of its strides unchecked. The new case runs the tiled path with three orbitals and padded ldc/ldo, poisons the padding to catch a stride that reads past nbf or writes past npts, and cross-checks against the contiguous decomposition. Replacing ldo_ with np in the scatter fails 3160 of its 3871 assertions. Grids with an axis of one point were also untested. from_molecule gives such an axis zero spacing, which make_tile_source has to treat as unconstrained rather than divide by. Both paths were additionally fuzzed offline against an unscreened reference over 400 random molecule/grid/nmo/ldc/ldo/tolerance combinations, with tiling forced on for every grid, with no failures; shrinking the screening radius to sqrt(0.6) of its value failed 185 of the 400, so the sweep was sensitive to the thing it was checking.
Several public entry points had no test at all. They are cheap to get wrong
in a refactor and none of them would have been caught:
- basis(), never called. Now checks that the accessor returns the
evaluator's own retuned copy rather than the caller's: given a basis at
1e-12 and a tolerance of 1e-4, the exposed shells must have strictly
shorter cutoff radii while the caller's keep theirs. Dropping the retune
loop fails it.
- The move constructor and move assignment, never exercised. Both are
defaulted, but nothing checked the pimpl actually travels. Assignment is
made over an evaluator built at a different tolerance so a surviving
target would change the answer.
- A non-Host ExecutionSpace, which is documented to throw.
- Null C, D, out and points on all four eval entry points, and ldd < nbf.
Only ldc, ldo and the INT_MAX ceiling had been covered.
- The early returns for npts == 0, nmo < 1 and an empty grid, which must
be no-ops rather than errors and must not touch the output.
- CubeGrid::from_molecule's guards on an empty molecule and a zero extent,
and points_into against points().
- write_cube on a null field, an empty grid and an unopenable path, plus
the default comment line; write_cube_hdf5's two guards where built.
Every density test so far used a diagonal D and passed ldd == nbf. Both choices are blind to real defects. A diagonal D cannot detect a scrambled compressed-submatrix map, because permuting an identity on both sides leaves it unchanged. Building the map from a reversed shell list is therefore silent under the old tests and wrong under any dense D. And with ldd == nbf, dropping ldd in favour of nbf when calling eval_xmat is a no-op. The new case runs a dense symmetric D through a padded, poisoned buffer with screening active, against a reference contracted directly from the unscreened AO block. It fails 664 of its 898 assertions for the ldd defect and 566 for the reversed map. Both were also confirmed against a random sweep of 300 configurations varying basis (cc-pVDZ and 6-31G(d)), Cartesian and spherical shells, grid shape and sign of spacing, thread count, nmo, ldc, ldo and ldd, including grids placed far enough away that every shell screens out. On the current code it is clean; the earlier identity-D sweep reported zero failures against both defects, which is what prompted this test.
The class doc said results may differ between thread counts "by at most the basis-set shell tolerance", and the test pinned exactly that. Measured, the bound is wrong: on benzene in cc-pVDZ at 24^3 the orbital deviation between 1 and 16 threads reached 2.15x the tolerance. It should never have been stated as a hard bound. Each shell dropped by screening contributes up to the tolerance and the contributions accumulate, so the deviation grows with how many are dropped and with the coefficient magnitudes. The doc now says "of the order of" and gives the measured figure, and the test's margin scales with nbf as the generous ceiling on the number of shells that can be dropped. That is still eight orders below the field itself, so it remains a real constraint rather than a rubber stamp; it is also no longer sensitive to the machine's thread count, which the old 1x margin was. Bit-reproducibility at a fixed thread count, the other half of the claim, does hold: 8 repeats at each of 1, 2, 4, 8 and 16 threads over four grid sizes were bitwise identical every time.
Master gained a C API, Windows build fixes and gau2grid 2.0.9 since this
branch started. All three conflicts were additive: both sides appended to
the same CMake source lists, so both sides are kept.
src/CMakeLists.txt atomic_radii.cxx alongside cube_grid.cxx and
orbital_evaluator.cxx
src/external/CMakeLists.txt cube_hdf5.cxx alongside the new
GAUXC_ENABLE_C block
tests/CMakeLists.txt orbital_evaluator_test.cxx alongside
c_api_test.cxx
Verified by configuring and building from scratch and running the whole
suite: 3641389 assertions in 54 test cases.
Incidentally, gau2grid 2.0.9 fixes an aligned_alloc call that asked for
64-byte alignment with a size that was not a multiple of it, which C11
forbids and AddressSanitizer rejects. Sanitizer runs on this branch no
longer need a workaround for it.
from_molecule applies the same margin on every axis, which is the PySCF cubegen convention and is what it says it does. The consequence is easy to walk into: benzene's nuclei span 0.007 Bohr in z, so the default gives a 6.0 Bohr box that way against 15.5 and 14.2 in x and y. Its HOMO is still 0.038 at the z face, roughly sixty times the amplitude at the y face, so a plot at a typical isovalue of 0.02 comes out sliced flat and 3.5% of the norm falls outside the box. Behaviour is unchanged. Matching cubegen is deliberate, and callers who want a different box can set the fields directly.
A batch spanning more than one row is boxed by the whole extent of the faster axes. Shrinking that box by one grid step went unnoticed by the suite: hi_idx[1] = ny - 2 in place of ny - 1 left every test green. It went unnoticed because a one-step shrink only drops shells lying between cutoff-step and cutoff of the box, and a shell that far out contributes about the shell tolerance by construction. On a fine grid the error is therefore smaller than the tolerance it would be measured against. The new case sets the y spacing to 1.5x the largest cutoff radius, measured from the basis rather than assumed, so a shell on the far face is either fully kept or fully lost, and puts an atom exactly there. The mutant fails 48 of its 259 assertions. The z equivalents of the same mutation cannot be made to bite and are left alone. A z spacing above the cutoff radius gives a z extent past the tiling threshold, so the grid takes the tiled path and never builds this box; and at nz == 1 the shrunk corner drops below the low corner, where the min/max that follows widens the box rather than narrowing it.
eval_xmat only compresses D into scratch when the surviving shells span more than one contiguous AO block; a single block is read in place. Sizing the scratch unconditionally cost nbe*nbe per thread for nothing, which at nbf=2000 is 32 MiB a thread on the worst-case batch.
ldd reaches blas::gemm through eval_xmat, which narrows it to int32_t. The orbital path already range-checked ldo for the same reason; the density path did not.
The predicate selecting between a contiguous and a tiled traversal turned tiling off for the default from_molecule grid, which is where it helps most. Counting nbe*np and nbe^2*np per batch, a deterministic proxy for collocation and GEMM cost, tiling never did more work anywhere measured: benzene and taxol, 64^3 to 200^3, margins of 3 and 6 Bohr, and a single-plane grid where the scatter degenerates to runs of one point. It ranged from a 2-3% saving on benzene at 128^3 to 53% on taxol at 200^3, so the second traversal is gone rather than retuned. This also removes a geometry-dependent flip: which traversal ran depended on grid extent, so adding one grid plane could change every value in the output. The grid and pointer overloads still decompose differently, which the header now states. make_tile_source could previously exit its rescale loop still above target, since each pass only shrinks the overshoot as r^(2/3). Halving the longest axis afterwards converges regardless.
The ~1 MiB figure covers only the blocks proportional to the batch. The nbe*nbe and nbe*nmo gathers do not scale with the batch size, so nothing chosen here bounds them, and on a large basis with weak screening the nbe*nbe gather is the larger of the two.
…ount Batch shape decides which shells screen out, so taking omp_get_max_threads() into the batch size made the returned numbers depend on the thread count. The load-balance term is now a fixed minimum batch count, covering 256 threads at four batches each; measured against a run with the bound lifted, water is faster on both paths at 1 and 16 threads and benzene trades ~4% on orbitals for ~10% on density. The two thread-count tests now demand bitwise equality rather than agreement within the shell tolerance.
The evaluator holds no mutable state, so const has to mean safe from several threads at once and from inside an enclosing parallel region. Both modes previously returned wrong numbers with no crash, which is exactly the failure a test has to catch since nothing else will.
The cache bound keeps shrinking with nbf, reaching six points at ubiquitin/cc-pVDZ, so above nbf ~340 it is the floor that sets the batch. The floor was never measured. Serial CPU time over batches of 32/64/128/256/512 puts the optimum at 32-64 for both taxol (nbf 1099) and ubiquitin (nbf 11577) despite the tenfold difference in basis size; 128 costs ~10% and 512 costs 2.5x. Halving it also halves the per-thread scratch on those systems. Below nbf ~340 the cache bound still wins and nothing changes, which is why no test moves.
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.
Summary
Adds
OrbitalEvaluatorandwrite_cube/CubeGridto GauXC. The evaluator wraps the host collocation driver with thread-parallel batched evaluation and per-batch shell screening; the cube writer is self-contained and has no kernel dependency. Either can be used independently. Also addswrite_cube_hdf5behindGAUXC_HAS_HDF5.API sketch
Performance
Versus PySCF
cubegen2.12.1, cc-pVDZ, 16 threads on an AMD Ryzen AI 7 PRO 350 under WSL2. Single-threaded BLAS on both sides, grid evaluation only — no SCF, no file I/O. The harness parses geometries out oftests/standards.cxxso the two codes cannot drift apart, and basis size is asserted identical for every row.The advantage is per-batch shell screening.
cubegenevaluates the full basis at every point — verified directly:eval_gto500 Bohr from the molecule, where every AO is identically zero, costs the same as on top of it. So its cost isO(npts · nbf)for orbitals andO(npts · nbf²)for density, while GauXC's tracks the locally survivingnbe. Concretely: taxol at 80³ has 10× benzene's basis but takes only 1.3× longer.Grids are walked in spatially compact tiles rather than contiguous index ranges, since a contiguous range spans the full z extent as soon as it crosses a row boundary and screens poorly. Tiling is the only grid traversal: counting the collocation and GEMM work each one actually does, tiling never lost anywhere measured — benzene and taxol, 64³ to 200³, margins of 3 and 6 Bohr, and a single-plane grid where its scatter degenerates to runs of one point — ranging from a 2–3% saving on benzene at 128³ to 53% on taxol at 200³.
Density is compared against a GEMM-contracted
cubegen. As shipped it contracts withnumpy.einsum('pi,ij,pj->p', ao, dm, ao), which costs a further 15× on taxol and 5.6× on benzene; against that, taxol at 40³ is ~970×.Accuracy against an unscreened PySCF reference evaluated at identical points (taxol, 40³, 1e-10 screening tolerance): density agrees to 4.9e-16 L2 relative, orbitals to 1.0e-10 absolute on a field spanning [-0.71, 0.32].
On the timing methodology
Wall-clock timing on this host is unreliable at the tens-of-milliseconds scale and simply taking more repeats does not fix it. Over 30 independent runs of each code, individual timings spanned a factor of 20 (
cubegenwater orbital: 27 ms to 545 ms, median 143 ms), and the running minimum was still falling at N=30 rather than converging — under WSL2 the available vCPU capacity drifts with host activity, so there is no fixed floor for a minimum to converge to. Bootstrapped spread on a 30-sample minimum ranged from 2% to 188% depending on the quantity.The water and benzene rows are therefore the median of 25 tightly interleaved pairs, each code measured immediately after the other so both see the same machine state and the drift cancels in the ratio. Interquartile ranges: water 2.2–4.1× (orbital) and 7.5–12.6× (density); benzene 3.1–4.1× and 9.9–12.2×. Quoted GauXC times are the fastest observed over those runs. The taxol rows are best of 3 and their margins are orders of magnitude outside any of this.
Tests
77,123 assertions across 22 cases for
[orbital_evaluator]and[cube]; 3,640,979 across 56 for the full suite. Covers the evaluator against a directeval_collocationreference, the screened paths (nbe == 0and0 < nbe < nbf), multi-batch grids, the tile bounding boxes and the scatter they require, multi-orbital scatter into a padded output, single-plane grids, accessors and move semantics, every argument guard and empty-work no-op, a dense density matrix through a padded leading dimension, invariance of the results to the thread count, safety of concurrentconstinvocation, CubeGrid overload consistency, cube and HDF5 round-trips, multi-chunk cube writes, and%13.5Eformatter edge cases including NaN, infinities, subnormals and 3-digit exponents.Known limitations
format_e13_5matches glibc%13.5Eexcept for values within a few ulp of a rounding boundary in the 6th significant digit, where scaling the mantissa can tip it across; the two then differ by one unit in the last printed digit, never more. Fuzzed againstsnprintf: 3M random bit patterns were byte-identical, and of 3M values deliberately parked within 3 ulp of a 6th-digit tie, 2.7% differed, all by exactly one in the last digit.ldoandlddabove the range BLAS is handed are rejected, since BLAS takes leading dimensions asint.Follow-ups
Device backends via the existing CUDA/HIP collocation drivers;
eval_orbital_gradient/eval_density_gradient; native-layout collocation to avoid Gau2Grid's transpose; NUMA-aware first touch.Closes #128