From 3f456929e2a0f1f62885437cb7a49f6c9637690c Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 16:37:32 +0000 Subject: [PATCH 01/31] Add OrbitalEvaluator and Gaussian cube writer utility 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. --- include/gauxc/external/cube.hpp | 95 ++++++++ include/gauxc/orbital_evaluator.hpp | 105 +++++++++ src/CMakeLists.txt | 1 + src/external/CMakeLists.txt | 3 + src/external/cube.cxx | 299 +++++++++++++++++++++++++ src/orbital_evaluator.cxx | 211 ++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/orbital_evaluator_test.cxx | 332 ++++++++++++++++++++++++++++ 8 files changed, 1047 insertions(+) create mode 100644 include/gauxc/external/cube.hpp create mode 100644 include/gauxc/orbital_evaluator.hpp create mode 100644 src/external/cube.cxx create mode 100644 src/orbital_evaluator.cxx create mode 100644 tests/orbital_evaluator_test.cxx diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp new file mode 100644 index 00000000..15ceb034 --- /dev/null +++ b/include/gauxc/external/cube.hpp @@ -0,0 +1,95 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include +#include + +#include + +namespace GauXC { + +/** @brief Specification of a Gaussian-cube-format 3D rectangular grid. + * + * The grid is axis-aligned with the Cartesian frame (the standard + * cube-format axis vectors are simply (spacing[0], 0, 0), (0, spacing[1], 0), + * (0, 0, spacing[2])). All quantities are in atomic units (Bohr). + * + * Total number of points: nx*ny*nz. Storage cost per scalar field: + * 8*nx*ny*nz bytes (double precision). + */ +struct CubeGrid { + std::array origin{0.0, 0.0, 0.0}; ///< (0,0,0) corner, Bohr + std::array spacing{0.2, 0.2, 0.2}; ///< Step on each axis, Bohr + int64_t nx = 80; + int64_t ny = 80; + int64_t nz = 80; + + /// Total number of grid points. + int64_t num_points() const noexcept { return nx * ny * nz; } + + /** @brief Build a default grid that tightly encloses a molecule. + * + * The bounding box of the atomic centres is extended by `margin` Bohr on + * each side and discretised with the requested number of points along + * each axis. Spacing is chosen so that the first and last grid points + * coincide with the extended bounding-box corners (PySCF cubegen + * convention). + * + * @param mol Molecule whose atomic centres define the bounding box. + * @param nx,ny,nz Number of grid points along each axis. + * @param margin Margin (Bohr) added on each side. Default 3.0 matches + * the PySCF cubegen default. + */ + static CubeGrid from_molecule(const Molecule& mol, + int64_t nx = 80, int64_t ny = 80, + int64_t nz = 80, double margin = 3.0); + + /** @brief Materialise the grid points as an AoS coordinate array. + * + * Returns a vector of length 3*num_points() laid out in row-major + * (ix, iy, iz) order with iz varying fastest, matching the field-element + * ordering expected by `write_cube`. Suitable to be passed directly as + * the `points` argument to `OrbitalEvaluator::eval_orbital` / + * `eval_density`. + */ + std::vector points() const; +}; + +/** @brief Write a Gaussian cube file. + * + * Field layout: row-major (ix, iy, iz) with iz varying fastest. Length must + * equal `grid.num_points()`. Values are written in fixed `%13.5E` format + * with six values per line, matching the standard cube-file convention + * produced by PySCF / Gaussian / NWChem. + * + * This routine performs only file I/O; it has no dependency on the GauXC + * numerical kernels and may be used independently of `OrbitalEvaluator`. + * + * @param path Output path. Parent directory must exist. + * @param mol Molecule (atomic numbers + Cartesian coordinates in Bohr) + * written into the cube-file header. + * @param grid Grid specification. + * @param field Length-`grid.num_points()` scalar field. + * @param comment Optional first comment line (the second comment line is + * always "Generated by GauXC"). If empty, a default first + * line is used. + */ +void write_cube(const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = ""); + +} // namespace GauXC diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp new file mode 100644 index 00000000..f0542e3e --- /dev/null +++ b/include/gauxc/orbital_evaluator.hpp @@ -0,0 +1,105 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include + +#include +#include + +namespace GauXC { + +/** @brief Evaluate molecular orbitals and densities on arbitrary point sets. + * + * Wraps the host collocation kernel exposed by the local work driver with a + * thread-parallel batched evaluation loop, hiding the driver factory and + * the per-thread AO scratch from callers. The class is designed to be + * constructed once per molecule and reused across many evaluations (e.g. + * one per active-space orbital) without re-initialising the driver. + * + * The evaluator is intentionally decoupled from any particular file format: + * it returns plain numerical arrays. For Gaussian cube file I/O, see + * `gauxc/external/cube.hpp`. + * + * Currently only `ExecutionSpace::Host` is supported; passing any other + * value will throw on construction. Device support can be added later by + * routing through a device-side collocation driver while keeping the same + * signatures. + */ +class OrbitalEvaluator { + public: + /** @brief Construct an evaluator bound to a basis set. + * + * @param basis Basis set (copied internally). + * @param exec Execution space; only `ExecutionSpace::Host` is supported. + */ + explicit OrbitalEvaluator(BasisSet basis, + ExecutionSpace exec = ExecutionSpace::Host); + + ~OrbitalEvaluator() noexcept; + + // Non-copyable, movable + OrbitalEvaluator(const OrbitalEvaluator&) = delete; + OrbitalEvaluator& operator=(const OrbitalEvaluator&) = delete; + OrbitalEvaluator(OrbitalEvaluator&&) noexcept; + OrbitalEvaluator& operator=(OrbitalEvaluator&&) noexcept; + + /// Number of basis functions (rows of the AO matrix). + int32_t nbf() const noexcept; + + /// Underlying basis set. + const BasisSet& basis() const noexcept; + + /** @brief Evaluate a single MO chi(r) = sum_mu C[mu] * phi_mu(r). + * + * @param[in] npts Number of evaluation points. + * @param[in] points AoS array of length 3*npts: (x0,y0,z0,x1,y1,z1,...) + * in atomic units (Bohr). + * @param[in] C MO coefficient vector, length nbf(). + * @param[out] out Length-npts array of MO values. + */ + void eval_orbital(int64_t npts, const double* points, + const double* C, double* out) const; + + /** @brief Evaluate `nmo` MOs simultaneously. + * + * Storage layout: + * - `C` : (nbf, nmo) column-major, leading dimension `ldc` (>= nbf). + * - `out` : (npts, nmo) column-major, leading dimension `ldo` (>= npts). + * + * Equivalent to calling `eval_orbital` `nmo` times but amortises the AO + * collocation evaluation across all MOs (single AO buffer, GEMM contraction). + */ + void eval_orbitals(int64_t npts, const double* points, + int32_t nmo, const double* C, int64_t ldc, + double* out, int64_t ldo) const; + + /** @brief Evaluate the electron density + * rho(r) = sum_{mu,nu} D[mu,nu] * phi_mu(r) * phi_nu(r). + * + * @param[in] npts Number of evaluation points. + * @param[in] points AoS array of length 3*npts in Bohr. + * @param[in] D (nbf, nbf) symmetric density matrix, column-major, + * leading dimension `ldd` (>= nbf). + * @param[out] out Length-npts array of density values. + */ + void eval_density(int64_t npts, const double* points, + const double* D, int64_t ldd, + double* out) const; + + private: + struct Impl; + std::unique_ptr pimpl_; +}; + +} // namespace GauXC diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1aed4b42..8cd463ce 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -41,6 +41,7 @@ add_library( gauxc molgrid_impl.cxx molgrid_defaults.cxx atomic_radii.cxx + orbital_evaluator.cxx ) target_include_directories( gauxc diff --git a/src/external/CMakeLists.txt b/src/external/CMakeLists.txt index 46612c81..b0add283 100644 --- a/src/external/CMakeLists.txt +++ b/src/external/CMakeLists.txt @@ -9,6 +9,9 @@ # # See LICENSE.txt for details # +# Cube file writer is a self-contained utility (no third-party deps); always build it. +target_sources( gauxc PRIVATE cube.cxx ) + if( GAUXC_ENABLE_HDF5 ) include(FetchContent) find_package(HDF5) diff --git a/src/external/cube.cxx b/src/external/cube.cxx new file mode 100644 index 00000000..018ad0c5 --- /dev/null +++ b/src/external/cube.cxx @@ -0,0 +1,299 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _OPENMP +#include +#endif + +#include + +namespace GauXC { + +// ============================================================================= +// CubeGrid +// ============================================================================= + +CubeGrid CubeGrid::from_molecule(const Molecule& mol, int64_t nx, int64_t ny, + int64_t nz, double margin) { + if (mol.empty()) { + GAUXC_GENERIC_EXCEPTION( + "CubeGrid::from_molecule: molecule has no atoms."); + } + if (nx < 1 || ny < 1 || nz < 1) { + GAUXC_GENERIC_EXCEPTION( + "CubeGrid::from_molecule: nx, ny, nz must be >= 1."); + } + + double xmin = mol[0].x, xmax = mol[0].x; + double ymin = mol[0].y, ymax = mol[0].y; + double zmin = mol[0].z, zmax = mol[0].z; + for (const auto& a : mol) { + xmin = std::min(xmin, a.x); + xmax = std::max(xmax, a.x); + ymin = std::min(ymin, a.y); + ymax = std::max(ymax, a.y); + zmin = std::min(zmin, a.z); + zmax = std::max(zmax, a.z); + } + + CubeGrid grid; + grid.origin = {xmin - margin, ymin - margin, zmin - margin}; + grid.nx = nx; + grid.ny = ny; + grid.nz = nz; + const double ex = (xmax - xmin) + 2.0 * margin; + const double ey = (ymax - ymin) + 2.0 * margin; + const double ez = (zmax - zmin) + 2.0 * margin; + grid.spacing[0] = nx > 1 ? ex / static_cast(nx - 1) : 0.0; + grid.spacing[1] = ny > 1 ? ey / static_cast(ny - 1) : 0.0; + grid.spacing[2] = nz > 1 ? ez / static_cast(nz - 1) : 0.0; + return grid; +} + +std::vector CubeGrid::points() const { + std::vector pts(static_cast(num_points()) * 3); + size_t k = 0; + for (int64_t ix = 0; ix < nx; ++ix) { + const double x = origin[0] + spacing[0] * static_cast(ix); + for (int64_t iy = 0; iy < ny; ++iy) { + const double y = origin[1] + spacing[1] * static_cast(iy); + for (int64_t iz = 0; iz < nz; ++iz) { + const double z = origin[2] + spacing[2] * static_cast(iz); + pts[3 * k + 0] = x; + pts[3 * k + 1] = y; + pts[3 * k + 2] = z; + ++k; + } + } + } + return pts; +} + +// Hand-rolled %13.5E formatter. Required field is fixed 13 chars +// (sign + D + '.' + 5d + 'E' + sign + 2d). snprintf in the inner loop +// dominates write time on large grids; this matches glibc snprintf +// bit-for-bit on the fast path (1-2 digit exponents) and falls back to +// snprintf for 3-digit-exponent edge cases. Output buffer must be +// exactly 13 bytes (no trailing NUL). + +namespace { + +inline void format_e13_5(double v, char* out) { + if (std::isnan(v)) { + std::memcpy(out, " NaN", 13); + return; + } + if (std::isinf(v)) { + std::memcpy(out, v < 0 ? " -Inf" : " Inf", 13); + return; + } + + bool negative = std::signbit(v); + double absv = std::fabs(v); + + if (absv == 0.0) { + // glibc "%13.5E": 0.0 -> " 0.00000E+00"; -0.0 -> " -0.00000E+00". + out[0] = ' '; + out[1] = negative ? '-' : ' '; + out[2] = '0'; + out[3] = '.'; + out[4] = '0'; + out[5] = '0'; + out[6] = '0'; + out[7] = '0'; + out[8] = '0'; + out[9] = 'E'; + out[10] = '+'; + out[11] = '0'; + out[12] = '0'; + return; + } + + // Exponent via floor(log10), with corrections for FP edge cases + // (e.g. 9.99999 rounding up across a power-of-ten boundary). + int exp10 = static_cast(std::floor(std::log10(absv))); + double scale = std::pow(10.0, -exp10); + double mant = absv * scale; + + long long mant_int = static_cast(std::llround(mant * 1e5)); + if (mant_int >= 1000000) { + mant_int = 100000; + ++exp10; + } else if (mant_int < 100000) { + --exp10; + scale = std::pow(10.0, -exp10); + mant = absv * scale; + mant_int = static_cast(std::llround(mant * 1e5)); + if (mant_int >= 1000000) { + mant_int = 999999; + } else if (mant_int < 100000) { + mant_int = 100000; + } + } + + // 3+ digit exponents have a different field layout; defer to snprintf. + if (exp10 > 99 || exp10 < -99) { + char tmp[32]; + const int n = std::snprintf(tmp, sizeof(tmp), "%13.5E", v); + if (n >= 13) { + std::memcpy(out, tmp + (n - 13), 13); + } else { + const int pad = 13 - n; + for (int i = 0; i < pad; ++i) out[i] = ' '; + std::memcpy(out + pad, tmp, static_cast(n)); + } + return; + } + + // Fast path. Layout: [sp][sign][D][.][d4 d3 d2 d1 d0][E][esign][e1 e0] + out[0] = ' '; + out[1] = negative ? '-' : ' '; + + char digits[6]; + for (int i = 5; i >= 0; --i) { + digits[i] = static_cast('0' + (mant_int % 10)); + mant_int /= 10; + } + out[2] = digits[0]; + out[3] = '.'; + out[4] = digits[1]; + out[5] = digits[2]; + out[6] = digits[3]; + out[7] = digits[4]; + out[8] = digits[5]; + out[9] = 'E'; + out[10] = exp10 < 0 ? '-' : '+'; + + const int aexp = exp10 < 0 ? -exp10 : exp10; + out[11] = static_cast('0' + (aexp / 10)); + out[12] = static_cast('0' + (aexp % 10)); +} + +} // namespace + +// ============================================================================= +// write_cube +// ============================================================================= + +void write_cube(const std::string& path, const Molecule& mol, + const CubeGrid& grid, const double* field, + const std::string& comment) { + if (field == nullptr) { + GAUXC_GENERIC_EXCEPTION("write_cube: field pointer is null."); + } + if (grid.num_points() <= 0) { + GAUXC_GENERIC_EXCEPTION("write_cube: grid has zero points."); + } + + std::FILE* f = std::fopen(path.c_str(), "w"); + if (f == nullptr) { + GAUXC_GENERIC_EXCEPTION("write_cube: failed to open output file: " + path); + } + + // --- Header --- + std::fprintf(f, "%s\n", + comment.empty() ? "GauXC cube file" : comment.c_str()); + std::fprintf(f, "Generated by GauXC\n"); + + // natoms + origin (Bohr). + std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(mol.size()), grid.origin[0], + grid.origin[1], grid.origin[2]); + + // Three voxel-axis lines (axis-aligned grid). + std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nx), grid.spacing[0], 0.0, 0.0); + std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.ny), 0.0, grid.spacing[1], 0.0); + std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nz), 0.0, 0.0, grid.spacing[2]); + + // One line per atom: Z, partial charge (0.0), x, y, z (Bohr). + for (const auto& atom : mol) { + std::fprintf(f, "%5lld %12.6f %12.6f %12.6f %12.6f\n", + static_cast(atom.Z.get()), 0.0, atom.x, atom.y, + atom.z); + } + + // --- Data block --- + // Each (ix, iy) row is grid.nz values, six per line, %13.5E. The cube + // format requires a newline at the end of every (ix, iy) row regardless + // of how many values land on the last line. Rows are independent, so we + // format them in parallel into a pre-sized buffer and commit with a + // single fwrite. + const int64_t nz = grid.nz; + const int64_t lines_per_row = (nz + 5) / 6; + // Worst case 6*13 + 1 = 79 bytes per line; over-allocates the trailing + // line of each row but avoids a precise sizing pass. + const int64_t bytes_per_row = lines_per_row * (6 * 13 + 1); + const int64_t n_rows = grid.nx * grid.ny; + std::vector buf(static_cast(bytes_per_row * n_rows)); + std::vector row_byte_count(static_cast(n_rows), 0); + +#pragma omp parallel for schedule(static) + for (int64_t row = 0; row < n_rows; ++row) { + const double* row_data = + field + static_cast(row) * static_cast(nz); + char* dst = buf.data() + static_cast(row) * + static_cast(bytes_per_row); + int64_t off = 0; + for (int64_t iz = 0; iz < nz; ++iz) { + format_e13_5(row_data[iz], dst + off); + off += 13; + // Every 6 values OR at the end of the row → newline. + if (((iz + 1) % 6 == 0) || (iz + 1 == nz)) { + dst[off++] = '\n'; + } + } + row_byte_count[static_cast(row)] = off; + } + + // Compact rows in place (worst-case padding between them) and emit + // with a single fwrite. + if (n_rows > 1) { + int64_t write_off = row_byte_count[0]; + for (int64_t row = 1; row < n_rows; ++row) { + const int64_t src_off = row * bytes_per_row; + const int64_t len = row_byte_count[static_cast(row)]; + std::memmove(buf.data() + write_off, buf.data() + src_off, + static_cast(len)); + write_off += len; + } + if (std::fwrite(buf.data(), 1, static_cast(write_off), f) != + static_cast(write_off)) { + std::fclose(f); + GAUXC_GENERIC_EXCEPTION("write_cube: short write to " + path); + } + } else { + if (std::fwrite(buf.data(), 1, static_cast(row_byte_count[0]), + f) != static_cast(row_byte_count[0])) { + std::fclose(f); + GAUXC_GENERIC_EXCEPTION("write_cube: short write to " + path); + } + } + + if (std::fclose(f) != 0) { + GAUXC_GENERIC_EXCEPTION("write_cube: failed to close " + path); + } +} + +} // namespace GauXC diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx new file mode 100644 index 00000000..39f5b61c --- /dev/null +++ b/src/orbital_evaluator.cxx @@ -0,0 +1,211 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include + +#include +#include +#include +#include + +#include +#include + +#include "xc_integrator/local_work_driver/host/blas.hpp" +#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" + +namespace GauXC { + +namespace { + +/// Per-thread point batch size, sized so the AO scratch buffer +/// (nbf*batch doubles) fits comfortably in L2 (~16 MB target). +inline int64_t choose_batch_size(int32_t nbf) { + constexpr int64_t kTargetBytesPerThread = 16ll * 1024 * 1024; + constexpr int64_t kMinBatch = 256; + constexpr int64_t kMaxBatch = 8192; + if (nbf <= 0) return kMaxBatch; + const int64_t b = + kTargetBytesPerThread / (static_cast(sizeof(double)) * nbf); + return std::clamp(b, kMinBatch, kMaxBatch); +} + +/// Single-block submat_map for the no-screening case (nbe == nbf); makes +/// `eval_xmat` route directly to GEMM without invoking the gather path. +inline LocalHostWorkDriver::submat_map_t full_submat_map(int32_t nbf) { + return {{ {int32_t{0}, nbf, int32_t{0}} }}; +} + +} // namespace + +struct OrbitalEvaluator::Impl { + BasisSet basis; + std::unique_ptr driver_owner; + LocalHostWorkDriver* host_driver = nullptr; // non-owning view + std::vector shell_list; + int32_t nbf_ = 0; + + void init(BasisSet bs, ExecutionSpace exec) { + if (exec != ExecutionSpace::Host) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator: only ExecutionSpace::Host is currently supported."); + } + + basis = std::move(bs); + + driver_owner = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference"); + host_driver = dynamic_cast(driver_owner.get()); + if (host_driver == nullptr) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator: LocalWorkDriverFactory did not return a " + "LocalHostWorkDriver."); + } + + shell_list.resize(basis.size()); + std::iota(shell_list.begin(), shell_list.end(), int32_t{0}); + nbf_ = basis.nbf(); + } +}; + +OrbitalEvaluator::OrbitalEvaluator(BasisSet basis, ExecutionSpace exec) + : pimpl_(std::make_unique()) { + pimpl_->init(std::move(basis), exec); +} + +OrbitalEvaluator::~OrbitalEvaluator() noexcept = default; +OrbitalEvaluator::OrbitalEvaluator(OrbitalEvaluator&&) noexcept = default; +OrbitalEvaluator& OrbitalEvaluator::operator=(OrbitalEvaluator&&) noexcept = + default; + +int32_t OrbitalEvaluator::nbf() const noexcept { return pimpl_->nbf_; } + +const BasisSet& OrbitalEvaluator::basis() const noexcept { + return pimpl_->basis; +} + +void OrbitalEvaluator::eval_orbital(int64_t npts, const double* points, + const double* C, double* out) const { + eval_orbitals(npts, points, /*nmo=*/1, C, /*ldc=*/pimpl_->nbf_, out, + /*ldo=*/npts); +} + +void OrbitalEvaluator::eval_orbitals(int64_t npts, const double* points, + int32_t nmo, const double* C, int64_t ldc, + double* out, int64_t ldo) const { + if (npts == 0 || nmo == 0) return; + if (points == nullptr || C == nullptr || out == nullptr) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals: null pointer argument."); + } + const int32_t nbf = pimpl_->nbf_; + if (ldc < nbf) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals: ldc must be >= nbf()."); + } + if (ldo < npts) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals: ldo must be >= npts."); + } + + const int32_t nshells = pimpl_->basis.nshells(); + const int64_t batch_size = choose_batch_size(nbf); + const int64_t n_batches = (npts + batch_size - 1) / batch_size; + LocalHostWorkDriver* driver = pimpl_->host_driver; + const int32_t* shell_list = pimpl_->shell_list.data(); + const BasisSet& basis = pimpl_->basis; + +#pragma omp parallel + { + std::vector ao_buf(static_cast(nbf) * batch_size); + +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < n_batches; ++b) { + const int64_t p0 = b * batch_size; + const int64_t np = std::min(batch_size, npts - p0); + + driver->eval_collocation(static_cast(np), + static_cast(nshells), + static_cast(nbf), points + 3 * p0, + basis, shell_list, ao_buf.data()); + + // out_slab(np, nmo) = ao^T(np, nbf) @ C(nbf, nmo); j-th MO column + // lives at out + j*ldo + p0. + blas::gemm( + /*TA=*/'T', /*TB=*/'N', + /*M=*/static_cast(np), /*N=*/static_cast(nmo), + /*K=*/nbf, /*ALPHA=*/1.0, /*A=*/ao_buf.data(), /*LDA=*/nbf, + /*B=*/C, /*LDB=*/static_cast(ldc), /*BETA=*/0.0, + /*C=*/out + p0, /*LDC=*/static_cast(ldo)); + } + } +} + +void OrbitalEvaluator::eval_density(int64_t npts, const double* points, + const double* D, int64_t ldd, + double* out) const { + if (npts == 0) return; + if (points == nullptr || D == nullptr || out == nullptr) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_density: null pointer argument."); + } + const int32_t nbf = pimpl_->nbf_; + if (ldd < nbf) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_density: ldd must be >= nbf()."); + } + + const int32_t nshells = pimpl_->basis.nshells(); + const int64_t batch_size = choose_batch_size(nbf); + const int64_t n_batches = (npts + batch_size - 1) / batch_size; + LocalHostWorkDriver* driver = pimpl_->host_driver; + const int32_t* shell_list = pimpl_->shell_list.data(); + const BasisSet& basis = pimpl_->basis; + +#pragma omp parallel + { + std::vector ao_buf(static_cast(nbf) * batch_size); + std::vector dm_ao_buf(static_cast(nbf) * batch_size); + // Sized for the eval_xmat submat-gather contract; unused on the + // nbe == nbf fast path but allocated unconditionally for safety. + std::vector xmat_scr(static_cast(nbf) * nbf); + auto submat_map = full_submat_map(nbf); + +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < n_batches; ++b) { + const int64_t p0 = b * batch_size; + const int64_t np = std::min(batch_size, npts - p0); + + driver->eval_collocation(static_cast(np), + static_cast(nshells), + static_cast(nbf), points + 3 * p0, + basis, shell_list, ao_buf.data()); + + // dm_ao = D @ ao + driver->eval_xmat( + /*npts=*/static_cast(np), /*nbf=*/static_cast(nbf), + /*nbe=*/static_cast(nbf), submat_map, /*fac=*/1.0, + /*P=*/D, /*ldp=*/static_cast(ldd), + /*basis_eval=*/ao_buf.data(), /*ldb=*/static_cast(nbf), + /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbf), + /*scr=*/xmat_scr.data()); + + // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) + driver->eval_uvvar_lda_rks( + /*npts=*/static_cast(np), /*nbe=*/static_cast(nbf), + /*basis_eval=*/ao_buf.data(), + /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbf), + /*den_eval=*/out + p0); + } + } +} + +} // namespace GauXC diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 46dbe487..3d4c0dfd 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -69,6 +69,7 @@ add_executable( gauxc_test basis/parse_basis.cxx dd_psi_potential_test.cxx 2nd_derivative_test.cxx + orbital_evaluator_test.cxx ) target_link_libraries( gauxc_test PUBLIC gauxc gauxc_catch2 Eigen3::Eigen ) if(GAUXC_ENABLE_CUTLASS) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx new file mode 100644 index 00000000..97562323 --- /dev/null +++ b/tests/orbital_evaluator_test.cxx @@ -0,0 +1,332 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include "ut_common.hpp" +#include "catch2/catch.hpp" + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include "standards.hpp" + +// Reference implementation lives in src/, behind the public header surface. +#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" + +using namespace GauXC; + +namespace { + +/// Build a deterministic PRNG-based set of points within a small bounding box +/// around the molecule. Avoids putting samples too close to nuclei to keep +/// values numerically well-behaved. +std::vector make_random_points(int64_t npts, unsigned seed = 1234u) { + std::mt19937 gen(seed); + std::uniform_real_distribution u(-2.5, 2.5); + std::vector pts(static_cast(npts) * 3); + for (int64_t p = 0; p < npts; ++p) { + pts[3 * p + 0] = u(gen); + pts[3 * p + 1] = u(gen); + pts[3 * p + 2] = u(gen); + } + return pts; +} + +} // namespace + +TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", + "[orbital_evaluator]") { + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + + const int32_t nbf = basis.nbf(); + REQUIRE(nbf > 0); + + const int64_t npts = 137; // not a multiple of any batch size + const auto pts = make_random_points(npts); + + // Reference: AO collocation directly via the LocalHostWorkDriver. + std::vector ao_ref(static_cast(nbf) * npts); + { + auto drv = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference"); + auto* host_drv = dynamic_cast(drv.get()); + REQUIRE(host_drv != nullptr); + std::vector shell_list(basis.size()); + for (size_t i = 0; i < shell_list.size(); ++i) + shell_list[i] = static_cast(i); + host_drv->eval_collocation(static_cast(npts), + static_cast(basis.nshells()), + static_cast(nbf), pts.data(), basis, + shell_list.data(), ao_ref.data()); + } + + OrbitalEvaluator eval(basis); + REQUIRE(eval.nbf() == nbf); + + SECTION("eval_orbital with one-hot coefficient reproduces single AO column") { + std::vector C(nbf, 0.0); + std::vector out(static_cast(npts), 0.0); + for (int32_t mu : {0, nbf / 3, nbf / 2, nbf - 1}) { + std::fill(C.begin(), C.end(), 0.0); + C[static_cast(mu)] = 1.0; + std::fill(out.begin(), out.end(), 0.0); + eval.eval_orbital(npts, pts.data(), C.data(), out.data()); + for (int64_t p = 0; p < npts; ++p) { + const double ref = + ao_ref[static_cast(p) * nbf + static_cast(mu)]; + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-12)); + } + } + } + + SECTION("eval_orbitals with random C matches AO ^T @ C") { + const int32_t nmo = 4; + std::vector C(static_cast(nbf) * nmo); + std::mt19937 gen(99); + std::uniform_real_distribution u(-1.0, 1.0); + for (auto& v : C) v = u(gen); + + std::vector out(static_cast(npts) * nmo, 0.0); + eval.eval_orbitals(npts, pts.data(), nmo, C.data(), nbf, out.data(), npts); + + for (int32_t j = 0; j < nmo; ++j) { + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + ref += C[static_cast(j) * nbf + mu] * + ao_ref[static_cast(p) * nbf + mu]; + } + const double got = + out[static_cast(j) * npts + static_cast(p)]; + CHECK(got == Approx(ref).margin(1e-10)); + } + } + } + + SECTION("eval_density with identity D equals sum of squared AO values") { + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t mu = 0; mu < nbf; ++mu) { + D[static_cast(mu) * nbf + mu] = 1.0; + } + std::vector out(static_cast(npts), 0.0); + eval.eval_density(npts, pts.data(), D.data(), nbf, out.data()); + + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + ref += a * a; + } + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-10)); + } + } + + SECTION("eval_density with rank-1 D = c c^T equals (c.AO)^2") { + std::vector c(nbf); + std::mt19937 gen(7); + std::uniform_real_distribution u(-1.0, 1.0); + for (auto& v : c) v = u(gen); + + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t mu = 0; mu < nbf; ++mu) { + for (int32_t nu = 0; nu < nbf; ++nu) { + D[static_cast(mu) * nbf + nu] = c[mu] * c[nu]; + } + } + std::vector out(static_cast(npts), 0.0); + eval.eval_density(npts, pts.data(), D.data(), nbf, out.data()); + + std::vector orb(static_cast(npts), 0.0); + eval.eval_orbital(npts, pts.data(), c.data(), orb.data()); + + for (int64_t p = 0; p < npts; ++p) { + const double ref = orb[static_cast(p)] * orb[static_cast(p)]; + CHECK(out[static_cast(p)] == Approx(ref).margin(1e-10)); + } + } +} + +TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { + auto mol = make_water(); + CubeGrid g = CubeGrid::from_molecule(mol, /*nx=*/16, /*ny=*/12, /*nz=*/8, + /*margin=*/2.0); + REQUIRE(g.num_points() == 16 * 12 * 8); + + auto pts = g.points(); + REQUIRE(pts.size() == static_cast(g.num_points()) * 3); + + // Check first and last grid points. + CHECK(pts[0] == Approx(g.origin[0])); + CHECK(pts[1] == Approx(g.origin[1])); + CHECK(pts[2] == Approx(g.origin[2])); + + const size_t last = static_cast(g.num_points()) - 1; + CHECK(pts[3 * last + 0] == + Approx(g.origin[0] + g.spacing[0] * (g.nx - 1))); + CHECK(pts[3 * last + 1] == + Approx(g.origin[1] + g.spacing[1] * (g.ny - 1))); + CHECK(pts[3 * last + 2] == + Approx(g.origin[2] + g.spacing[2] * (g.nz - 1))); + + // Check ordering: iz varies fastest. Point (1, 0, 0) should be at offset + // ny*nz, point (0, 1, 0) at nz, point (0, 0, 1) at 1. + const int64_t off_x = g.ny * g.nz; + const int64_t off_y = g.nz; + CHECK(pts[3 * static_cast(off_x) + 0] == + Approx(g.origin[0] + g.spacing[0])); + CHECK(pts[3 * static_cast(off_y) + 1] == + Approx(g.origin[1] + g.spacing[1])); + CHECK(pts[3 * 1 + 2] == Approx(g.origin[2] + g.spacing[2])); +} + +TEST_CASE("write_cube round-trips header and field data", "[cube]") { + auto mol = make_water(); + CubeGrid grid = CubeGrid::from_molecule(mol, /*nx=*/5, /*ny=*/4, /*nz=*/7, + /*margin=*/2.0); + std::vector field(static_cast(grid.num_points())); + for (int64_t i = 0; i < grid.num_points(); ++i) { + // Mix of magnitudes to exercise the formatter (negative, near-zero, etc.) + field[static_cast(i)] = std::sin(0.13 * static_cast(i)) * + std::pow(10.0, (i % 5) - 2); + } + + // Use a tmp path under /tmp. + const std::string path = std::string("/tmp/gauxc_cube_test_") + + std::to_string(::getpid()) + ".cube"; + write_cube(path, mol, grid, field.data(), "Test cube"); + + // Parse back. + std::ifstream in(path); + REQUIRE(in.is_open()); + std::string l1, l2; + std::getline(in, l1); + std::getline(in, l2); + CHECK(l1 == "Test cube"); + CHECK(l2 == "Generated by GauXC"); + + long long natoms_read; + double ox, oy, oz; + in >> natoms_read >> ox >> oy >> oz; + CHECK(natoms_read == static_cast(mol.size())); + CHECK(ox == Approx(grid.origin[0])); + CHECK(oy == Approx(grid.origin[1])); + CHECK(oz == Approx(grid.origin[2])); + + // 3 axis lines. + for (int axis = 0; axis < 3; ++axis) { + long long n; + double a, b, c; + in >> n >> a >> b >> c; + if (axis == 0) { + CHECK(n == grid.nx); + CHECK(a == Approx(grid.spacing[0])); + CHECK(b == Approx(0.0)); + CHECK(c == Approx(0.0)); + } else if (axis == 1) { + CHECK(n == grid.ny); + CHECK(a == Approx(0.0)); + CHECK(b == Approx(grid.spacing[1])); + CHECK(c == Approx(0.0)); + } else { + CHECK(n == grid.nz); + CHECK(a == Approx(0.0)); + CHECK(b == Approx(0.0)); + CHECK(c == Approx(grid.spacing[2])); + } + } + + // Atoms. + for (size_t i = 0; i < mol.size(); ++i) { + long long Z; + double q, x, y, z; + in >> Z >> q >> x >> y >> z; + CHECK(Z == mol[i].Z.get()); + CHECK(q == Approx(0.0)); + CHECK(x == Approx(mol[i].x)); + CHECK(y == Approx(mol[i].y)); + CHECK(z == Approx(mol[i].z)); + } + + // Field. Read all remaining whitespace-separated doubles and compare. + std::vector field_read; + field_read.reserve(static_cast(grid.num_points())); + double v; + while (in >> v) field_read.push_back(v); + + REQUIRE(field_read.size() == field.size()); + // %13.5E gives 5 significant digits → relative tolerance ~1e-5 for the + // round-trip. + for (size_t i = 0; i < field.size(); ++i) { + CHECK(field_read[i] == Approx(field[i]).epsilon(1e-4).margin(1e-30)); + } + + std::remove(path.c_str()); +} + +TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { + // Spot-check the custom formatter against snprintf for a battery of values + // by writing a tiny cube file and parsing it back. This is a stronger check + // than the round-trip above since we compare the byte stream. + auto mol = make_water(); + CubeGrid grid; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {0.1, 0.1, 0.1}; + grid.nx = 1; + grid.ny = 1; + grid.nz = 12; + + std::vector field = {0.0, -0.0, 1.23456e-10, -9.99995e-1, + 1.0e+99, -1.0e+99, 3.14159265358979, + -2.71828, 1.0, -1.0, 1e-300, + 1.234e+05}; + + const std::string path = std::string("/tmp/gauxc_cube_format_") + + std::to_string(::getpid()) + ".cube"; + write_cube(path, mol, grid, field.data(), "fmt"); + + std::ifstream in(path); + REQUIRE(in.is_open()); + // Skip header (2 comments + 4 grid lines + natoms atom lines). + std::string skip; + for (int i = 0; i < 2; ++i) std::getline(in, skip); + std::getline(in, skip); // natoms line + for (int i = 0; i < 3; ++i) std::getline(in, skip); // 3 axis lines + for (size_t i = 0; i < mol.size(); ++i) std::getline(in, skip); + + // Read the field-block as raw text and compare against snprintf + // line-by-line. Six values per line, then row terminator after 12 (single + // (ix, iy) row here so no early newline). + std::string data_block((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + + std::ostringstream expected; + for (size_t i = 0; i < field.size(); ++i) { + char buf[32]; + std::snprintf(buf, sizeof(buf), "%13.5E", field[i]); + expected << buf; + if ((i + 1) % 6 == 0 || (i + 1) == field.size()) expected << '\n'; + } + + CHECK(data_block == expected.str()); + + std::remove(path.c_str()); +} From 7ff6c623e1e049835ea5fc58a869b6b19737cf89 Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 18:40:23 +0000 Subject: [PATCH 02/31] OrbitalEvaluator: add per-batch shell screening 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). --- src/orbital_evaluator.cxx | 184 ++++++++++++++++++++++++++++++++------ 1 file changed, 157 insertions(+), 27 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 39f5b61c..1879f684 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -12,13 +12,18 @@ #include #include +#include #include #include +#include #include +#include #include +#include #include +#include "xc_integrator/integrator_util/integrator_common.hpp" #include "xc_integrator/local_work_driver/host/blas.hpp" #include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" @@ -44,6 +49,42 @@ inline LocalHostWorkDriver::submat_map_t full_submat_map(int32_t nbf) { return {{ {int32_t{0}, nbf, int32_t{0}} }}; } +/// Axis-aligned bbox of a set of npts AoS points (length 3*npts). +struct PointBbox { + std::array lo; + std::array hi; +}; +inline PointBbox compute_bbox(const double* points, int64_t npts) { + PointBbox b{{points[0], points[1], points[2]}, + {points[0], points[1], points[2]}}; + for (int64_t p = 1; p < npts; ++p) { + const double* xyz = points + 3 * p; + for (int k = 0; k < 3; ++k) { + if (xyz[k] < b.lo[k]) b.lo[k] = xyz[k]; + if (xyz[k] > b.hi[k]) b.hi[k] = xyz[k]; + } + } + return b; +} + +/// Squared distance from `center` to the nearest point in the axis-aligned +/// bbox `[lo, hi]^3`. Zero if the center lies inside the bbox. +inline double dist2_center_to_bbox(const double* center, + const PointBbox& bbox) { + double d2 = 0.0; + for (int k = 0; k < 3; ++k) { + const double c = center[k]; + if (c < bbox.lo[k]) { + const double dx = bbox.lo[k] - c; + d2 += dx * dx; + } else if (c > bbox.hi[k]) { + const double dx = c - bbox.hi[k]; + d2 += dx * dx; + } + } + return d2; +} + } // namespace struct OrbitalEvaluator::Impl { @@ -51,6 +92,13 @@ struct OrbitalEvaluator::Impl { std::unique_ptr driver_owner; LocalHostWorkDriver* host_driver = nullptr; // non-owning view std::vector shell_list; + // BasisSetMap powers shell -> AO range lookups for the screened + // submat_map. We don't have a Molecule available so we pass an empty + // one; the only field that needs it (shell_to_center) is unused here. + std::unique_ptr basis_map; + // Per-shell squared cutoff radius, cached so the screening loop is just + // a comparison instead of a square root per shell per batch. + std::vector shell_cutoff_r2; int32_t nbf_ = 0; void init(BasisSet bs, ExecutionSpace exec) { @@ -73,6 +121,14 @@ struct OrbitalEvaluator::Impl { shell_list.resize(basis.size()); std::iota(shell_list.begin(), shell_list.end(), int32_t{0}); nbf_ = basis.nbf(); + + basis_map = std::make_unique(basis, Molecule{}); + + shell_cutoff_r2.resize(basis.size()); + for (size_t s = 0; s < basis.size(); ++s) { + const double r = basis[s].cutoff_radius(); + shell_cutoff_r2[s] = r * r; + } } }; @@ -116,34 +172,79 @@ void OrbitalEvaluator::eval_orbitals(int64_t npts, const double* points, "OrbitalEvaluator::eval_orbitals: ldo must be >= npts."); } - const int32_t nshells = pimpl_->basis.nshells(); + const int32_t nshells_total = pimpl_->basis.nshells(); const int64_t batch_size = choose_batch_size(nbf); const int64_t n_batches = (npts + batch_size - 1) / batch_size; LocalHostWorkDriver* driver = pimpl_->host_driver; - const int32_t* shell_list = pimpl_->shell_list.data(); const BasisSet& basis = pimpl_->basis; + const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; + const BasisSetMap& basis_map = *pimpl_->basis_map; #pragma omp parallel { std::vector ao_buf(static_cast(nbf) * batch_size); + std::vector C_compressed(static_cast(nbf) * nmo); + std::vector screened_shells; + screened_shells.reserve(nshells_total); #pragma omp for schedule(dynamic, 1) for (int64_t b = 0; b < n_batches; ++b) { const int64_t p0 = b * batch_size; const int64_t np = std::min(batch_size, npts - p0); - driver->eval_collocation(static_cast(np), - static_cast(nshells), - static_cast(nbf), points + 3 * p0, - basis, shell_list, ao_buf.data()); - - // out_slab(np, nmo) = ao^T(np, nbf) @ C(nbf, nmo); j-th MO column - // lives at out + j*ldo + p0. + // Per-batch shell screening: keep shells whose cutoff_radius reaches + // any point of this batch's bounding box. + const PointBbox bbox = compute_bbox(points + 3 * p0, np); + screened_shells.clear(); + int32_t nbe = 0; + for (int32_t s = 0; s < nshells_total; ++s) { + if (dist2_center_to_bbox(basis[s].O_data(), bbox) < + shell_cutoff_r2[s]) { + screened_shells.push_back(s); + nbe += basis[s].size(); + } + } + if (nbe == 0) { + // All shells screen out for this batch -> orbital is identically + // zero. Zero the slab and skip eval. + for (int32_t j = 0; j < nmo; ++j) { + double* out_col = out + static_cast(j) * ldo + p0; + std::fill(out_col, out_col + np, 0.0); + } + continue; + } + + driver->eval_collocation( + static_cast(np), + static_cast(screened_shells.size()), + static_cast(nbe), points + 3 * p0, basis, + screened_shells.data(), ao_buf.data()); + + // Gather the rows of C corresponding to surviving shells into a + // contiguous (nbe, nmo) col-major buffer so the contraction is a + // dense GEMM. + { + int32_t row = 0; + for (int32_t s : screened_shells) { + const auto rng = basis_map.shell_to_ao_range(s); + const int32_t shell_nbe = rng.second - rng.first; + for (int32_t j = 0; j < nmo; ++j) { + const double* C_col = C + static_cast(j) * ldc; + double* dst = C_compressed.data() + + static_cast(j) * nbe + row; + std::copy(C_col + rng.first, C_col + rng.second, dst); + } + row += shell_nbe; + } + } + + // out_slab(np, nmo) = ao^T(np, nbe) @ C_compressed(nbe, nmo); j-th MO + // column lives at out + j*ldo + p0. blas::gemm( /*TA=*/'T', /*TB=*/'N', /*M=*/static_cast(np), /*N=*/static_cast(nmo), - /*K=*/nbf, /*ALPHA=*/1.0, /*A=*/ao_buf.data(), /*LDA=*/nbf, - /*B=*/C, /*LDB=*/static_cast(ldc), /*BETA=*/0.0, + /*K=*/nbe, /*ALPHA=*/1.0, /*A=*/ao_buf.data(), /*LDA=*/nbe, + /*B=*/C_compressed.data(), /*LDB=*/nbe, /*BETA=*/0.0, /*C=*/out + p0, /*LDC=*/static_cast(ldo)); } } @@ -163,46 +264,75 @@ void OrbitalEvaluator::eval_density(int64_t npts, const double* points, "OrbitalEvaluator::eval_density: ldd must be >= nbf()."); } - const int32_t nshells = pimpl_->basis.nshells(); + const int32_t nshells_total = pimpl_->basis.nshells(); const int64_t batch_size = choose_batch_size(nbf); const int64_t n_batches = (npts + batch_size - 1) / batch_size; LocalHostWorkDriver* driver = pimpl_->host_driver; - const int32_t* shell_list = pimpl_->shell_list.data(); const BasisSet& basis = pimpl_->basis; + const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; + const BasisSetMap& basis_map = *pimpl_->basis_map; #pragma omp parallel { std::vector ao_buf(static_cast(nbf) * batch_size); std::vector dm_ao_buf(static_cast(nbf) * batch_size); - // Sized for the eval_xmat submat-gather contract; unused on the - // nbe == nbf fast path but allocated unconditionally for safety. + // eval_xmat scratch when the submat-gather path is taken; sized for + // worst case (nbe == nbf). std::vector xmat_scr(static_cast(nbf) * nbf); - auto submat_map = full_submat_map(nbf); + std::vector screened_shells; + screened_shells.reserve(nshells_total); #pragma omp for schedule(dynamic, 1) for (int64_t b = 0; b < n_batches; ++b) { const int64_t p0 = b * batch_size; const int64_t np = std::min(batch_size, npts - p0); - driver->eval_collocation(static_cast(np), - static_cast(nshells), - static_cast(nbf), points + 3 * p0, - basis, shell_list, ao_buf.data()); - - // dm_ao = D @ ao + const PointBbox bbox = compute_bbox(points + 3 * p0, np); + screened_shells.clear(); + int32_t nbe = 0; + for (int32_t s = 0; s < nshells_total; ++s) { + if (dist2_center_to_bbox(basis[s].O_data(), bbox) < + shell_cutoff_r2[s]) { + screened_shells.push_back(s); + nbe += basis[s].size(); + } + } + if (nbe == 0) { + std::fill(out + p0, out + p0 + np, 0.0); + continue; + } + + driver->eval_collocation( + static_cast(np), + static_cast(screened_shells.size()), + static_cast(nbe), points + 3 * p0, basis, + screened_shells.data(), ao_buf.data()); + + // Build the compressed submat_map for this batch (gather the rows + // of D corresponding to surviving shells). Falls through to a single + // direct-gemm submat_map when nbe == nbf. + LocalHostWorkDriver::submat_map_t submat_map; + if (nbe == nbf) { + submat_map = full_submat_map(nbf); + } else { + std::tie(submat_map, std::ignore) = + gen_compressed_submat_map(basis_map, screened_shells, nbf, nbf); + } + + // dm_ao = D_compressed @ ao driver->eval_xmat( /*npts=*/static_cast(np), /*nbf=*/static_cast(nbf), - /*nbe=*/static_cast(nbf), submat_map, /*fac=*/1.0, + /*nbe=*/static_cast(nbe), submat_map, /*fac=*/1.0, /*P=*/D, /*ldp=*/static_cast(ldd), - /*basis_eval=*/ao_buf.data(), /*ldb=*/static_cast(nbf), - /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbf), + /*basis_eval=*/ao_buf.data(), /*ldb=*/static_cast(nbe), + /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbe), /*scr=*/xmat_scr.data()); // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) driver->eval_uvvar_lda_rks( - /*npts=*/static_cast(np), /*nbe=*/static_cast(nbf), + /*npts=*/static_cast(np), /*nbe=*/static_cast(nbe), /*basis_eval=*/ao_buf.data(), - /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbf), + /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbe), /*den_eval=*/out + p0); } } From 32e6a0ee14c99c78de5a49700f5b7d23802ad726 Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 18:46:41 +0000 Subject: [PATCH 03/31] OrbitalEvaluator/CubeGrid: add grid-native eval overloads 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 --- include/gauxc/external/cube.hpp | 8 ++ include/gauxc/orbital_evaluator.hpp | 21 +++ src/external/cube.cxx | 12 +- src/orbital_evaluator.cxx | 202 ++++++++++++++++++++++++++++ 4 files changed, 239 insertions(+), 4 deletions(-) diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp index 15ceb034..47037fe4 100644 --- a/include/gauxc/external/cube.hpp +++ b/include/gauxc/external/cube.hpp @@ -65,6 +65,14 @@ struct CubeGrid { * `eval_density`. */ std::vector points() const; + + /** @brief Write grid-point coordinates into a caller-supplied buffer. + * + * Same layout as `points()` but writes into `out`, which must have room + * for at least 3*num_points() doubles. Use this to avoid a heap + * allocation when the caller already owns a suitably sized buffer. + */ + void points_into(double* out) const; }; /** @brief Write a Gaussian cube file. diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index f0542e3e..04106266 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -16,6 +16,7 @@ #include #include +#include namespace GauXC { @@ -97,6 +98,26 @@ class OrbitalEvaluator { const double* D, int64_t ldd, double* out) const; + /** @brief Evaluate a single MO on a CubeGrid without materialising all 3*N + * grid-point coordinates. + */ + void eval_orbital(const CubeGrid& grid, + const double* C, double* out) const; + + /** @brief Evaluate `nmo` MOs on a CubeGrid without materialising all 3*N + * grid-point coordinates. + */ + void eval_orbitals(const CubeGrid& grid, + int32_t nmo, const double* C, int64_t ldc, + double* out, int64_t ldo) const; + + /** @brief Evaluate the electron density on a CubeGrid without materialising + * all 3*N grid-point coordinates. + */ + void eval_density(const CubeGrid& grid, + const double* D, int64_t ldd, + double* out) const; + private: struct Impl; std::unique_ptr pimpl_; diff --git a/src/external/cube.cxx b/src/external/cube.cxx index 018ad0c5..aead8206 100644 --- a/src/external/cube.cxx +++ b/src/external/cube.cxx @@ -71,6 +71,11 @@ CubeGrid CubeGrid::from_molecule(const Molecule& mol, int64_t nx, int64_t ny, std::vector CubeGrid::points() const { std::vector pts(static_cast(num_points()) * 3); + points_into(pts.data()); + return pts; +} + +void CubeGrid::points_into(double* out) const { size_t k = 0; for (int64_t ix = 0; ix < nx; ++ix) { const double x = origin[0] + spacing[0] * static_cast(ix); @@ -78,14 +83,13 @@ std::vector CubeGrid::points() const { const double y = origin[1] + spacing[1] * static_cast(iy); for (int64_t iz = 0; iz < nz; ++iz) { const double z = origin[2] + spacing[2] * static_cast(iz); - pts[3 * k + 0] = x; - pts[3 * k + 1] = y; - pts[3 * k + 2] = z; + out[3 * k + 0] = x; + out[3 * k + 1] = y; + out[3 * k + 2] = z; ++k; } } } - return pts; } // Hand-rolled %13.5E formatter. Required field is fixed 13 chars diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 1879f684..0e5ce22d 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -49,6 +50,22 @@ inline LocalHostWorkDriver::submat_map_t full_submat_map(int32_t nbf) { return {{ {int32_t{0}, nbf, int32_t{0}} }}; } +/// Fill `pts[3*np]` with grid-point coordinates for batch [p0, p0+np). +/// Avoids materialising the full 3*N point array for grid-based evaluation. +inline void fill_batch_pts(const CubeGrid& g, int64_t p0, int64_t np, + double* pts) { + const int64_t nz = g.nz, ny = g.ny; + for (int64_t i = 0; i < np; ++i) { + const int64_t k = p0 + i; + const int64_t iz = k % nz; + const int64_t iy = (k / nz) % ny; + const int64_t ix = k / (ny * nz); + pts[3 * i + 0] = g.origin[0] + static_cast(ix) * g.spacing[0]; + pts[3 * i + 1] = g.origin[1] + static_cast(iy) * g.spacing[1]; + pts[3 * i + 2] = g.origin[2] + static_cast(iz) * g.spacing[2]; + } +} + /// Axis-aligned bbox of a set of npts AoS points (length 3*npts). struct PointBbox { std::array lo; @@ -338,4 +355,189 @@ void OrbitalEvaluator::eval_density(int64_t npts, const double* points, } } +// --------------------------------------------------------------------------- +// CubeGrid overloads (3b): generate per-batch point coordinates on-the-fly, +// avoiding the 3*num_points()*8 byte temporary coordinate array. +// --------------------------------------------------------------------------- + +void OrbitalEvaluator::eval_orbital(const CubeGrid& grid, + const double* C, double* out) const { + eval_orbitals(grid, /*nmo=*/1, C, /*ldc=*/pimpl_->nbf_, out, + /*ldo=*/grid.num_points()); +} + +void OrbitalEvaluator::eval_orbitals(const CubeGrid& grid, + int32_t nmo, const double* C, int64_t ldc, + double* out, int64_t ldo) const { + const int64_t npts = grid.num_points(); + if (npts == 0 || nmo == 0) return; + if (C == nullptr || out == nullptr) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals(grid): null pointer argument."); + } + const int32_t nbf = pimpl_->nbf_; + if (ldc < nbf) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals(grid): ldc must be >= nbf()."); + } + if (ldo < npts) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals(grid): ldo must be >= npts."); + } + + const int32_t nshells_total = pimpl_->basis.nshells(); + const int64_t batch_size = choose_batch_size(nbf); + const int64_t n_batches = (npts + batch_size - 1) / batch_size; + LocalHostWorkDriver* driver = pimpl_->host_driver; + const BasisSet& basis = pimpl_->basis; + const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; + const BasisSetMap& basis_map = *pimpl_->basis_map; + +#pragma omp parallel + { + // batch_pts replaces the full 3*npts coordinate array; only batch_size + // coordinates are live at a time (~192 KB for batch_size=8192). + std::vector batch_pts(static_cast(batch_size) * 3); + std::vector ao_buf(static_cast(nbf) * batch_size); + std::vector C_compressed(static_cast(nbf) * nmo); + std::vector screened_shells; + screened_shells.reserve(nshells_total); + +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < n_batches; ++b) { + const int64_t p0 = b * batch_size; + const int64_t np = std::min(batch_size, npts - p0); + + fill_batch_pts(grid, p0, np, batch_pts.data()); + const PointBbox bbox = compute_bbox(batch_pts.data(), np); + screened_shells.clear(); + int32_t nbe = 0; + for (int32_t s = 0; s < nshells_total; ++s) { + if (dist2_center_to_bbox(basis[s].O_data(), bbox) < + shell_cutoff_r2[s]) { + screened_shells.push_back(s); + nbe += basis[s].size(); + } + } + if (nbe == 0) { + for (int32_t j = 0; j < nmo; ++j) { + double* out_col = out + static_cast(j) * ldo + p0; + std::fill(out_col, out_col + np, 0.0); + } + continue; + } + + driver->eval_collocation( + static_cast(np), + static_cast(screened_shells.size()), + static_cast(nbe), batch_pts.data(), basis, + screened_shells.data(), ao_buf.data()); + + { + int32_t row = 0; + for (int32_t s : screened_shells) { + const auto rng = basis_map.shell_to_ao_range(s); + const int32_t shell_nbe = rng.second - rng.first; + for (int32_t j = 0; j < nmo; ++j) { + const double* C_col = C + static_cast(j) * ldc; + double* dst = C_compressed.data() + + static_cast(j) * nbe + row; + std::copy(C_col + rng.first, C_col + rng.second, dst); + } + row += shell_nbe; + } + } + + blas::gemm( + 'T', 'N', static_cast(np), static_cast(nmo), nbe, + 1.0, ao_buf.data(), nbe, C_compressed.data(), nbe, + 0.0, out + p0, static_cast(ldo)); + } + } +} + +void OrbitalEvaluator::eval_density(const CubeGrid& grid, + const double* D, int64_t ldd, + double* out) const { + const int64_t npts = grid.num_points(); + if (npts == 0) return; + if (D == nullptr || out == nullptr) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_density(grid): null pointer argument."); + } + const int32_t nbf = pimpl_->nbf_; + if (ldd < nbf) { + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_density(grid): ldd must be >= nbf()."); + } + + const int32_t nshells_total = pimpl_->basis.nshells(); + const int64_t batch_size = choose_batch_size(nbf); + const int64_t n_batches = (npts + batch_size - 1) / batch_size; + LocalHostWorkDriver* driver = pimpl_->host_driver; + const BasisSet& basis = pimpl_->basis; + const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; + const BasisSetMap& basis_map = *pimpl_->basis_map; + +#pragma omp parallel + { + std::vector batch_pts(static_cast(batch_size) * 3); + std::vector ao_buf(static_cast(nbf) * batch_size); + std::vector dm_ao_buf(static_cast(nbf) * batch_size); + std::vector xmat_scr(static_cast(nbf) * nbf); + std::vector screened_shells; + screened_shells.reserve(nshells_total); + +#pragma omp for schedule(dynamic, 1) + for (int64_t b = 0; b < n_batches; ++b) { + const int64_t p0 = b * batch_size; + const int64_t np = std::min(batch_size, npts - p0); + + fill_batch_pts(grid, p0, np, batch_pts.data()); + const PointBbox bbox = compute_bbox(batch_pts.data(), np); + screened_shells.clear(); + int32_t nbe = 0; + for (int32_t s = 0; s < nshells_total; ++s) { + if (dist2_center_to_bbox(basis[s].O_data(), bbox) < + shell_cutoff_r2[s]) { + screened_shells.push_back(s); + nbe += basis[s].size(); + } + } + if (nbe == 0) { + std::fill(out + p0, out + p0 + np, 0.0); + continue; + } + + driver->eval_collocation( + static_cast(np), + static_cast(screened_shells.size()), + static_cast(nbe), batch_pts.data(), basis, + screened_shells.data(), ao_buf.data()); + + LocalHostWorkDriver::submat_map_t submat_map; + if (nbe == nbf) { + submat_map = full_submat_map(nbf); + } else { + std::tie(submat_map, std::ignore) = + gen_compressed_submat_map(basis_map, screened_shells, nbf, nbf); + } + + driver->eval_xmat( + static_cast(np), static_cast(nbf), + static_cast(nbe), submat_map, 1.0, + D, static_cast(ldd), + ao_buf.data(), static_cast(nbe), + dm_ao_buf.data(), static_cast(nbe), + xmat_scr.data()); + + driver->eval_uvvar_lda_rks( + static_cast(np), static_cast(nbe), + ao_buf.data(), + dm_ao_buf.data(), static_cast(nbe), + out + p0); + } + } +} + } // namespace GauXC From 9d417ca1c819ffff44cd6d57df6b0b3639146006 Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 19:20:47 +0000 Subject: [PATCH 04/31] OrbitalEvaluator: pipelined collocation+GEMM via OMP tasks 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%) --- src/orbital_evaluator.cxx | 114 +++++++++++++++++++++++++++++++++++++- 1 file changed, 112 insertions(+), 2 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 0e5ce22d..30ad0a29 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -18,6 +18,8 @@ #include #include +#include + #include #include #include @@ -393,10 +395,117 @@ void OrbitalEvaluator::eval_orbitals(const CubeGrid& grid, const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; const BasisSetMap& basis_map = *pimpl_->basis_map; + // Choose between task pipeline and omp-for based on saturation. + // When n_batches >> nthreads, omp-for already fully saturates all cores + // and the task pipeline's scheduling overhead is a net negative. + const int nthreads = omp_get_max_threads(); + const bool use_pipeline = (n_batches < 2 * nthreads); + + if (use_pipeline) { + // Pipelined evaluation (3d): overlap collocation of batch B+1 with + // GEMM of batch B using OMP tasks with dependency tags. + struct Slot { + std::vector pts; + std::vector ao; + std::vector C_comp; + std::vector shells; + int32_t nbe; + int64_t p0, np; + }; + + const int nslots = + std::min(nthreads, n_batches); + std::vector slots(nslots); + for (auto& s : slots) { + s.pts.resize(static_cast(batch_size) * 3); + s.ao.resize(static_cast(nbf) * batch_size); + s.C_comp.resize(static_cast(nbf) * nmo); + s.shells.reserve(nshells_total); + s.nbe = 0; + s.p0 = 0; + s.np = 0; + } + + // Dependency tag array — raw array for omp depend(). + std::unique_ptr dtag(new char[nslots]{}); + char* dtag_ptr = dtag.get(); + (void)dtag_ptr; // used only in omp depend() clauses + +#pragma omp parallel +#pragma omp single + { + for (int64_t b = 0; b < n_batches; ++b) { + const int si = static_cast(b % nslots); + + // Phase A: generate points, screen shells, eval collocation. +#pragma omp task shared(slots, grid, basis, shell_cutoff_r2) \ + depend(out: dtag_ptr[si]) firstprivate(b, si) + { + Slot& sl = slots[si]; + sl.p0 = b * batch_size; + sl.np = std::min(batch_size, npts - sl.p0); + + fill_batch_pts(grid, sl.p0, sl.np, sl.pts.data()); + const PointBbox bbox = compute_bbox(sl.pts.data(), sl.np); + sl.shells.clear(); + sl.nbe = 0; + for (int32_t s = 0; s < nshells_total; ++s) { + if (dist2_center_to_bbox(basis[s].O_data(), bbox) < + shell_cutoff_r2[s]) { + sl.shells.push_back(s); + sl.nbe += basis[s].size(); + } + } + if (sl.nbe > 0) { + driver->eval_collocation( + static_cast(sl.np), + static_cast(sl.shells.size()), + static_cast(sl.nbe), sl.pts.data(), basis, + sl.shells.data(), sl.ao.data()); + } + } + + // Phase B: gather C + GEMM → out. +#pragma omp task shared(slots, out, C, basis_map) \ + depend(in: dtag_ptr[si]) firstprivate(b, si) + { + Slot& sl = slots[si]; + if (sl.nbe == 0) { + for (int32_t j = 0; j < nmo; ++j) { + double* oc = out + static_cast(j) * ldo + sl.p0; + std::fill(oc, oc + sl.np, 0.0); + } + } else { + const int32_t nbe = sl.nbe; + { + int32_t row = 0; + for (int32_t s : sl.shells) { + const auto rng = basis_map.shell_to_ao_range(s); + const int32_t shell_nbe = rng.second - rng.first; + for (int32_t j = 0; j < nmo; ++j) { + const double* C_col = C + static_cast(j) * ldc; + double* dst = sl.C_comp.data() + + static_cast(j) * nbe + row; + std::copy(C_col + rng.first, C_col + rng.second, dst); + } + row += shell_nbe; + } + } + + blas::gemm( + 'T', 'N', static_cast(sl.np), static_cast(nmo), + nbe, 1.0, sl.ao.data(), nbe, sl.C_comp.data(), nbe, + 0.0, out + sl.p0, static_cast(ldo)); + } + } + } + } + + } else { + // Saturated path: plain omp-for, lower overhead at high thread counts. + #pragma omp parallel { - // batch_pts replaces the full 3*npts coordinate array; only batch_size - // coordinates are live at a time (~192 KB for batch_size=8192). std::vector batch_pts(static_cast(batch_size) * 3); std::vector ao_buf(static_cast(nbf) * batch_size); std::vector C_compressed(static_cast(nbf) * nmo); @@ -454,6 +563,7 @@ void OrbitalEvaluator::eval_orbitals(const CubeGrid& grid, 0.0, out + p0, static_cast(ldo)); } } + } // end saturated fallback } void OrbitalEvaluator::eval_density(const CubeGrid& grid, From 5dc6806f2854f99f296e9c5bba9d1729c61c8fbb Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 19:40:16 +0000 Subject: [PATCH 05/31] OrbitalEvaluator: hoist xmat_scr to per-evaluator scratch pool 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. --- src/orbital_evaluator.cxx | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 30ad0a29..8f909422 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -120,6 +120,20 @@ struct OrbitalEvaluator::Impl { std::vector shell_cutoff_r2; int32_t nbf_ = 0; + // Per-thread scratch for eval_xmat in eval_density. Allocated once + // per evaluator rather than once per eval_density call, saving + // nthreads * nbf^2 * 8 bytes of allocation churn per call. + mutable std::vector> xmat_scratch_pool; + + /// Return a pointer to thread-local xmat scratch of size nbf*nbf. + /// Must be called from inside an omp parallel region whose thread + /// count does not exceed the value of omp_get_max_threads() at + /// construction time. + double* xmat_scratch() const { + auto& buf = xmat_scratch_pool[omp_get_thread_num()]; + return buf.data(); + } + void init(BasisSet bs, ExecutionSpace exec) { if (exec != ExecutionSpace::Host) { GAUXC_GENERIC_EXCEPTION( @@ -148,6 +162,12 @@ struct OrbitalEvaluator::Impl { const double r = basis[s].cutoff_radius(); shell_cutoff_r2[s] = r * r; } + + // Pre-size the scratch pool for the current OMP thread count. + xmat_scratch_pool.resize(omp_get_max_threads()); + for (auto& buf : xmat_scratch_pool) { + buf.resize(static_cast(nbf_) * nbf_); + } } }; @@ -295,9 +315,6 @@ void OrbitalEvaluator::eval_density(int64_t npts, const double* points, { std::vector ao_buf(static_cast(nbf) * batch_size); std::vector dm_ao_buf(static_cast(nbf) * batch_size); - // eval_xmat scratch when the submat-gather path is taken; sized for - // worst case (nbe == nbf). - std::vector xmat_scr(static_cast(nbf) * nbf); std::vector screened_shells; screened_shells.reserve(nshells_total); @@ -345,7 +362,7 @@ void OrbitalEvaluator::eval_density(int64_t npts, const double* points, /*P=*/D, /*ldp=*/static_cast(ldd), /*basis_eval=*/ao_buf.data(), /*ldb=*/static_cast(nbe), /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbe), - /*scr=*/xmat_scr.data()); + /*scr=*/pimpl_->xmat_scratch()); // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) driver->eval_uvvar_lda_rks( @@ -594,7 +611,6 @@ void OrbitalEvaluator::eval_density(const CubeGrid& grid, std::vector batch_pts(static_cast(batch_size) * 3); std::vector ao_buf(static_cast(nbf) * batch_size); std::vector dm_ao_buf(static_cast(nbf) * batch_size); - std::vector xmat_scr(static_cast(nbf) * nbf); std::vector screened_shells; screened_shells.reserve(nshells_total); @@ -639,7 +655,7 @@ void OrbitalEvaluator::eval_density(const CubeGrid& grid, D, static_cast(ldd), ao_buf.data(), static_cast(nbe), dm_ao_buf.data(), static_cast(nbe), - xmat_scr.data()); + pimpl_->xmat_scratch()); driver->eval_uvvar_lda_rks( static_cast(np), static_cast(nbe), From 40afaaf4439759c969ec426d87ffa664db9e53ff Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 20:21:52 +0000 Subject: [PATCH 06/31] Add write_cube_hdf5 for HDF5 cube field output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- include/gauxc/external/cube.hpp | 30 ++++++++++++ src/external/CMakeLists.txt | 2 +- src/external/cube_hdf5.cxx | 81 ++++++++++++++++++++++++++++++++ tests/orbital_evaluator_test.cxx | 78 ++++++++++++++++++++++++++++++ 4 files changed, 190 insertions(+), 1 deletion(-) create mode 100644 src/external/cube_hdf5.cxx diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp index 47037fe4..0cdd8f7f 100644 --- a/include/gauxc/external/cube.hpp +++ b/include/gauxc/external/cube.hpp @@ -100,4 +100,34 @@ void write_cube(const std::string& path, const double* field, const std::string& comment = ""); +#ifdef GAUXC_HAS_HDF5 +/** @brief Write a cube field to an HDF5 file. + * + * Stores the grid specification, molecular geometry, and scalar field in a + * single HDF5 file for downstream analysis. The field is stored as a 3D + * dataset of shape (nx, ny, nz) in row-major order with iz varying fastest + * (matching the cube-file convention). + * + * HDF5 layout: + * /field (nx, ny, nz) float64 — the scalar field + * /grid/origin (3,) float64 + * /grid/spacing (3,) float64 + * /grid/shape (3,) int64 — {nx, ny, nz} + * /atoms/Z (natom,) int64 + * /atoms/coords (natom, 3) float64 + * /comment scalar string attribute on /field + * + * @param path Output path. Parent directory must exist. + * @param mol Molecule (atomic numbers + Cartesian coordinates in Bohr). + * @param grid Grid specification. + * @param field Length-`grid.num_points()` scalar field. + * @param comment Optional comment stored as an attribute on /field. + */ +void write_cube_hdf5(const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = ""); +#endif + } // namespace GauXC diff --git a/src/external/CMakeLists.txt b/src/external/CMakeLists.txt index b0add283..9df2df95 100644 --- a/src/external/CMakeLists.txt +++ b/src/external/CMakeLists.txt @@ -35,7 +35,7 @@ if( GAUXC_ENABLE_HDF5 ) FetchContent_MakeAvailable( HighFive ) endif() - target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx ) + target_sources( gauxc PRIVATE hdf5_write.cxx hdf5_read.cxx cube_hdf5.cxx ) target_link_libraries( gauxc PUBLIC HighFive ) else() message(WARNING "GAUXC_ENABLE_HDF5 was enabled, but HDF5 was not found, Disabling HDF5 Bindings") diff --git a/src/external/cube_hdf5.cxx b/src/external/cube_hdf5.cxx new file mode 100644 index 00000000..13b8f0d1 --- /dev/null +++ b/src/external/cube_hdf5.cxx @@ -0,0 +1,81 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2026, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include +#ifdef GAUXC_HAS_HDF5 + +#include +#include + +#include + +#include + +namespace GauXC { + +void write_cube_hdf5(const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment) { + if (grid.num_points() <= 0) { + GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: grid has zero points."); + } + if (field == nullptr) { + GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: null field pointer."); + } + + HighFive::File file(path, HighFive::File::Overwrite); + + // --- Field: 3D dataset (nx, ny, nz) --- + const auto nx = static_cast(grid.nx); + const auto ny = static_cast(grid.ny); + const auto nz = static_cast(grid.nz); + + auto ds_field = file.createDataSet( + "field", HighFive::DataSpace({nx, ny, nz})); + ds_field.write_raw(field); + + // Comment attribute on /field. + const std::string cmt = + comment.empty() ? "Cube field generated by GauXC" : comment; + ds_field.createAttribute("comment", cmt); + + // --- Grid metadata --- + auto grp_grid = file.createGroup("grid"); + grp_grid.createDataSet("origin", + std::vector(grid.origin.begin(), grid.origin.end())); + grp_grid.createDataSet("spacing", + std::vector(grid.spacing.begin(), grid.spacing.end())); + grp_grid.createDataSet("shape", + std::vector{grid.nx, grid.ny, grid.nz}); + + // --- Atoms --- + const size_t natom = mol.size(); + auto grp_atoms = file.createGroup("atoms"); + + std::vector atomic_numbers(natom); + std::vector coords(natom * 3); + for (size_t i = 0; i < natom; ++i) { + atomic_numbers[i] = static_cast(mol[i].Z.get()); + coords[3 * i + 0] = mol[i].x; + coords[3 * i + 1] = mol[i].y; + coords[3 * i + 2] = mol[i].z; + } + + grp_atoms.createDataSet("Z", atomic_numbers); + grp_atoms.createDataSet("coords", + HighFive::DataSpace({natom, 3ul})).write_raw(coords.data()); +} + +} // namespace GauXC + +#endif // GAUXC_HAS_HDF5 diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 97562323..82971187 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -24,6 +24,10 @@ #include #include #include +#include +#ifdef GAUXC_HAS_HDF5 +#include +#endif #include "standards.hpp" @@ -330,3 +334,77 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { std::remove(path.c_str()); } + +#ifdef GAUXC_HAS_HDF5 +TEST_CASE("write_cube_hdf5 round-trip", "[cube]") { + auto mol = make_water(); + auto grid = CubeGrid::from_molecule(mol, 4, 5, 6); + + // Fill a small field with known values. + const int64_t npts = grid.num_points(); + std::vector field(npts); + for (int64_t i = 0; i < npts; ++i) field[i] = 0.01 * i - 0.5; + + const std::string path = "/tmp/_gauxc_test_cube.h5"; + write_cube_hdf5(path, mol, grid, field.data(), "test cube"); + + // Read back and verify. + HighFive::File file(path, HighFive::File::ReadOnly); + + // Field shape and values. + auto ds = file.getDataSet("field"); + auto dims = ds.getDimensions(); + REQUIRE(dims.size() == 3); + CHECK(dims[0] == static_cast(grid.nx)); + CHECK(dims[1] == static_cast(grid.ny)); + CHECK(dims[2] == static_cast(grid.nz)); + + std::vector read_field(npts); + ds.read(read_field.data()); + for (int64_t i = 0; i < npts; ++i) { + CHECK(read_field[i] == Approx(field[i]).epsilon(1e-14)); + } + + // Comment attribute. + std::string cmt; + ds.getAttribute("comment").read(cmt); + CHECK(cmt == "test cube"); + + // Grid metadata. + auto grp_grid = file.getGroup("grid"); + std::vector origin, spacing; + std::vector shape; + grp_grid.getDataSet("origin").read(origin); + grp_grid.getDataSet("spacing").read(spacing); + grp_grid.getDataSet("shape").read(shape); + REQUIRE(origin.size() == 3); + REQUIRE(spacing.size() == 3); + REQUIRE(shape.size() == 3); + for (int k = 0; k < 3; ++k) { + CHECK(origin[k] == Approx(grid.origin[k]).epsilon(1e-14)); + CHECK(spacing[k] == Approx(grid.spacing[k]).epsilon(1e-14)); + } + CHECK(shape[0] == grid.nx); + CHECK(shape[1] == grid.ny); + CHECK(shape[2] == grid.nz); + + // Atoms. + auto grp_atoms = file.getGroup("atoms"); + std::vector Z; + grp_atoms.getDataSet("Z").read(Z); + REQUIRE(Z.size() == mol.size()); + for (size_t i = 0; i < mol.size(); ++i) { + CHECK(Z[i] == static_cast(mol[i].Z.get())); + } + + std::vector coords(mol.size() * 3); + grp_atoms.getDataSet("coords").read(coords.data()); + for (size_t i = 0; i < mol.size(); ++i) { + CHECK(coords[3 * i + 0] == Approx(mol[i].x).epsilon(1e-14)); + CHECK(coords[3 * i + 1] == Approx(mol[i].y).epsilon(1e-14)); + CHECK(coords[3 * i + 2] == Approx(mol[i].z).epsilon(1e-14)); + } + + std::remove(path.c_str()); +} +#endif From 10f59a8f42ae5f66e72c287831980a6448e36aa2 Mon Sep 17 00:00:00 2001 From: Cody Johnston Date: Wed, 29 Apr 2026 20:56:16 +0000 Subject: [PATCH 07/31] Fix portability issues and add CubeGrid overload tests - Guard #include 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. --- src/orbital_evaluator.cxx | 9 ++++- tests/orbital_evaluator_test.cxx | 67 ++++++++++++++++++++++++++++---- 2 files changed, 66 insertions(+), 10 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 8f909422..05e7064a 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -18,7 +18,12 @@ #include #include +#ifdef _OPENMP #include +#else +inline int omp_get_max_threads() { return 1; } +inline int omp_get_thread_num() { return 0; } +#endif #include #include @@ -375,7 +380,7 @@ void OrbitalEvaluator::eval_density(int64_t npts, const double* points, } // --------------------------------------------------------------------------- -// CubeGrid overloads (3b): generate per-batch point coordinates on-the-fly, +// CubeGrid overloads: generate per-batch point coordinates on-the-fly, // avoiding the 3*num_points()*8 byte temporary coordinate array. // --------------------------------------------------------------------------- @@ -419,7 +424,7 @@ void OrbitalEvaluator::eval_orbitals(const CubeGrid& grid, const bool use_pipeline = (n_batches < 2 * nthreads); if (use_pipeline) { - // Pipelined evaluation (3d): overlap collocation of batch B+1 with + // Pipelined evaluation: overlap collocation of batch B+1 with // GEMM of batch B using OMP tasks with dependency tags. struct Slot { std::vector pts; diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 82971187..de7a6118 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -17,10 +17,9 @@ #include #include #include +#include #include -#include - #include #include #include @@ -38,6 +37,12 @@ using namespace GauXC; namespace { +/// Generate a unique temp path for test files (avoids tmpnam warning). +std::string make_temp_path(const char* suffix) { + static int counter = 0; + return std::string("/tmp/gauxc_test_") + std::to_string(++counter) + suffix; +} + /// Build a deterministic PRNG-based set of points within a small bounding box /// around the molecule. Avoids putting samples too close to nuclei to keep /// values numerically well-behaved. @@ -169,6 +174,54 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", } } +TEST_CASE("CubeGrid eval overloads match pointer-based eval", + "[orbital_evaluator]") { + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + const int32_t nbf = basis.nbf(); + OrbitalEvaluator eval(basis); + + // Small grid for fast testing. + auto grid = CubeGrid::from_molecule(mol, 6, 7, 8); + const int64_t npts = grid.num_points(); + auto pts = grid.points(); + + // Random MO coefficients. + std::mt19937 rng(42u); + std::uniform_real_distribution dist(-1.0, 1.0); + std::vector C(static_cast(nbf)); + for (auto& v : C) v = dist(rng); + + SECTION("eval_orbital(grid) matches eval_orbital(npts, points)") { + std::vector ref(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); + + std::vector out(static_cast(npts)); + eval.eval_orbital(grid, C.data(), out.data()); + + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } + } + + SECTION("eval_density(grid) matches eval_density(npts, points)") { + // Identity density. + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector ref(static_cast(npts)); + eval.eval_density(npts, pts.data(), D.data(), nbf, ref.data()); + + std::vector out(static_cast(npts)); + eval.eval_density(grid, D.data(), nbf, out.data()); + + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } + } +} + TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { auto mol = make_water(); CubeGrid g = CubeGrid::from_molecule(mol, /*nx=*/16, /*ny=*/12, /*nz=*/8, @@ -213,9 +266,8 @@ TEST_CASE("write_cube round-trips header and field data", "[cube]") { std::pow(10.0, (i % 5) - 2); } - // Use a tmp path under /tmp. - const std::string path = std::string("/tmp/gauxc_cube_test_") + - std::to_string(::getpid()) + ".cube"; + // Use a tmp path. + const std::string path = make_temp_path(".cube"); write_cube(path, mol, grid, field.data(), "Test cube"); // Parse back. @@ -303,8 +355,7 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { -2.71828, 1.0, -1.0, 1e-300, 1.234e+05}; - const std::string path = std::string("/tmp/gauxc_cube_format_") + - std::to_string(::getpid()) + ".cube"; + const std::string path = make_temp_path(".cube"); write_cube(path, mol, grid, field.data(), "fmt"); std::ifstream in(path); @@ -345,7 +396,7 @@ TEST_CASE("write_cube_hdf5 round-trip", "[cube]") { std::vector field(npts); for (int64_t i = 0; i < npts; ++i) field[i] = 0.01 * i - 0.5; - const std::string path = "/tmp/_gauxc_test_cube.h5"; + const std::string path = make_temp_path(".h5"); write_cube_hdf5(path, mol, grid, field.data(), "test cube"); // Read back and verify. From 2a1f46e9a658151ea53850d01053b286e9ff5582 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Mon, 4 May 2026 10:46:41 -0700 Subject: [PATCH 08/31] Tests: guard cube-file I/O cases against MPI rank collisions 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). --- tests/orbital_evaluator_test.cxx | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index de7a6118..30ec1f54 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -20,6 +20,8 @@ #include #include +#include // getpid + #include #include #include @@ -27,6 +29,9 @@ #ifdef GAUXC_HAS_HDF5 #include #endif +#ifdef GAUXC_HAS_MPI +#include +#endif #include "standards.hpp" @@ -38,9 +43,29 @@ using namespace GauXC; namespace { /// Generate a unique temp path for test files (avoids tmpnam warning). +/// Includes the PID so concurrent MPI ranks (which share /tmp on a single +/// node) do not collide on the same file. std::string make_temp_path(const char* suffix) { static int counter = 0; - return std::string("/tmp/gauxc_test_") + std::to_string(++counter) + suffix; + return std::string("/tmp/gauxc_test_") + std::to_string(::getpid()) + + "_" + std::to_string(++counter) + suffix; +} + +/// Returns true on the root MPI rank (or always when MPI is disabled). +/// I/O test cases below run only on rank 0 to avoid concurrent ranks racing +/// on the same /tmp file paths during MPI test runs. +bool is_root_rank() { +#ifdef GAUXC_HAS_MPI + int rank = 0; + int initialized = 0; + MPI_Initialized(&initialized); + if (initialized) { + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + } + return rank == 0; +#else + return true; +#endif } /// Build a deterministic PRNG-based set of points within a small bounding box @@ -256,6 +281,7 @@ TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { } TEST_CASE("write_cube round-trips header and field data", "[cube]") { + if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. auto mol = make_water(); CubeGrid grid = CubeGrid::from_molecule(mol, /*nx=*/5, /*ny=*/4, /*nz=*/7, /*margin=*/2.0); @@ -339,6 +365,7 @@ TEST_CASE("write_cube round-trips header and field data", "[cube]") { } TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { + if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. // Spot-check the custom formatter against snprintf for a battery of values // by writing a tiny cube file and parsing it back. This is a stronger check // than the round-trip above since we compare the byte stream. @@ -388,6 +415,7 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { #ifdef GAUXC_HAS_HDF5 TEST_CASE("write_cube_hdf5 round-trip", "[cube]") { + if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. auto mol = make_water(); auto grid = CubeGrid::from_molecule(mol, 4, 5, 6); From e2dac06ce1777c6c122b98d3e65bc5f4315ef8f9 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Fri, 31 Jul 2026 11:51:55 -0700 Subject: [PATCH 09/31] Address PR review: fix eval_density scratch race, dedupe batch kernel 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. --- include/gauxc/cube_grid.hpp | 77 +++ include/gauxc/external/cube.hpp | 81 +-- include/gauxc/orbital_evaluator.hpp | 116 ++-- src/CMakeLists.txt | 1 + src/cube_grid.cxx | 77 +++ src/external/cube.cxx | 297 +++------ src/external/cube_hdf5.cxx | 18 +- src/orbital_evaluator.cxx | 948 ++++++++++++---------------- src/orbital_evaluator_impl.hpp | 42 ++ tests/CMakeLists.txt | 3 + tests/orbital_evaluator_test.cxx | 508 ++++++++++++--- tests/ut_common.hpp.in | 1 + 12 files changed, 1222 insertions(+), 947 deletions(-) create mode 100644 include/gauxc/cube_grid.hpp create mode 100644 src/cube_grid.cxx create mode 100644 src/orbital_evaluator_impl.hpp diff --git a/include/gauxc/cube_grid.hpp b/include/gauxc/cube_grid.hpp new file mode 100644 index 00000000..6b01f112 --- /dev/null +++ b/include/gauxc/cube_grid.hpp @@ -0,0 +1,77 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include +#include + +#include + +namespace GauXC { + +/** @brief Specification of an axis-aligned 3D rectangular point grid. + * + * The grid axes are parallel to the Cartesian frame, i.e. the axis vectors + * are (spacing[0], 0, 0), (0, spacing[1], 0) and (0, 0, spacing[2]). All + * quantities are in atomic units (Bohr). + * + * Total number of points: nx*ny*nz. Storage cost per scalar field: + * 8*nx*ny*nz bytes (double precision). + */ +struct CubeGrid { + std::array origin{0.0, 0.0, 0.0}; ///< (0,0,0) corner, Bohr + std::array spacing{0.2, 0.2, 0.2}; ///< Step on each axis, Bohr + int64_t nx = 80; + int64_t ny = 80; + int64_t nz = 80; + + /// Total number of grid points. + int64_t num_points() const noexcept { return nx * ny * nz; } + + /** @brief Build a default grid that tightly encloses a molecule. + * + * The bounding box of the atomic centres is extended by `margin` Bohr on + * each side and discretised with the requested number of points along + * each axis. Spacing is chosen so that the first and last grid points + * coincide with the extended bounding-box corners (PySCF cubegen + * convention). + * + * @param mol Molecule whose atomic centres define the bounding box. + * @param nx,ny,nz Number of grid points along each axis. + * @param margin Margin (Bohr) added on each side. Default 3.0 matches + * the PySCF cubegen default. + */ + static CubeGrid from_molecule( const Molecule& mol, + int64_t nx = 80, int64_t ny = 80, + int64_t nz = 80, double margin = 3.0 ); + + /** @brief Materialise the grid points as an AoS coordinate array. + * + * Returns a vector of length 3*num_points() laid out in row-major + * (ix, iy, iz) order with iz varying fastest, matching the field-element + * ordering expected by `write_cube`. Suitable to be passed directly as + * the `points` argument to `OrbitalEvaluator::eval_orbital` / + * `eval_density`. + */ + std::vector points() const; + + /** @brief Write grid-point coordinates into a caller-supplied buffer. + * + * Same layout as `points()` but writes into `out`, which must have room + * for at least 3*num_points() doubles. Use this to avoid a heap + * allocation when the caller already owns a suitably sized buffer. + */ + void points_into( double* out ) const; +}; + +} // namespace GauXC diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp index 0cdd8f7f..78d29611 100644 --- a/include/gauxc/external/cube.hpp +++ b/include/gauxc/external/cube.hpp @@ -3,7 +3,7 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * @@ -11,70 +11,13 @@ */ #pragma once -#include -#include #include -#include +#include #include namespace GauXC { -/** @brief Specification of a Gaussian-cube-format 3D rectangular grid. - * - * The grid is axis-aligned with the Cartesian frame (the standard - * cube-format axis vectors are simply (spacing[0], 0, 0), (0, spacing[1], 0), - * (0, 0, spacing[2])). All quantities are in atomic units (Bohr). - * - * Total number of points: nx*ny*nz. Storage cost per scalar field: - * 8*nx*ny*nz bytes (double precision). - */ -struct CubeGrid { - std::array origin{0.0, 0.0, 0.0}; ///< (0,0,0) corner, Bohr - std::array spacing{0.2, 0.2, 0.2}; ///< Step on each axis, Bohr - int64_t nx = 80; - int64_t ny = 80; - int64_t nz = 80; - - /// Total number of grid points. - int64_t num_points() const noexcept { return nx * ny * nz; } - - /** @brief Build a default grid that tightly encloses a molecule. - * - * The bounding box of the atomic centres is extended by `margin` Bohr on - * each side and discretised with the requested number of points along - * each axis. Spacing is chosen so that the first and last grid points - * coincide with the extended bounding-box corners (PySCF cubegen - * convention). - * - * @param mol Molecule whose atomic centres define the bounding box. - * @param nx,ny,nz Number of grid points along each axis. - * @param margin Margin (Bohr) added on each side. Default 3.0 matches - * the PySCF cubegen default. - */ - static CubeGrid from_molecule(const Molecule& mol, - int64_t nx = 80, int64_t ny = 80, - int64_t nz = 80, double margin = 3.0); - - /** @brief Materialise the grid points as an AoS coordinate array. - * - * Returns a vector of length 3*num_points() laid out in row-major - * (ix, iy, iz) order with iz varying fastest, matching the field-element - * ordering expected by `write_cube`. Suitable to be passed directly as - * the `points` argument to `OrbitalEvaluator::eval_orbital` / - * `eval_density`. - */ - std::vector points() const; - - /** @brief Write grid-point coordinates into a caller-supplied buffer. - * - * Same layout as `points()` but writes into `out`, which must have room - * for at least 3*num_points() doubles. Use this to avoid a heap - * allocation when the caller already owns a suitably sized buffer. - */ - void points_into(double* out) const; -}; - /** @brief Write a Gaussian cube file. * * Field layout: row-major (ix, iy, iz) with iz varying fastest. Length must @@ -94,11 +37,11 @@ struct CubeGrid { * always "Generated by GauXC"). If empty, a default first * line is used. */ -void write_cube(const std::string& path, - const Molecule& mol, - const CubeGrid& grid, - const double* field, - const std::string& comment = ""); +void write_cube( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = "" ); #ifdef GAUXC_HAS_HDF5 /** @brief Write a cube field to an HDF5 file. @@ -123,11 +66,11 @@ void write_cube(const std::string& path, * @param field Length-`grid.num_points()` scalar field. * @param comment Optional comment stored as an attribute on /field. */ -void write_cube_hdf5(const std::string& path, - const Molecule& mol, - const CubeGrid& grid, - const double* field, - const std::string& comment = ""); +void write_cube_hdf5( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment = "" ); #endif } // namespace GauXC diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index 04106266..a510ee85 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -3,7 +3,7 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * @@ -11,18 +11,24 @@ */ #pragma once +#include #include #include #include +#include #include -#include namespace GauXC { +namespace detail { + /// OrbitalEvaluator Implementation class + class OrbitalEvaluatorImpl; +} + /** @brief Evaluate molecular orbitals and densities on arbitrary point sets. * - * Wraps the host collocation kernel exposed by the local work driver with a + * Wraps the collocation kernel exposed by the local work driver with a * thread-parallel batched evaluation loop, hiding the driver factory and * the per-thread AO scratch from callers. The class is designed to be * constructed once per molecule and reused across many evaluations (e.g. @@ -32,34 +38,41 @@ namespace GauXC { * it returns plain numerical arrays. For Gaussian cube file I/O, see * `gauxc/external/cube.hpp`. * - * Currently only `ExecutionSpace::Host` is supported; passing any other - * value will throw on construction. Device support can be added later by - * routing through a device-side collocation driver while keeping the same - * signatures. + * Points are evaluated in batches and each batch is screened against the + * per-shell cutoff radii. The batch size is derived in part from the OpenMP + * thread count, so results are bit-reproducible for a fixed thread count but + * may differ between thread counts by at most the basis-set shell tolerance. + * + * Instances are produced by OrbitalEvaluatorFactory, which selects the + * implementation for a given ExecutionSpace. */ class OrbitalEvaluator { - public: - /** @brief Construct an evaluator bound to a basis set. - * - * @param basis Basis set (copied internally). - * @param exec Execution space; only `ExecutionSpace::Host` is supported. - */ - explicit OrbitalEvaluator(BasisSet basis, - ExecutionSpace exec = ExecutionSpace::Host); + + using pimpl_type = detail::OrbitalEvaluatorImpl; + using pimpl_ptr_type = std::unique_ptr; + pimpl_ptr_type pimpl_; ///< Pointer to implementation instance + +public: + + // Delete default ctor + OrbitalEvaluator() = delete; + + /// Construct an OrbitalEvaluator instance from a preconstructed implementation + OrbitalEvaluator( pimpl_ptr_type&& pimpl ); ~OrbitalEvaluator() noexcept; // Non-copyable, movable - OrbitalEvaluator(const OrbitalEvaluator&) = delete; - OrbitalEvaluator& operator=(const OrbitalEvaluator&) = delete; - OrbitalEvaluator(OrbitalEvaluator&&) noexcept; - OrbitalEvaluator& operator=(OrbitalEvaluator&&) noexcept; + OrbitalEvaluator( const OrbitalEvaluator& ) = delete; + OrbitalEvaluator& operator=( const OrbitalEvaluator& ) = delete; + OrbitalEvaluator( OrbitalEvaluator&& ) noexcept; + OrbitalEvaluator& operator=( OrbitalEvaluator&& ) noexcept; /// Number of basis functions (rows of the AO matrix). - int32_t nbf() const noexcept; + int32_t nbf() const; /// Underlying basis set. - const BasisSet& basis() const noexcept; + const BasisSet& basis() const; /** @brief Evaluate a single MO chi(r) = sum_mu C[mu] * phi_mu(r). * @@ -69,8 +82,8 @@ class OrbitalEvaluator { * @param[in] C MO coefficient vector, length nbf(). * @param[out] out Length-npts array of MO values. */ - void eval_orbital(int64_t npts, const double* points, - const double* C, double* out) const; + void eval_orbital( size_t npts, const double* points, + const double* C, double* out ) const; /** @brief Evaluate `nmo` MOs simultaneously. * @@ -81,9 +94,9 @@ class OrbitalEvaluator { * Equivalent to calling `eval_orbital` `nmo` times but amortises the AO * collocation evaluation across all MOs (single AO buffer, GEMM contraction). */ - void eval_orbitals(int64_t npts, const double* points, - int32_t nmo, const double* C, int64_t ldc, - double* out, int64_t ldo) const; + void eval_orbitals( size_t npts, const double* points, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const; /** @brief Evaluate the electron density * rho(r) = sum_{mu,nu} D[mu,nu] * phi_mu(r) * phi_nu(r). @@ -94,33 +107,52 @@ class OrbitalEvaluator { * leading dimension `ldd` (>= nbf). * @param[out] out Length-npts array of density values. */ - void eval_density(int64_t npts, const double* points, - const double* D, int64_t ldd, - double* out) const; + void eval_density( size_t npts, const double* points, + const double* D, size_t ldd, + double* out ) const; /** @brief Evaluate a single MO on a CubeGrid without materialising all 3*N * grid-point coordinates. */ - void eval_orbital(const CubeGrid& grid, - const double* C, double* out) const; + void eval_orbital( const CubeGrid& grid, + const double* C, double* out ) const; /** @brief Evaluate `nmo` MOs on a CubeGrid without materialising all 3*N * grid-point coordinates. */ - void eval_orbitals(const CubeGrid& grid, - int32_t nmo, const double* C, int64_t ldc, - double* out, int64_t ldo) const; + void eval_orbitals( const CubeGrid& grid, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const; /** @brief Evaluate the electron density on a CubeGrid without materialising * all 3*N grid-point coordinates. */ - void eval_density(const CubeGrid& grid, - const double* D, int64_t ldd, - double* out) const; - - private: - struct Impl; - std::unique_ptr pimpl_; -}; + void eval_density( const CubeGrid& grid, + const double* D, size_t ldd, + double* out ) const; + +}; // class OrbitalEvaluator + + +/// A factory to generate OrbitalEvaluator instances +class OrbitalEvaluatorFactory { + +public: + + // Delete default ctor + OrbitalEvaluatorFactory() = delete; + + /** + * @brief Construct an OrbitalEvaluator for a given execution space + * + * @param[in] ex Execution space in which to evaluate orbitals/densities. + * Currently only ExecutionSpace::Host is implemented; + * anything else throws. + * @param[in] basis Basis set (copied into the evaluator). + */ + static OrbitalEvaluator make_orbital_evaluator( ExecutionSpace ex, + BasisSet basis ); + +}; // class OrbitalEvaluatorFactory } // namespace GauXC diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8cd463ce..0ef34edb 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -41,6 +41,7 @@ add_library( gauxc molgrid_impl.cxx molgrid_defaults.cxx atomic_radii.cxx + cube_grid.cxx orbital_evaluator.cxx ) diff --git a/src/cube_grid.cxx b/src/cube_grid.cxx new file mode 100644 index 00000000..bd59c1a3 --- /dev/null +++ b/src/cube_grid.cxx @@ -0,0 +1,77 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#include + +#include + +#include + +namespace GauXC { + +CubeGrid CubeGrid::from_molecule( const Molecule& mol, int64_t nx, int64_t ny, + int64_t nz, double margin ) { + if( mol.empty() ) { + GAUXC_GENERIC_EXCEPTION("CubeGrid::from_molecule: molecule has no atoms."); + } + if( nx < 1 or ny < 1 or nz < 1 ) { + GAUXC_GENERIC_EXCEPTION("CubeGrid::from_molecule: nx, ny, nz must be >= 1."); + } + + double xmin = mol[0].x, xmax = mol[0].x; + double ymin = mol[0].y, ymax = mol[0].y; + double zmin = mol[0].z, zmax = mol[0].z; + for( const auto& a : mol ) { + xmin = std::min(xmin, a.x); + xmax = std::max(xmax, a.x); + ymin = std::min(ymin, a.y); + ymax = std::max(ymax, a.y); + zmin = std::min(zmin, a.z); + zmax = std::max(zmax, a.z); + } + + CubeGrid grid; + grid.origin = {xmin - margin, ymin - margin, zmin - margin}; + grid.nx = nx; + grid.ny = ny; + grid.nz = nz; + const double ex = (xmax - xmin) + 2.0 * margin; + const double ey = (ymax - ymin) + 2.0 * margin; + const double ez = (zmax - zmin) + 2.0 * margin; + grid.spacing[0] = nx > 1 ? ex / static_cast(nx - 1) : 0.0; + grid.spacing[1] = ny > 1 ? ey / static_cast(ny - 1) : 0.0; + grid.spacing[2] = nz > 1 ? ez / static_cast(nz - 1) : 0.0; + return grid; +} + +std::vector CubeGrid::points() const { + std::vector pts(static_cast(num_points()) * 3); + points_into(pts.data()); + return pts; +} + +void CubeGrid::points_into( double* out ) const { + size_t k = 0; + for( int64_t ix = 0; ix < nx; ++ix ) { + const double x = origin[0] + spacing[0] * static_cast(ix); + for( int64_t iy = 0; iy < ny; ++iy ) { + const double y = origin[1] + spacing[1] * static_cast(iy); + for( int64_t iz = 0; iz < nz; ++iz ) { + out[3 * k + 0] = x; + out[3 * k + 1] = y; + out[3 * k + 2] = origin[2] + spacing[2] * static_cast(iz); + ++k; + } + } + } +} + +} // namespace GauXC diff --git a/src/external/cube.cxx b/src/external/cube.cxx index aead8206..626169ee 100644 --- a/src/external/cube.cxx +++ b/src/external/cube.cxx @@ -3,7 +3,7 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * @@ -14,156 +14,78 @@ #include #include #include -#include #include -#include +#include #include #include -#ifdef _OPENMP -#include -#endif - #include namespace GauXC { -// ============================================================================= -// CubeGrid -// ============================================================================= - -CubeGrid CubeGrid::from_molecule(const Molecule& mol, int64_t nx, int64_t ny, - int64_t nz, double margin) { - if (mol.empty()) { - GAUXC_GENERIC_EXCEPTION( - "CubeGrid::from_molecule: molecule has no atoms."); - } - if (nx < 1 || ny < 1 || nz < 1) { - GAUXC_GENERIC_EXCEPTION( - "CubeGrid::from_molecule: nx, ny, nz must be >= 1."); - } - - double xmin = mol[0].x, xmax = mol[0].x; - double ymin = mol[0].y, ymax = mol[0].y; - double zmin = mol[0].z, zmax = mol[0].z; - for (const auto& a : mol) { - xmin = std::min(xmin, a.x); - xmax = std::max(xmax, a.x); - ymin = std::min(ymin, a.y); - ymax = std::max(ymax, a.y); - zmin = std::min(zmin, a.z); - zmax = std::max(zmax, a.z); - } - - CubeGrid grid; - grid.origin = {xmin - margin, ymin - margin, zmin - margin}; - grid.nx = nx; - grid.ny = ny; - grid.nz = nz; - const double ex = (xmax - xmin) + 2.0 * margin; - const double ey = (ymax - ymin) + 2.0 * margin; - const double ez = (zmax - zmin) + 2.0 * margin; - grid.spacing[0] = nx > 1 ? ex / static_cast(nx - 1) : 0.0; - grid.spacing[1] = ny > 1 ? ey / static_cast(ny - 1) : 0.0; - grid.spacing[2] = nz > 1 ? ez / static_cast(nz - 1) : 0.0; - return grid; -} - -std::vector CubeGrid::points() const { - std::vector pts(static_cast(num_points()) * 3); - points_into(pts.data()); - return pts; -} - -void CubeGrid::points_into(double* out) const { - size_t k = 0; - for (int64_t ix = 0; ix < nx; ++ix) { - const double x = origin[0] + spacing[0] * static_cast(ix); - for (int64_t iy = 0; iy < ny; ++iy) { - const double y = origin[1] + spacing[1] * static_cast(iy); - for (int64_t iz = 0; iz < nz; ++iz) { - const double z = origin[2] + spacing[2] * static_cast(iz); - out[3 * k + 0] = x; - out[3 * k + 1] = y; - out[3 * k + 2] = z; - ++k; - } - } - } -} - -// Hand-rolled %13.5E formatter. Required field is fixed 13 chars -// (sign + D + '.' + 5d + 'E' + sign + 2d). snprintf in the inner loop -// dominates write time on large grids; this matches glibc snprintf -// bit-for-bit on the fast path (1-2 digit exponents) and falls back to -// snprintf for 3-digit-exponent edge cases. Output buffer must be -// exactly 13 bytes (no trailing NUL). - namespace { -inline void format_e13_5(double v, char* out) { - if (std::isnan(v)) { - std::memcpy(out, " NaN", 13); +/** @brief Format `v` as a fixed 13-character "%13.5E" field. + * + * Layout: sign + D + '.' + 5 digits + 'E' + sign + 2 digits. snprintf in + * the inner loop dominates write time on large grids; this reproduces glibc + * "%13.5E" for every value except exact half-way ties in the 6th significant + * digit, where glibc rounds the exact binary value half-to-even while this + * routine rounds half-away-from-zero. Both agree to within one unit of the + * last printed digit. 3-digit exponents defer to snprintf. + * + * @param[in] v Value to format. + * @param[out] out Exactly 13 bytes are written (no trailing NUL). + */ +inline void format_e13_5( double v, char* out ) { + + if( std::isnan(v) ) { + std::memcpy( out, std::signbit(v) ? " -NAN" : " NAN", 13 ); return; } - if (std::isinf(v)) { - std::memcpy(out, v < 0 ? " -Inf" : " Inf", 13); + if( std::isinf(v) ) { + std::memcpy( out, v < 0 ? " -INF" : " INF", 13 ); return; } - bool negative = std::signbit(v); - double absv = std::fabs(v); + const bool negative = std::signbit(v); + const double absv = std::fabs(v); - if (absv == 0.0) { + if( absv == 0.0 ) { // glibc "%13.5E": 0.0 -> " 0.00000E+00"; -0.0 -> " -0.00000E+00". - out[0] = ' '; - out[1] = negative ? '-' : ' '; - out[2] = '0'; - out[3] = '.'; - out[4] = '0'; - out[5] = '0'; - out[6] = '0'; - out[7] = '0'; - out[8] = '0'; - out[9] = 'E'; - out[10] = '+'; - out[11] = '0'; - out[12] = '0'; + std::memcpy( out, negative ? " -0.00000E+00" : " 0.00000E+00", 13 ); return; } // Exponent via floor(log10), with corrections for FP edge cases // (e.g. 9.99999 rounding up across a power-of-ten boundary). - int exp10 = static_cast(std::floor(std::log10(absv))); - double scale = std::pow(10.0, -exp10); + int exp10 = static_cast( std::floor( std::log10(absv) ) ); + double scale = std::pow( 10.0, -exp10 ); double mant = absv * scale; - long long mant_int = static_cast(std::llround(mant * 1e5)); - if (mant_int >= 1000000) { + long long mant_int = std::llround( mant * 1e5 ); + if( mant_int >= 1000000 ) { mant_int = 100000; ++exp10; - } else if (mant_int < 100000) { + } else if( mant_int < 100000 ) { --exp10; - scale = std::pow(10.0, -exp10); + scale = std::pow( 10.0, -exp10 ); mant = absv * scale; - mant_int = static_cast(std::llround(mant * 1e5)); - if (mant_int >= 1000000) { - mant_int = 999999; - } else if (mant_int < 100000) { - mant_int = 100000; - } + mant_int = std::llround( mant * 1e5 ); + if( mant_int >= 1000000 ) mant_int = 999999; + else if( mant_int < 100000 ) mant_int = 100000; } // 3+ digit exponents have a different field layout; defer to snprintf. - if (exp10 > 99 || exp10 < -99) { + if( exp10 > 99 or exp10 < -99 ) { char tmp[32]; - const int n = std::snprintf(tmp, sizeof(tmp), "%13.5E", v); - if (n >= 13) { - std::memcpy(out, tmp + (n - 13), 13); + const int n = std::snprintf( tmp, sizeof(tmp), "%13.5E", v ); + if( n >= 13 ) { + std::memcpy( out, tmp + (n - 13), 13 ); } else { const int pad = 13 - n; - for (int i = 0; i < pad; ++i) out[i] = ' '; - std::memcpy(out + pad, tmp, static_cast(n)); + for( int i = 0; i < pad; ++i ) out[i] = ' '; + std::memcpy( out + pad, tmp, static_cast(n) ); } return; } @@ -173,8 +95,8 @@ inline void format_e13_5(double v, char* out) { out[1] = negative ? '-' : ' '; char digits[6]; - for (int i = 5; i >= 0; --i) { - digits[i] = static_cast('0' + (mant_int % 10)); + for( int i = 5; i >= 0; --i ) { + digits[i] = static_cast( '0' + (mant_int % 10) ); mant_int /= 10; } out[2] = digits[0]; @@ -188,116 +110,105 @@ inline void format_e13_5(double v, char* out) { out[10] = exp10 < 0 ? '-' : '+'; const int aexp = exp10 < 0 ? -exp10 : exp10; - out[11] = static_cast('0' + (aexp / 10)); - out[12] = static_cast('0' + (aexp % 10)); + out[11] = static_cast( '0' + (aexp / 10) ); + out[12] = static_cast( '0' + (aexp % 10) ); + } } // namespace -// ============================================================================= -// write_cube -// ============================================================================= -void write_cube(const std::string& path, const Molecule& mol, - const CubeGrid& grid, const double* field, - const std::string& comment) { - if (field == nullptr) { +void write_cube( const std::string& path, const Molecule& mol, + const CubeGrid& grid, const double* field, + const std::string& comment ) { + + if( not field ) { GAUXC_GENERIC_EXCEPTION("write_cube: field pointer is null."); } - if (grid.num_points() <= 0) { + if( grid.num_points() <= 0 ) { GAUXC_GENERIC_EXCEPTION("write_cube: grid has zero points."); } - std::FILE* f = std::fopen(path.c_str(), "w"); - if (f == nullptr) { + std::FILE* f = std::fopen( path.c_str(), "w" ); + if( not f ) { GAUXC_GENERIC_EXCEPTION("write_cube: failed to open output file: " + path); } + std::unique_ptr fh( f, &std::fclose ); // --- Header --- - std::fprintf(f, "%s\n", - comment.empty() ? "GauXC cube file" : comment.c_str()); - std::fprintf(f, "Generated by GauXC\n"); + std::fprintf( f, "%s\n", + comment.empty() ? "GauXC cube file" : comment.c_str() ); + std::fprintf( f, "Generated by GauXC\n" ); // natoms + origin (Bohr). - std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", - static_cast(mol.size()), grid.origin[0], - grid.origin[1], grid.origin[2]); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(mol.size()), grid.origin[0], grid.origin[1], + grid.origin[2] ); // Three voxel-axis lines (axis-aligned grid). - std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", - static_cast(grid.nx), grid.spacing[0], 0.0, 0.0); - std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", - static_cast(grid.ny), 0.0, grid.spacing[1], 0.0); - std::fprintf(f, "%5lld %12.6f %12.6f %12.6f\n", - static_cast(grid.nz), 0.0, 0.0, grid.spacing[2]); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nx), grid.spacing[0], 0.0, 0.0 ); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.ny), 0.0, grid.spacing[1], 0.0 ); + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f\n", + static_cast(grid.nz), 0.0, 0.0, grid.spacing[2] ); // One line per atom: Z, partial charge (0.0), x, y, z (Bohr). - for (const auto& atom : mol) { - std::fprintf(f, "%5lld %12.6f %12.6f %12.6f %12.6f\n", - static_cast(atom.Z.get()), 0.0, atom.x, atom.y, - atom.z); + for( const auto& atom : mol ) { + std::fprintf( f, "%5lld %12.6f %12.6f %12.6f %12.6f\n", + static_cast(atom.Z.get()), 0.0, atom.x, atom.y, atom.z ); } // --- Data block --- - // Each (ix, iy) row is grid.nz values, six per line, %13.5E. The cube - // format requires a newline at the end of every (ix, iy) row regardless - // of how many values land on the last line. Rows are independent, so we - // format them in parallel into a pre-sized buffer and commit with a - // single fwrite. + // Each (ix, iy) row is grid.nz values, six per line, %13.5E. The cube format + // requires a newline at the end of every row regardless of how many values + // land on the last line, so a row occupies exactly nz*13 + ceil(nz/6) bytes. + // Rows are therefore equally sized and independent: a chunk of rows is + // formatted in parallel at known offsets and committed with a single fwrite, + // with no compaction pass and a staging buffer that stays bounded regardless + // of grid size. const int64_t nz = grid.nz; - const int64_t lines_per_row = (nz + 5) / 6; - // Worst case 6*13 + 1 = 79 bytes per line; over-allocates the trailing - // line of each row but avoids a precise sizing pass. - const int64_t bytes_per_row = lines_per_row * (6 * 13 + 1); + const int64_t bytes_per_row = nz * 13 + (nz + 5) / 6; const int64_t n_rows = grid.nx * grid.ny; - std::vector buf(static_cast(bytes_per_row * n_rows)); - std::vector row_byte_count(static_cast(n_rows), 0); + + // 8 MiB is well past the point where sequential fwrite stops caring about + // block size, and keeps the staging buffer small enough to stay cache- and + // NUMA-friendly on large grids. + constexpr int64_t target_chunk_bytes = 8ll * 1024 * 1024; + const int64_t rows_per_chunk = + std::clamp( target_chunk_bytes / bytes_per_row, 1, n_rows ); + std::vector buf( static_cast(rows_per_chunk * bytes_per_row) ); + + for( int64_t r0 = 0; r0 < n_rows; r0 += rows_per_chunk ) { + + const int64_t nr = std::min( rows_per_chunk, n_rows - r0 ); #pragma omp parallel for schedule(static) - for (int64_t row = 0; row < n_rows; ++row) { - const double* row_data = - field + static_cast(row) * static_cast(nz); - char* dst = buf.data() + static_cast(row) * - static_cast(bytes_per_row); - int64_t off = 0; - for (int64_t iz = 0; iz < nz; ++iz) { - format_e13_5(row_data[iz], dst + off); - off += 13; - // Every 6 values OR at the end of the row → newline. - if (((iz + 1) % 6 == 0) || (iz + 1 == nz)) { - dst[off++] = '\n'; + for( int64_t r = 0; r < nr; ++r ) { + const double* row_data = + field + static_cast(r0 + r) * static_cast(nz); + char* dst = + buf.data() + static_cast(r) * static_cast(bytes_per_row); + int64_t off = 0; + for( int64_t iz = 0; iz < nz; ++iz ) { + format_e13_5( row_data[iz], dst + off ); + off += 13; + // Every 6 values OR at the end of the row -> newline. + if( (iz + 1) % 6 == 0 or iz + 1 == nz ) dst[off++] = '\n'; } } - row_byte_count[static_cast(row)] = off; - } - // Compact rows in place (worst-case padding between them) and emit - // with a single fwrite. - if (n_rows > 1) { - int64_t write_off = row_byte_count[0]; - for (int64_t row = 1; row < n_rows; ++row) { - const int64_t src_off = row * bytes_per_row; - const int64_t len = row_byte_count[static_cast(row)]; - std::memmove(buf.data() + write_off, buf.data() + src_off, - static_cast(len)); - write_off += len; - } - if (std::fwrite(buf.data(), 1, static_cast(write_off), f) != - static_cast(write_off)) { - std::fclose(f); - GAUXC_GENERIC_EXCEPTION("write_cube: short write to " + path); - } - } else { - if (std::fwrite(buf.data(), 1, static_cast(row_byte_count[0]), - f) != static_cast(row_byte_count[0])) { - std::fclose(f); + const size_t nbytes = static_cast( nr * bytes_per_row ); + if( std::fwrite( buf.data(), 1, nbytes, f ) != nbytes ) { GAUXC_GENERIC_EXCEPTION("write_cube: short write to " + path); } + } - if (std::fclose(f) != 0) { + if( std::fclose( fh.release() ) != 0 ) { GAUXC_GENERIC_EXCEPTION("write_cube: failed to close " + path); } + } } // namespace GauXC diff --git a/src/external/cube_hdf5.cxx b/src/external/cube_hdf5.cxx index 13b8f0d1..db909d57 100644 --- a/src/external/cube_hdf5.cxx +++ b/src/external/cube_hdf5.cxx @@ -3,7 +3,7 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * @@ -21,15 +21,15 @@ namespace GauXC { -void write_cube_hdf5(const std::string& path, - const Molecule& mol, - const CubeGrid& grid, - const double* field, - const std::string& comment) { - if (grid.num_points() <= 0) { +void write_cube_hdf5( const std::string& path, + const Molecule& mol, + const CubeGrid& grid, + const double* field, + const std::string& comment ) { + if( grid.num_points() <= 0 ) { GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: grid has zero points."); } - if (field == nullptr) { + if( not field ) { GAUXC_GENERIC_EXCEPTION("write_cube_hdf5: null field pointer."); } @@ -64,7 +64,7 @@ void write_cube_hdf5(const std::string& path, std::vector atomic_numbers(natom); std::vector coords(natom * 3); - for (size_t i = 0; i < natom; ++i) { + for( size_t i = 0; i < natom; ++i ) { atomic_numbers[i] = static_cast(mol[i].Z.get()); coords[3 * i + 0] = mol[i].x; coords[3 * i + 1] = mol[i].y; diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 05e7064a..b3e68cc3 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -3,18 +3,18 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * * See LICENSE.txt for details */ -#include +#include "orbital_evaluator_impl.hpp" #include #include -#include -#include +#include +#include #include #include @@ -22,86 +22,144 @@ #include #else inline int omp_get_max_threads() { return 1; } -inline int omp_get_thread_num() { return 0; } #endif -#include -#include #include #include -#include #include "xc_integrator/integrator_util/integrator_common.hpp" #include "xc_integrator/local_work_driver/host/blas.hpp" -#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" namespace GauXC { -namespace { +namespace detail { -/// Per-thread point batch size, sized so the AO scratch buffer -/// (nbf*batch doubles) fits comfortably in L2 (~16 MB target). -inline int64_t choose_batch_size(int32_t nbf) { - constexpr int64_t kTargetBytesPerThread = 16ll * 1024 * 1024; - constexpr int64_t kMinBatch = 256; - constexpr int64_t kMaxBatch = 8192; - if (nbf <= 0) return kMaxBatch; - const int64_t b = - kTargetBytesPerThread / (static_cast(sizeof(double)) * nbf); - return std::clamp(b, kMinBatch, kMaxBatch); -} +OrbitalEvaluatorImpl::OrbitalEvaluatorImpl( BasisSet bs ) : + basis( std::move(bs) ) { + + driver_owner = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference" ); + host_driver = dynamic_cast( driver_owner.get() ); + if( not host_driver ) { + GAUXC_GENERIC_EXCEPTION("OrbitalEvaluator: LocalWorkDriverFactory did not " + "return a LocalHostWorkDriver."); + } + + nbf_ = basis.nbf(); + + // BasisSetMap powers shell -> AO range lookups for the screened submat_map. + // No Molecule is available here, so an empty one is passed; the only field + // that needs it (shell_to_center) is unused by this class. + basis_map = std::make_unique( basis, Molecule{} ); + + // Cache squared cutoff radii so the screening loop is a comparison rather + // than a square root per shell per batch. + shell_cutoff_r2.resize( basis.size() ); + for( size_t s = 0; s < basis.size(); ++s ) { + const double r = basis[s].cutoff_radius(); + shell_cutoff_r2[s] = r * r; + } -/// Single-block submat_map for the no-screening case (nbe == nbf); makes -/// `eval_xmat` route directly to GEMM without invoking the gather path. -inline LocalHostWorkDriver::submat_map_t full_submat_map(int32_t nbf) { - return {{ {int32_t{0}, nbf, int32_t{0}} }}; } -/// Fill `pts[3*np]` with grid-point coordinates for batch [p0, p0+np). -/// Avoids materialising the full 3*N point array for grid-based evaluation. -inline void fill_batch_pts(const CubeGrid& g, int64_t p0, int64_t np, - double* pts) { - const int64_t nz = g.nz, ny = g.ny; - for (int64_t i = 0; i < np; ++i) { - const int64_t k = p0 + i; - const int64_t iz = k % nz; - const int64_t iy = (k / nz) % ny; - const int64_t ix = k / (ny * nz); - pts[3 * i + 0] = g.origin[0] + static_cast(ix) * g.spacing[0]; - pts[3 * i + 1] = g.origin[1] + static_cast(iy) * g.spacing[1]; - pts[3 * i + 2] = g.origin[2] + static_cast(iz) * g.spacing[2]; +} // namespace detail + + +namespace { + +/** @brief Choose the number of points evaluated per batch. + * + * Two constraints, whichever is tighter: + * 1. Cache footprint: the AO scratch (nbf*batch doubles) should stay + * around 1 MiB per thread so it lives comfortably in L2/L3. + * 2. Load balance: enough batches that each thread gets several, i.e. + * batch <= npts / (4*nthreads). + */ +size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads ) { + constexpr size_t kTargetAOBytesPerThread = 1024 * 1024; + constexpr size_t kMinBatch = 128; + constexpr size_t kMaxBatch = 8192; + + size_t batch = kMaxBatch; + if( nbf > 0 ) { + batch = kTargetAOBytesPerThread / ( sizeof(double) * static_cast(nbf) ); } + if( nthreads > 0 ) { + batch = std::min( batch, npts / ( 4 * static_cast(nthreads) ) ); + } + return std::clamp( batch, kMinBatch, kMaxBatch ); } -/// Axis-aligned bbox of a set of npts AoS points (length 3*npts). +/// Axis-aligned bounding box of a point set. struct PointBbox { std::array lo; std::array hi; }; -inline PointBbox compute_bbox(const double* points, int64_t npts) { - PointBbox b{{points[0], points[1], points[2]}, - {points[0], points[1], points[2]}}; - for (int64_t p = 1; p < npts; ++p) { + +/// Bbox of `npts` AoS points (array of length 3*npts). +PointBbox compute_bbox( const double* points, size_t npts ) { + PointBbox b{ {points[0], points[1], points[2]}, + {points[0], points[1], points[2]} }; + for( size_t p = 1; p < npts; ++p ) { const double* xyz = points + 3 * p; - for (int k = 0; k < 3; ++k) { - if (xyz[k] < b.lo[k]) b.lo[k] = xyz[k]; - if (xyz[k] > b.hi[k]) b.hi[k] = xyz[k]; + for( int k = 0; k < 3; ++k ) { + if( xyz[k] < b.lo[k] ) b.lo[k] = xyz[k]; + if( xyz[k] > b.hi[k] ) b.hi[k] = xyz[k]; + } + } + return b; +} + +/** @brief Bbox of the contiguous CubeGrid index range [p0, p0+np). + * + * Exact (not merely conservative): with iz varying fastest, a range that + * crosses an ix boundary necessarily contains both iy=0 and iy=ny-1, and a + * range that crosses an iy boundary necessarily contains both iz=0 and + * iz=nz-1. Computing this analytically avoids rescanning the 3*np + * coordinates that were just generated. + */ +PointBbox bbox_for_range( const CubeGrid& g, size_t p0, size_t np ) { + const int64_t nz = g.nz, ny = g.ny; + const int64_t k0 = static_cast(p0); + const int64_t k1 = static_cast(p0 + np) - 1; + + int64_t lo_idx[3], hi_idx[3]; + lo_idx[0] = k0 / (ny * nz); + hi_idx[0] = k1 / (ny * nz); + if( hi_idx[0] > lo_idx[0] ) { + lo_idx[1] = 0; hi_idx[1] = ny - 1; + lo_idx[2] = 0; hi_idx[2] = nz - 1; + } else { + lo_idx[1] = (k0 / nz) % ny; + hi_idx[1] = (k1 / nz) % ny; + if( hi_idx[1] > lo_idx[1] ) { + lo_idx[2] = 0; hi_idx[2] = nz - 1; + } else { + lo_idx[2] = k0 % nz; + hi_idx[2] = k1 % nz; } } + + PointBbox b; + for( int k = 0; k < 3; ++k ) { + const double a = g.origin[k] + g.spacing[k] * static_cast(lo_idx[k]); + const double c = g.origin[k] + g.spacing[k] * static_cast(hi_idx[k]); + b.lo[k] = std::min(a, c); + b.hi[k] = std::max(a, c); + } return b; } -/// Squared distance from `center` to the nearest point in the axis-aligned -/// bbox `[lo, hi]^3`. Zero if the center lies inside the bbox. -inline double dist2_center_to_bbox(const double* center, - const PointBbox& bbox) { +/// Squared distance from `center` to the nearest point of the bbox. Zero if +/// the center lies inside. +double dist2_center_to_bbox( const double* center, const PointBbox& bbox ) { double d2 = 0.0; - for (int k = 0; k < 3; ++k) { + for( int k = 0; k < 3; ++k ) { const double c = center[k]; - if (c < bbox.lo[k]) { + if( c < bbox.lo[k] ) { const double dx = bbox.lo[k] - c; d2 += dx * dx; - } else if (c > bbox.hi[k]) { + } else if( c > bbox.hi[k] ) { const double dx = c - bbox.hi[k]; d2 += dx * dx; } @@ -109,566 +167,372 @@ inline double dist2_center_to_bbox(const double* center, return d2; } -} // namespace +/// Point source backed by a caller-supplied AoS coordinate array. +struct RawPointSource { + static constexpr bool needs_scratch = false; + const double* points; -struct OrbitalEvaluator::Impl { - BasisSet basis; - std::unique_ptr driver_owner; - LocalHostWorkDriver* host_driver = nullptr; // non-owning view - std::vector shell_list; - // BasisSetMap powers shell -> AO range lookups for the screened - // submat_map. We don't have a Molecule available so we pass an empty - // one; the only field that needs it (shell_to_center) is unused here. - std::unique_ptr basis_map; - // Per-shell squared cutoff radius, cached so the screening loop is just - // a comparison instead of a square root per shell per batch. - std::vector shell_cutoff_r2; - int32_t nbf_ = 0; - - // Per-thread scratch for eval_xmat in eval_density. Allocated once - // per evaluator rather than once per eval_density call, saving - // nthreads * nbf^2 * 8 bytes of allocation churn per call. - mutable std::vector> xmat_scratch_pool; - - /// Return a pointer to thread-local xmat scratch of size nbf*nbf. - /// Must be called from inside an omp parallel region whose thread - /// count does not exceed the value of omp_get_max_threads() at - /// construction time. - double* xmat_scratch() const { - auto& buf = xmat_scratch_pool[omp_get_thread_num()]; - return buf.data(); + const double* batch( size_t p0, size_t, double* ) const { + return points + 3 * p0; + } + PointBbox bbox( size_t, size_t np, const double* pts ) const { + return compute_bbox( pts, np ); } +}; - void init(BasisSet bs, ExecutionSpace exec) { - if (exec != ExecutionSpace::Host) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator: only ExecutionSpace::Host is currently supported."); +/// Point source that generates CubeGrid coordinates on the fly, avoiding the +/// 3*num_points()*8 byte temporary coordinate array. +struct GridPointSource { + static constexpr bool needs_scratch = true; + const CubeGrid* grid; + + const double* batch( size_t p0, size_t np, double* scr ) const { + // Incremental (ix,iy,iz) counters; coordinates are recomputed from the + // index (rather than accumulated) to stay bit-identical to + // CubeGrid::points_into. + const CubeGrid& g = *grid; + const int64_t nz = g.nz, ny = g.ny; + int64_t iz = static_cast(p0) % nz; + int64_t iy = (static_cast(p0) / nz) % ny; + int64_t ix = static_cast(p0) / (ny * nz); + double x = g.origin[0] + g.spacing[0] * static_cast(ix); + double y = g.origin[1] + g.spacing[1] * static_cast(iy); + + for( size_t i = 0; i < np; ++i ) { + scr[3 * i + 0] = x; + scr[3 * i + 1] = y; + scr[3 * i + 2] = g.origin[2] + g.spacing[2] * static_cast(iz); + if( ++iz == nz ) { + iz = 0; + if( ++iy == ny ) { + iy = 0; + ++ix; + x = g.origin[0] + g.spacing[0] * static_cast(ix); + } + y = g.origin[1] + g.spacing[1] * static_cast(iy); + } } + return scr; + } + + PointBbox bbox( size_t p0, size_t np, const double* ) const { + return bbox_for_range( *grid, p0, np ); + } +}; - basis = std::move(bs); +/// Contraction of the AO batch against MO coefficients: +/// out(np,nmo) = ao^T(np,nbe) @ C(nbe,nmo). +class OrbitalContractor { - driver_owner = LocalWorkDriverFactory::make_local_work_driver( - ExecutionSpace::Host, "Reference"); - host_driver = dynamic_cast(driver_owner.get()); - if (host_driver == nullptr) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator: LocalWorkDriverFactory did not return a " - "LocalHostWorkDriver."); - } + const detail::OrbitalEvaluatorImpl& impl_; + int32_t nmo_; + const double* C_; + size_t ldc_; + double* out_; + size_t ldo_; + +public: + + struct Scratch { std::vector C_compressed; }; - shell_list.resize(basis.size()); - std::iota(shell_list.begin(), shell_list.end(), int32_t{0}); - nbf_ = basis.nbf(); + OrbitalContractor( const detail::OrbitalEvaluatorImpl& impl, int32_t nmo, + const double* C, size_t ldc, double* out, size_t ldo ) : + impl_(impl), nmo_(nmo), C_(C), ldc_(ldc), out_(out), ldo_(ldo) {} - basis_map = std::make_unique(basis, Molecule{}); + void init( Scratch& scr, size_t ) const { + scr.C_compressed.resize( static_cast(impl_.nbf_) * nmo_ ); + } - shell_cutoff_r2.resize(basis.size()); - for (size_t s = 0; s < basis.size(); ++s) { - const double r = basis[s].cutoff_radius(); - shell_cutoff_r2[s] = r * r; + void zero( size_t p0, size_t np ) const { + for( int32_t j = 0; j < nmo_; ++j ) { + double* out_col = out_ + static_cast(j) * ldo_ + p0; + std::fill( out_col, out_col + np, 0.0 ); } + } - // Pre-size the scratch pool for the current OMP thread count. - xmat_scratch_pool.resize(omp_get_max_threads()); - for (auto& buf : xmat_scratch_pool) { - buf.resize(static_cast(nbf_) * nbf_); + void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + // Gather the rows of C for the surviving shells into a contiguous + // (nbe, nmo) col-major buffer so the contraction is a dense GEMM. + const BasisSetMap& basis_map = *impl_.basis_map; + int32_t row = 0; + for( int32_t ish : shells ) { + const auto rng = basis_map.shell_to_ao_range(ish); + for( int32_t j = 0; j < nmo_; ++j ) { + const double* C_col = C_ + static_cast(j) * ldc_; + double* dst = scr.C_compressed.data() + + static_cast(j) * nbe + row; + std::copy( C_col + rng.first, C_col + rng.second, dst ); + } + row += rng.second - rng.first; } + + blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, + ao, nbe, scr.C_compressed.data(), nbe, 0.0, out_ + p0, + static_cast(ldo_) ); } + }; -OrbitalEvaluator::OrbitalEvaluator(BasisSet basis, ExecutionSpace exec) - : pimpl_(std::make_unique()) { - pimpl_->init(std::move(basis), exec); -} +/// Contraction of the AO batch against the density matrix: +/// rho(r) = sum_{mu,nu} D[mu,nu] phi_mu(r) phi_nu(r). +class DensityContractor { -OrbitalEvaluator::~OrbitalEvaluator() noexcept = default; -OrbitalEvaluator::OrbitalEvaluator(OrbitalEvaluator&&) noexcept = default; -OrbitalEvaluator& OrbitalEvaluator::operator=(OrbitalEvaluator&&) noexcept = - default; + const detail::OrbitalEvaluatorImpl& impl_; + const double* D_; + size_t ldd_; + double* out_; -int32_t OrbitalEvaluator::nbf() const noexcept { return pimpl_->nbf_; } +public: -const BasisSet& OrbitalEvaluator::basis() const noexcept { - return pimpl_->basis; -} + struct Scratch { + std::vector dm_ao; + std::vector xmat_scr; + }; -void OrbitalEvaluator::eval_orbital(int64_t npts, const double* points, - const double* C, double* out) const { - eval_orbitals(npts, points, /*nmo=*/1, C, /*ldc=*/pimpl_->nbf_, out, - /*ldo=*/npts); -} + DensityContractor( const detail::OrbitalEvaluatorImpl& impl, const double* D, + size_t ldd, double* out ) : + impl_(impl), D_(D), ldd_(ldd), out_(out) {} -void OrbitalEvaluator::eval_orbitals(int64_t npts, const double* points, - int32_t nmo, const double* C, int64_t ldc, - double* out, int64_t ldo) const { - if (npts == 0 || nmo == 0) return; - if (points == nullptr || C == nullptr || out == nullptr) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals: null pointer argument."); + void init( Scratch& scr, size_t batch_size ) const { + scr.dm_ao.resize( static_cast(impl_.nbf_) * batch_size ); } - const int32_t nbf = pimpl_->nbf_; - if (ldc < nbf) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals: ldc must be >= nbf()."); + + void zero( size_t p0, size_t np ) const { + std::fill( out_ + p0, out_ + p0 + np, 0.0 ); } - if (ldo < npts) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals: ldo must be >= npts."); + + void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + const std::vector& shells, const double* ao ) const { + + // eval_xmat needs nbe*nbe scratch. Growing it on demand inside the + // parallel region keeps the high-water mark at the largest nbe this + // thread actually saw (not nbf) and first-touches the pages on the + // owning thread. + const size_t scr_size = static_cast(nbe) * nbe; + if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); + + const int32_t nbf = impl_.nbf_; + LocalHostWorkDriver::submat_map_t submat_map; + std::tie( submat_map, std::ignore ) = + gen_compressed_submat_map( *impl_.basis_map, shells, nbf, nbf ); + + // dm_ao = D_compressed @ ao + impl_.host_driver->eval_xmat( np, static_cast(nbf), + static_cast(nbe), submat_map, 1.0, D_, ldd_, + ao, static_cast(nbe), + scr.dm_ao.data(), static_cast(nbe), scr.xmat_scr.data() ); + + // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) + impl_.host_driver->eval_uvvar_lda_rks( np, static_cast(nbe), ao, + scr.dm_ao.data(), static_cast(nbe), out_ + p0 ); } - const int32_t nshells_total = pimpl_->basis.nshells(); - const int64_t batch_size = choose_batch_size(nbf); - const int64_t n_batches = (npts + batch_size - 1) / batch_size; - LocalHostWorkDriver* driver = pimpl_->host_driver; - const BasisSet& basis = pimpl_->basis; - const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; - const BasisSetMap& basis_map = *pimpl_->basis_map; +}; + +/** @brief The single batched evaluation loop shared by all entry points. + * + * Batches of points are screened against the per-shell cutoff radii, the + * surviving shells are collocated into a per-thread AO buffer, and the + * result is handed to `contract` for the orbital- or density-specific + * reduction. All scratch is per-thread and per-call. + */ +template +void batched_eval( const detail::OrbitalEvaluatorImpl& impl, size_t npts, + const PointSource& src, const Contractor& contract ) { + + const BasisSet& basis = impl.basis; + const auto& shell_cutoff_r2 = impl.shell_cutoff_r2; + const int32_t nbf = impl.nbf_; + const int32_t nshells_total = basis.nshells(); + + const size_t batch_size = choose_batch_size( nbf, npts, omp_get_max_threads() ); + const int64_t n_batches = + static_cast( (npts + batch_size - 1) / batch_size ); #pragma omp parallel { - std::vector ao_buf(static_cast(nbf) * batch_size); - std::vector C_compressed(static_cast(nbf) * nmo); + std::vector pt_buf; + if constexpr( PointSource::needs_scratch ) pt_buf.resize( 3 * batch_size ); + std::vector ao_buf( static_cast(nbf) * batch_size ); std::vector screened_shells; - screened_shells.reserve(nshells_total); + screened_shells.reserve( nshells_total ); + typename Contractor::Scratch scr; + contract.init( scr, batch_size ); #pragma omp for schedule(dynamic, 1) - for (int64_t b = 0; b < n_batches; ++b) { - const int64_t p0 = b * batch_size; - const int64_t np = std::min(batch_size, npts - p0); + for( int64_t b = 0; b < n_batches; ++b ) { + + const size_t p0 = static_cast(b) * batch_size; + const size_t np = std::min( batch_size, npts - p0 ); - // Per-batch shell screening: keep shells whose cutoff_radius reaches + const double* pts = src.batch( p0, np, pt_buf.data() ); + + // Per-batch shell screening: keep shells whose cutoff radius reaches // any point of this batch's bounding box. - const PointBbox bbox = compute_bbox(points + 3 * p0, np); + const PointBbox bbox = src.bbox( p0, np, pts ); screened_shells.clear(); int32_t nbe = 0; - for (int32_t s = 0; s < nshells_total; ++s) { - if (dist2_center_to_bbox(basis[s].O_data(), bbox) < - shell_cutoff_r2[s]) { - screened_shells.push_back(s); - nbe += basis[s].size(); - } - } - if (nbe == 0) { - // All shells screen out for this batch -> orbital is identically - // zero. Zero the slab and skip eval. - for (int32_t j = 0; j < nmo; ++j) { - double* out_col = out + static_cast(j) * ldo + p0; - std::fill(out_col, out_col + np, 0.0); - } - continue; + for( int32_t s = 0; s < nshells_total; ++s ) + if( dist2_center_to_bbox( basis[s].O_data(), bbox ) < shell_cutoff_r2[s] ) { + screened_shells.push_back(s); + nbe += basis[s].size(); } - driver->eval_collocation( - static_cast(np), - static_cast(screened_shells.size()), - static_cast(nbe), points + 3 * p0, basis, - screened_shells.data(), ao_buf.data()); - - // Gather the rows of C corresponding to surviving shells into a - // contiguous (nbe, nmo) col-major buffer so the contraction is a - // dense GEMM. - { - int32_t row = 0; - for (int32_t s : screened_shells) { - const auto rng = basis_map.shell_to_ao_range(s); - const int32_t shell_nbe = rng.second - rng.first; - for (int32_t j = 0; j < nmo; ++j) { - const double* C_col = C + static_cast(j) * ldc; - double* dst = C_compressed.data() + - static_cast(j) * nbe + row; - std::copy(C_col + rng.first, C_col + rng.second, dst); - } - row += shell_nbe; - } - } + // All shells screen out -> the result is identically zero here. + if( not nbe ) { contract.zero( p0, np ); continue; } - // out_slab(np, nmo) = ao^T(np, nbe) @ C_compressed(nbe, nmo); j-th MO - // column lives at out + j*ldo + p0. - blas::gemm( - /*TA=*/'T', /*TB=*/'N', - /*M=*/static_cast(np), /*N=*/static_cast(nmo), - /*K=*/nbe, /*ALPHA=*/1.0, /*A=*/ao_buf.data(), /*LDA=*/nbe, - /*B=*/C_compressed.data(), /*LDB=*/nbe, /*BETA=*/0.0, - /*C=*/out + p0, /*LDC=*/static_cast(ldo)); - } - } -} + // Both contractions assume eval_collocation packs AO rows by cumulative + // shell size in shell_list order. That holds for the gau2grid path; the + // non-gau2grid fallback in gau2grid_collocation.cxx instead uses the + // global shell_to_first_ao offset, which is only equivalent when no + // shell is screened out. This is a pre-existing inconsistency in the + // reference driver rather than one introduced here. + impl.host_driver->eval_collocation( np, screened_shells.size(), + static_cast(nbe), pts, basis, screened_shells.data(), + ao_buf.data() ); -void OrbitalEvaluator::eval_density(int64_t npts, const double* points, - const double* D, int64_t ldd, - double* out) const { - if (npts == 0) return; - if (points == nullptr || D == nullptr || out == nullptr) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_density: null pointer argument."); - } - const int32_t nbf = pimpl_->nbf_; - if (ldd < nbf) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_density: ldd must be >= nbf()."); + contract.apply( scr, p0, np, nbe, screened_shells, ao_buf.data() ); + + } } - const int32_t nshells_total = pimpl_->basis.nshells(); - const int64_t batch_size = choose_batch_size(nbf); - const int64_t n_batches = (npts + batch_size - 1) / batch_size; - LocalHostWorkDriver* driver = pimpl_->host_driver; - const BasisSet& basis = pimpl_->basis; - const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; - const BasisSetMap& basis_map = *pimpl_->basis_map; +} -#pragma omp parallel - { - std::vector ao_buf(static_cast(nbf) * batch_size); - std::vector dm_ao_buf(static_cast(nbf) * batch_size); - std::vector screened_shells; - screened_shells.reserve(nshells_total); +void check_orbital_args( const std::string& ctx, size_t npts, int32_t nbf, + const double* C, size_t ldc, const double* out, + size_t ldo ) { + if( not C or not out ) + GAUXC_GENERIC_EXCEPTION( ctx + ": null pointer argument." ); + if( ldc < static_cast(nbf) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldc must be >= nbf()." ); + if( ldo < npts ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldo must be >= npts." ); + // ldo is passed to blas::gemm, whose LDC parameter is int. + if( ldo > static_cast(std::numeric_limits::max()) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldo exceeds BLAS int range." ); +} -#pragma omp for schedule(dynamic, 1) - for (int64_t b = 0; b < n_batches; ++b) { - const int64_t p0 = b * batch_size; - const int64_t np = std::min(batch_size, npts - p0); +void check_density_args( const std::string& ctx, int32_t nbf, const double* D, + size_t ldd, const double* out ) { + if( not D or not out ) + GAUXC_GENERIC_EXCEPTION( ctx + ": null pointer argument." ); + if( ldd < static_cast(nbf) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldd must be >= nbf()." ); +} - const PointBbox bbox = compute_bbox(points + 3 * p0, np); - screened_shells.clear(); - int32_t nbe = 0; - for (int32_t s = 0; s < nshells_total; ++s) { - if (dist2_center_to_bbox(basis[s].O_data(), bbox) < - shell_cutoff_r2[s]) { - screened_shells.push_back(s); - nbe += basis[s].size(); - } - } - if (nbe == 0) { - std::fill(out + p0, out + p0 + np, 0.0); - continue; - } +} // namespace - driver->eval_collocation( - static_cast(np), - static_cast(screened_shells.size()), - static_cast(nbe), points + 3 * p0, basis, - screened_shells.data(), ao_buf.data()); - - // Build the compressed submat_map for this batch (gather the rows - // of D corresponding to surviving shells). Falls through to a single - // direct-gemm submat_map when nbe == nbf. - LocalHostWorkDriver::submat_map_t submat_map; - if (nbe == nbf) { - submat_map = full_submat_map(nbf); - } else { - std::tie(submat_map, std::ignore) = - gen_compressed_submat_map(basis_map, screened_shells, nbf, nbf); - } - // dm_ao = D_compressed @ ao - driver->eval_xmat( - /*npts=*/static_cast(np), /*nbf=*/static_cast(nbf), - /*nbe=*/static_cast(nbe), submat_map, /*fac=*/1.0, - /*P=*/D, /*ldp=*/static_cast(ldd), - /*basis_eval=*/ao_buf.data(), /*ldb=*/static_cast(nbe), - /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbe), - /*scr=*/pimpl_->xmat_scratch()); - - // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) - driver->eval_uvvar_lda_rks( - /*npts=*/static_cast(np), /*nbe=*/static_cast(nbe), - /*basis_eval=*/ao_buf.data(), - /*X=*/dm_ao_buf.data(), /*ldx=*/static_cast(nbe), - /*den_eval=*/out + p0); - } - } +OrbitalEvaluator::OrbitalEvaluator( pimpl_ptr_type&& pimpl ) : + pimpl_( std::move(pimpl) ) { + if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED(); } -// --------------------------------------------------------------------------- -// CubeGrid overloads: generate per-batch point coordinates on-the-fly, -// avoiding the 3*num_points()*8 byte temporary coordinate array. -// --------------------------------------------------------------------------- +OrbitalEvaluator::~OrbitalEvaluator() noexcept = default; +OrbitalEvaluator::OrbitalEvaluator( OrbitalEvaluator&& ) noexcept = default; +OrbitalEvaluator& OrbitalEvaluator::operator=( OrbitalEvaluator&& ) noexcept = + default; -void OrbitalEvaluator::eval_orbital(const CubeGrid& grid, - const double* C, double* out) const { - eval_orbitals(grid, /*nmo=*/1, C, /*ldc=*/pimpl_->nbf_, out, - /*ldo=*/grid.num_points()); +int32_t OrbitalEvaluator::nbf() const { return pimpl_->nbf_; } + +const BasisSet& OrbitalEvaluator::basis() const { + return pimpl_->basis; } -void OrbitalEvaluator::eval_orbitals(const CubeGrid& grid, - int32_t nmo, const double* C, int64_t ldc, - double* out, int64_t ldo) const { - const int64_t npts = grid.num_points(); - if (npts == 0 || nmo == 0) return; - if (C == nullptr || out == nullptr) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals(grid): null pointer argument."); - } - const int32_t nbf = pimpl_->nbf_; - if (ldc < nbf) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals(grid): ldc must be >= nbf()."); - } - if (ldo < npts) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_orbitals(grid): ldo must be >= npts."); - } - const int32_t nshells_total = pimpl_->basis.nshells(); - const int64_t batch_size = choose_batch_size(nbf); - const int64_t n_batches = (npts + batch_size - 1) / batch_size; - LocalHostWorkDriver* driver = pimpl_->host_driver; - const BasisSet& basis = pimpl_->basis; - const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; - const BasisSetMap& basis_map = *pimpl_->basis_map; - - // Choose between task pipeline and omp-for based on saturation. - // When n_batches >> nthreads, omp-for already fully saturates all cores - // and the task pipeline's scheduling overhead is a net negative. - const int nthreads = omp_get_max_threads(); - const bool use_pipeline = (n_batches < 2 * nthreads); - - if (use_pipeline) { - // Pipelined evaluation: overlap collocation of batch B+1 with - // GEMM of batch B using OMP tasks with dependency tags. - struct Slot { - std::vector pts; - std::vector ao; - std::vector C_comp; - std::vector shells; - int32_t nbe; - int64_t p0, np; - }; +OrbitalEvaluator OrbitalEvaluatorFactory::make_orbital_evaluator( + ExecutionSpace ex, BasisSet basis ) { - const int nslots = - std::min(nthreads, n_batches); - std::vector slots(nslots); - for (auto& s : slots) { - s.pts.resize(static_cast(batch_size) * 3); - s.ao.resize(static_cast(nbf) * batch_size); - s.C_comp.resize(static_cast(nbf) * nmo); - s.shells.reserve(nshells_total); - s.nbe = 0; - s.p0 = 0; - s.np = 0; - } + switch(ex) { - // Dependency tag array — raw array for omp depend(). - std::unique_ptr dtag(new char[nslots]{}); - char* dtag_ptr = dtag.get(); - (void)dtag_ptr; // used only in omp depend() clauses + case ExecutionSpace::Host: + return OrbitalEvaluator( + std::make_unique( std::move(basis) ) + ); -#pragma omp parallel -#pragma omp single - { - for (int64_t b = 0; b < n_batches; ++b) { - const int si = static_cast(b % nslots); - - // Phase A: generate points, screen shells, eval collocation. -#pragma omp task shared(slots, grid, basis, shell_cutoff_r2) \ - depend(out: dtag_ptr[si]) firstprivate(b, si) - { - Slot& sl = slots[si]; - sl.p0 = b * batch_size; - sl.np = std::min(batch_size, npts - sl.p0); - - fill_batch_pts(grid, sl.p0, sl.np, sl.pts.data()); - const PointBbox bbox = compute_bbox(sl.pts.data(), sl.np); - sl.shells.clear(); - sl.nbe = 0; - for (int32_t s = 0; s < nshells_total; ++s) { - if (dist2_center_to_bbox(basis[s].O_data(), bbox) < - shell_cutoff_r2[s]) { - sl.shells.push_back(s); - sl.nbe += basis[s].size(); - } - } - if (sl.nbe > 0) { - driver->eval_collocation( - static_cast(sl.np), - static_cast(sl.shells.size()), - static_cast(sl.nbe), sl.pts.data(), basis, - sl.shells.data(), sl.ao.data()); - } - } + default: + GAUXC_GENERIC_EXCEPTION("OrbitalEvaluator: only ExecutionSpace::Host is " + "currently supported."); - // Phase B: gather C + GEMM → out. -#pragma omp task shared(slots, out, C, basis_map) \ - depend(in: dtag_ptr[si]) firstprivate(b, si) - { - Slot& sl = slots[si]; - if (sl.nbe == 0) { - for (int32_t j = 0; j < nmo; ++j) { - double* oc = out + static_cast(j) * ldo + sl.p0; - std::fill(oc, oc + sl.np, 0.0); - } - } else { - const int32_t nbe = sl.nbe; - { - int32_t row = 0; - for (int32_t s : sl.shells) { - const auto rng = basis_map.shell_to_ao_range(s); - const int32_t shell_nbe = rng.second - rng.first; - for (int32_t j = 0; j < nmo; ++j) { - const double* C_col = C + static_cast(j) * ldc; - double* dst = sl.C_comp.data() + - static_cast(j) * nbe + row; - std::copy(C_col + rng.first, C_col + rng.second, dst); - } - row += shell_nbe; - } - } - - blas::gemm( - 'T', 'N', static_cast(sl.np), static_cast(nmo), - nbe, 1.0, sl.ao.data(), nbe, sl.C_comp.data(), nbe, - 0.0, out + sl.p0, static_cast(ldo)); - } - } - } } - } else { - // Saturated path: plain omp-for, lower overhead at high thread counts. +} -#pragma omp parallel - { - std::vector batch_pts(static_cast(batch_size) * 3); - std::vector ao_buf(static_cast(nbf) * batch_size); - std::vector C_compressed(static_cast(nbf) * nmo); - std::vector screened_shells; - screened_shells.reserve(nshells_total); -#pragma omp for schedule(dynamic, 1) - for (int64_t b = 0; b < n_batches; ++b) { - const int64_t p0 = b * batch_size; - const int64_t np = std::min(batch_size, npts - p0); +// --------------------------------------------------------------------------- +// Caller-supplied point sets +// --------------------------------------------------------------------------- - fill_batch_pts(grid, p0, np, batch_pts.data()); - const PointBbox bbox = compute_bbox(batch_pts.data(), np); - screened_shells.clear(); - int32_t nbe = 0; - for (int32_t s = 0; s < nshells_total; ++s) { - if (dist2_center_to_bbox(basis[s].O_data(), bbox) < - shell_cutoff_r2[s]) { - screened_shells.push_back(s); - nbe += basis[s].size(); - } - } - if (nbe == 0) { - for (int32_t j = 0; j < nmo; ++j) { - double* out_col = out + static_cast(j) * ldo + p0; - std::fill(out_col, out_col + np, 0.0); - } - continue; - } +void OrbitalEvaluator::eval_orbital( size_t npts, const double* points, + const double* C, double* out ) const { + eval_orbitals( npts, points, /*nmo=*/1, C, + /*ldc=*/static_cast(pimpl_->nbf_), out, /*ldo=*/npts ); +} - driver->eval_collocation( - static_cast(np), - static_cast(screened_shells.size()), - static_cast(nbe), batch_pts.data(), basis, - screened_shells.data(), ao_buf.data()); - - { - int32_t row = 0; - for (int32_t s : screened_shells) { - const auto rng = basis_map.shell_to_ao_range(s); - const int32_t shell_nbe = rng.second - rng.first; - for (int32_t j = 0; j < nmo; ++j) { - const double* C_col = C + static_cast(j) * ldc; - double* dst = C_compressed.data() + - static_cast(j) * nbe + row; - std::copy(C_col + rng.first, C_col + rng.second, dst); - } - row += shell_nbe; - } - } +void OrbitalEvaluator::eval_orbitals( size_t npts, const double* points, + int32_t nmo, const double* C, size_t ldc, + double* out, size_t ldo ) const { + if( not npts or nmo < 1 ) return; + if( not points ) + GAUXC_GENERIC_EXCEPTION( + "OrbitalEvaluator::eval_orbitals: null pointer argument."); + check_orbital_args( "OrbitalEvaluator::eval_orbitals", npts, pimpl_->nbf_, C, + ldc, out, ldo ); - blas::gemm( - 'T', 'N', static_cast(np), static_cast(nmo), nbe, - 1.0, ao_buf.data(), nbe, C_compressed.data(), nbe, - 0.0, out + p0, static_cast(ldo)); - } - } - } // end saturated fallback + batched_eval( *pimpl_, npts, RawPointSource{points}, + OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); } -void OrbitalEvaluator::eval_density(const CubeGrid& grid, - const double* D, int64_t ldd, - double* out) const { - const int64_t npts = grid.num_points(); - if (npts == 0) return; - if (D == nullptr || out == nullptr) { - GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_density(grid): null pointer argument."); - } - const int32_t nbf = pimpl_->nbf_; - if (ldd < nbf) { +void OrbitalEvaluator::eval_density( size_t npts, const double* points, + const double* D, size_t ldd, + double* out ) const { + if( not npts ) return; + if( not points ) GAUXC_GENERIC_EXCEPTION( - "OrbitalEvaluator::eval_density(grid): ldd must be >= nbf()."); - } + "OrbitalEvaluator::eval_density: null pointer argument."); + check_density_args( "OrbitalEvaluator::eval_density", pimpl_->nbf_, D, ldd, + out ); + + batched_eval( *pimpl_, npts, RawPointSource{points}, + DensityContractor( *pimpl_, D, ldd, out ) ); +} - const int32_t nshells_total = pimpl_->basis.nshells(); - const int64_t batch_size = choose_batch_size(nbf); - const int64_t n_batches = (npts + batch_size - 1) / batch_size; - LocalHostWorkDriver* driver = pimpl_->host_driver; - const BasisSet& basis = pimpl_->basis; - const auto& shell_cutoff_r2 = pimpl_->shell_cutoff_r2; - const BasisSetMap& basis_map = *pimpl_->basis_map; -#pragma omp parallel - { - std::vector batch_pts(static_cast(batch_size) * 3); - std::vector ao_buf(static_cast(nbf) * batch_size); - std::vector dm_ao_buf(static_cast(nbf) * batch_size); - std::vector screened_shells; - screened_shells.reserve(nshells_total); +// --------------------------------------------------------------------------- +// CubeGrid overloads: generate per-batch point coordinates on-the-fly, +// avoiding the 3*num_points()*8 byte temporary coordinate array. +// --------------------------------------------------------------------------- -#pragma omp for schedule(dynamic, 1) - for (int64_t b = 0; b < n_batches; ++b) { - const int64_t p0 = b * batch_size; - const int64_t np = std::min(batch_size, npts - p0); +void OrbitalEvaluator::eval_orbital( const CubeGrid& grid, const double* C, + double* out ) const { + eval_orbitals( grid, /*nmo=*/1, C, + /*ldc=*/static_cast(pimpl_->nbf_), out, + /*ldo=*/static_cast(grid.num_points()) ); +} - fill_batch_pts(grid, p0, np, batch_pts.data()); - const PointBbox bbox = compute_bbox(batch_pts.data(), np); - screened_shells.clear(); - int32_t nbe = 0; - for (int32_t s = 0; s < nshells_total; ++s) { - if (dist2_center_to_bbox(basis[s].O_data(), bbox) < - shell_cutoff_r2[s]) { - screened_shells.push_back(s); - nbe += basis[s].size(); - } - } - if (nbe == 0) { - std::fill(out + p0, out + p0 + np, 0.0); - continue; - } +void OrbitalEvaluator::eval_orbitals( const CubeGrid& grid, int32_t nmo, + const double* C, size_t ldc, double* out, + size_t ldo ) const { + if( grid.num_points() <= 0 or nmo < 1 ) return; + const size_t npts = static_cast( grid.num_points() ); + check_orbital_args( "OrbitalEvaluator::eval_orbitals(grid)", npts, + pimpl_->nbf_, C, ldc, out, ldo ); - driver->eval_collocation( - static_cast(np), - static_cast(screened_shells.size()), - static_cast(nbe), batch_pts.data(), basis, - screened_shells.data(), ao_buf.data()); - - LocalHostWorkDriver::submat_map_t submat_map; - if (nbe == nbf) { - submat_map = full_submat_map(nbf); - } else { - std::tie(submat_map, std::ignore) = - gen_compressed_submat_map(basis_map, screened_shells, nbf, nbf); - } + batched_eval( *pimpl_, npts, GridPointSource{&grid}, + OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); +} - driver->eval_xmat( - static_cast(np), static_cast(nbf), - static_cast(nbe), submat_map, 1.0, - D, static_cast(ldd), - ao_buf.data(), static_cast(nbe), - dm_ao_buf.data(), static_cast(nbe), - pimpl_->xmat_scratch()); - - driver->eval_uvvar_lda_rks( - static_cast(np), static_cast(nbe), - ao_buf.data(), - dm_ao_buf.data(), static_cast(nbe), - out + p0); - } - } +void OrbitalEvaluator::eval_density( const CubeGrid& grid, const double* D, + size_t ldd, double* out ) const { + if( grid.num_points() <= 0 ) return; + const size_t npts = static_cast( grid.num_points() ); + check_density_args( "OrbitalEvaluator::eval_density(grid)", pimpl_->nbf_, D, + ldd, out ); + + batched_eval( *pimpl_, npts, GridPointSource{&grid}, + DensityContractor( *pimpl_, D, ldd, out ) ); } } // namespace GauXC diff --git a/src/orbital_evaluator_impl.hpp b/src/orbital_evaluator_impl.hpp new file mode 100644 index 00000000..fab041ea --- /dev/null +++ b/src/orbital_evaluator_impl.hpp @@ -0,0 +1,42 @@ +/** + * GauXC Copyright (c) 2020-2024, The Regents of the University of California, + * through Lawrence Berkeley National Laboratory (subject to receipt of + * any required approvals from the U.S. Dept. of Energy). + * + * (c) 2024-2025, Microsoft Corporation + * + * All rights reserved. + * + * See LICENSE.txt for details + */ +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include "xc_integrator/local_work_driver/host/local_host_work_driver.hpp" + +namespace GauXC::detail { + +/// Host implementation state for OrbitalEvaluator +class OrbitalEvaluatorImpl { + +public: + + BasisSet basis; ///< Basis set to evaluate + std::unique_ptr driver_owner; + LocalHostWorkDriver* host_driver = nullptr; ///< Non-owning view of the driver + std::unique_ptr basis_map; ///< shell -> AO range lookups + std::vector shell_cutoff_r2; ///< Per-shell squared cutoff radius + int32_t nbf_ = 0; + + OrbitalEvaluatorImpl( BasisSet bs ); + +}; // class OrbitalEvaluatorImpl + +} // namespace GauXC::detail diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3d4c0dfd..3198d63c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -79,6 +79,9 @@ endif() set( GAUXC_REF_DATA_PATH "${PROJECT_SOURCE_DIR}/tests/ref_data" ) +# Scratch directory for tests which write files (cube I/O) +set( GAUXC_TEST_TMP_PATH "${PROJECT_BINARY_DIR}/tests/tmp" ) +file( MAKE_DIRECTORY "${GAUXC_TEST_TMP_PATH}" ) configure_file( ut_common.hpp.in ${PROJECT_BINARY_DIR}/tests/ut_common.hpp ) target_include_directories( gauxc_test PRIVATE ${PROJECT_BINARY_DIR}/tests ) target_include_directories( gauxc_test PRIVATE ${PROJECT_SOURCE_DIR}/tests ) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 30ec1f54..65c59a75 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -3,7 +3,7 @@ * through Lawrence Berkeley National Laboratory (subject to receipt of * any required approvals from the U.S. Dept. of Energy). * - * (c) 2024-2026, Microsoft Corporation + * (c) 2024-2025, Microsoft Corporation * * All rights reserved. * @@ -12,20 +12,22 @@ #include "ut_common.hpp" #include "catch2/catch.hpp" +#include #include -#include #include +#include #include #include #include #include -#include // getpid - #include #include #include #include +#ifdef _OPENMP +#include +#endif #ifdef GAUXC_HAS_HDF5 #include #endif @@ -42,30 +44,36 @@ using namespace GauXC; namespace { -/// Generate a unique temp path for test files (avoids tmpnam warning). -/// Includes the PID so concurrent MPI ranks (which share /tmp on a single -/// node) do not collide on the same file. +/// Unique path for a test output file, inside the build-tree scratch +/// directory configured by CMake. std::string make_temp_path(const char* suffix) { static int counter = 0; - return std::string("/tmp/gauxc_test_") + std::to_string(::getpid()) + - "_" + std::to_string(++counter) + suffix; + return std::string(GAUXC_TEST_TMP_PATH) + "/gauxc_test_" + + std::to_string(++counter) + suffix; } -/// Returns true on the root MPI rank (or always when MPI is disabled). -/// I/O test cases below run only on rank 0 to avoid concurrent ranks racing -/// on the same /tmp file paths during MPI test runs. -bool is_root_rank() { -#ifdef GAUXC_HAS_MPI - int rank = 0; - int initialized = 0; - MPI_Initialized(&initialized); - if (initialized) { - MPI_Comm_rank(MPI_COMM_WORLD, &rank); - } - return rank == 0; -#else - return true; -#endif +OrbitalEvaluator make_evaluator(const BasisSet& basis) { + return OrbitalEvaluatorFactory::make_orbital_evaluator(ExecutionSpace::Host, + basis); +} + +/// AO collocation over every shell, i.e. the unscreened reference. +std::vector reference_collocation(const BasisSet& basis, + int64_t npts, const double* pts) { + const int32_t nbf = basis.nbf(); + std::vector ao(static_cast(nbf) * npts); + auto drv = LocalWorkDriverFactory::make_local_work_driver( + ExecutionSpace::Host, "Reference"); + auto* host_drv = dynamic_cast(drv.get()); + REQUIRE(host_drv != nullptr); + std::vector shell_list(basis.size()); + for (size_t i = 0; i < shell_list.size(); ++i) + shell_list[i] = static_cast(i); + host_drv->eval_collocation(static_cast(npts), + static_cast(basis.nshells()), + static_cast(nbf), pts, basis, + shell_list.data(), ao.data()); + return ao; } /// Build a deterministic PRNG-based set of points within a small bounding box @@ -83,6 +91,14 @@ std::vector make_random_points(int64_t npts, unsigned seed = 1234u) { return pts; } +std::vector make_random_vector(size_t n, unsigned seed) { + std::mt19937 gen(seed); + std::uniform_real_distribution u(-1.0, 1.0); + std::vector v(n); + for (auto& x : v) x = u(gen); + return v; +} + } // namespace TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", @@ -98,22 +114,9 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", const auto pts = make_random_points(npts); // Reference: AO collocation directly via the LocalHostWorkDriver. - std::vector ao_ref(static_cast(nbf) * npts); - { - auto drv = LocalWorkDriverFactory::make_local_work_driver( - ExecutionSpace::Host, "Reference"); - auto* host_drv = dynamic_cast(drv.get()); - REQUIRE(host_drv != nullptr); - std::vector shell_list(basis.size()); - for (size_t i = 0; i < shell_list.size(); ++i) - shell_list[i] = static_cast(i); - host_drv->eval_collocation(static_cast(npts), - static_cast(basis.nshells()), - static_cast(nbf), pts.data(), basis, - shell_list.data(), ao_ref.data()); - } + const auto ao_ref = reference_collocation(basis, npts, pts.data()); - OrbitalEvaluator eval(basis); + auto eval = make_evaluator(basis); REQUIRE(eval.nbf() == nbf); SECTION("eval_orbital with one-hot coefficient reproduces single AO column") { @@ -134,10 +137,7 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", SECTION("eval_orbitals with random C matches AO ^T @ C") { const int32_t nmo = 4; - std::vector C(static_cast(nbf) * nmo); - std::mt19937 gen(99); - std::uniform_real_distribution u(-1.0, 1.0); - for (auto& v : C) v = u(gen); + const auto C = make_random_vector(static_cast(nbf) * nmo, 99u); std::vector out(static_cast(npts) * nmo, 0.0); eval.eval_orbitals(npts, pts.data(), nmo, C.data(), nbf, out.data(), npts); @@ -175,10 +175,7 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", } SECTION("eval_density with rank-1 D = c c^T equals (c.AO)^2") { - std::vector c(nbf); - std::mt19937 gen(7); - std::uniform_real_distribution u(-1.0, 1.0); - for (auto& v : c) v = u(gen); + const auto c = make_random_vector(static_cast(nbf), 7u); std::vector D(static_cast(nbf) * nbf, 0.0); for (int32_t mu = 0; mu < nbf; ++mu) { @@ -197,6 +194,21 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", CHECK(out[static_cast(p)] == Approx(ref).margin(1e-10)); } } + + SECTION("invalid leading dimensions throw") { + std::vector C(nbf, 0.0); + std::vector out(static_cast(npts), 0.0); + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf - 1, + out.data(), npts)); + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf, + out.data(), npts - 1)); + // ldo is narrowed to int for the BLAS call; oversized values must throw + // rather than silently truncate. + const size_t huge_ldo = + static_cast(std::numeric_limits::max()) + 1; + CHECK_THROWS(eval.eval_orbitals(npts, pts.data(), 1, C.data(), nbf, + out.data(), huge_ldo)); + } } TEST_CASE("CubeGrid eval overloads match pointer-based eval", @@ -205,28 +217,30 @@ TEST_CASE("CubeGrid eval overloads match pointer-based eval", auto basis = make_ccpvdz(mol, SphericalType(true)); for (auto& sh : basis) sh.set_shell_tolerance(1e-12); const int32_t nbf = basis.nbf(); - OrbitalEvaluator eval(basis); + auto eval = make_evaluator(basis); - // Small grid for fast testing. - auto grid = CubeGrid::from_molecule(mol, 6, 7, 8); - const int64_t npts = grid.num_points(); - auto pts = grid.points(); + // 6x7x8 fits in a single batch; 20x20x20 = 8000 points spreads over many + // batches at any thread count, exercising the trailing partial batch and + // the dynamic schedule. + const std::vector grids = {CubeGrid::from_molecule(mol, 6, 7, 8), + CubeGrid::from_molecule(mol, 20, 20, 20)}; - // Random MO coefficients. - std::mt19937 rng(42u); - std::uniform_real_distribution dist(-1.0, 1.0); - std::vector C(static_cast(nbf)); - for (auto& v : C) v = dist(rng); + const auto C = make_random_vector(static_cast(nbf), 42u); SECTION("eval_orbital(grid) matches eval_orbital(npts, points)") { - std::vector ref(static_cast(npts)); - eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); - std::vector out(static_cast(npts)); - eval.eval_orbital(grid, C.data(), out.data()); + std::vector ref(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); - for (int64_t p = 0; p < npts; ++p) { - CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + std::vector out(static_cast(npts)); + eval.eval_orbital(grid, C.data(), out.data()); + + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } } } @@ -235,15 +249,198 @@ TEST_CASE("CubeGrid eval overloads match pointer-based eval", std::vector D(static_cast(nbf) * nbf, 0.0); for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; - std::vector ref(static_cast(npts)); - eval.eval_density(npts, pts.data(), D.data(), nbf, ref.data()); + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + + std::vector ref(static_cast(npts)); + eval.eval_density(npts, pts.data(), D.data(), nbf, ref.data()); - std::vector out(static_cast(npts)); - eval.eval_density(grid, D.data(), nbf, out.data()); + std::vector out(static_cast(npts)); + eval.eval_density(grid, D.data(), nbf, out.data()); + for (int64_t p = 0; p < npts; ++p) { + CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + } + } + } +} + +TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { + // Two well-separated centres of different elements: batches near one centre + // screen out every shell of the other, so 0 < nbe < nbf is reached, and + // batches in the gap screen out everything (nbe == 0). + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-10); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis); + + const auto C = make_random_vector(static_cast(nbf), 2024u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + SECTION("a grid far from every atom evaluates to exactly zero") { + CubeGrid far_grid; + far_grid.origin = {100.0, 100.0, 100.0}; + far_grid.spacing = {0.5, 0.5, 0.5}; + far_grid.nx = far_grid.ny = far_grid.nz = 12; + const int64_t npts = far_grid.num_points(); + + std::vector orb(static_cast(npts), 1.0); + eval.eval_orbital(far_grid, C.data(), orb.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(orb[p] == 0.0); + + std::vector rho(static_cast(npts), 1.0); + eval.eval_density(far_grid, D.data(), nbf, rho.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(rho[p] == 0.0); + } + + SECTION("partial screening agrees with the unscreened reference") { + // Grid elongated along x (the slow axis), so a batch of consecutive + // points is a narrow x-slab: near one centre, near the other, or in the + // empty gap between them. + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.8, 2.0, 2.0}; + grid.nx = 40; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + std::vector orb(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + + std::vector rho(static_cast(npts)); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + // The grid overload screens against an analytically derived batch bbox + // while the pointer overload scans the coordinates; with screening active + // the two must still agree exactly. + std::vector orb_pts(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + std::vector rho_pts(static_cast(npts)); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); + + bool any_nonzero = false; for (int64_t p = 0; p < npts; ++p) { - CHECK(out[p] == Approx(ref[p]).margin(1e-12)); + CHECK(orb[p] == orb_pts[p]); + CHECK(rho[p] == rho_pts[p]); + + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + if (std::fabs(orb_ref) > 1e-3) any_nonzero = true; + // Screening discards shells whose magnitude is below the shell + // tolerance across the batch bbox, so agreement is at that level. + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); } + CHECK(any_nonzero); + } +} + +TEST_CASE("OrbitalEvaluator is invariant to the OpenMP thread count", + "[orbital_evaluator]") { + // Bit-exact here because the grid encloses the molecule at a 1e-12 shell + // tolerance, so nothing screens and the batch decomposition cannot change + // the arithmetic. The screened case is covered separately below. + auto mol = make_water(); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + const int32_t nbf = basis.nbf(); + + auto grid = CubeGrid::from_molecule(mol, 16, 16, 16); + const int64_t npts = grid.num_points(); + + const auto C = make_random_vector(static_cast(nbf), 5150u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + // Construct while the thread count is 1, then raise it. Per-call scratch + // must follow the thread count in force at evaluation time. + omp_set_num_threads(1); +#endif + auto eval = make_evaluator(basis); + + std::vector orb_serial(static_cast(npts)); + std::vector rho_serial(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb_serial.data()); + eval.eval_density(grid, D.data(), nbf, rho_serial.data()); + +#ifdef _OPENMP + omp_set_num_threads(saved_threads > 1 ? saved_threads : 2); +#endif + + std::vector orb_par(static_cast(npts)); + std::vector rho_par(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb_par.data()); + eval.eval_density(grid, D.data(), nbf, rho_par.data()); + +#ifdef _OPENMP + omp_set_num_threads(saved_threads); +#endif + + for (int64_t p = 0; p < npts; ++p) { + CHECK(orb_par[p] == orb_serial[p]); + CHECK(rho_par[p] == rho_serial[p]); + } +} + +TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", + "[orbital_evaluator]") { + // Batch size is derived from the thread count, so when screening is active + // different thread counts screen against different batch bounding boxes and + // the results are not bit-identical. The discrepancy is bounded by the shell + // tolerance, which is what this pins down. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + auto basis = make_ccpvdz(mol, SphericalType(true)); + constexpr double shell_tol = 1e-10; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.8, 2.0, 2.0}; + grid.nx = 40; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector rho_serial(static_cast(npts)); + std::vector rho_par(static_cast(npts)); + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + omp_set_num_threads(1); +#endif + eval.eval_density(grid, D.data(), nbf, rho_serial.data()); +#ifdef _OPENMP + omp_set_num_threads(saved_threads > 1 ? saved_threads : 2); +#endif + eval.eval_density(grid, D.data(), nbf, rho_par.data()); +#ifdef _OPENMP + omp_set_num_threads(saved_threads); +#endif + + for (int64_t p = 0; p < npts; ++p) { + CHECK(rho_par[p] == Approx(rho_serial[p]).margin(shell_tol)); } } @@ -281,7 +478,11 @@ TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { } TEST_CASE("write_cube round-trips header and field data", "[cube]") { - if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif auto mol = make_water(); CubeGrid grid = CubeGrid::from_molecule(mol, /*nx=*/5, /*ny=*/4, /*nz=*/7, /*margin=*/2.0); @@ -292,7 +493,6 @@ TEST_CASE("write_cube round-trips header and field data", "[cube]") { std::pow(10.0, (i % 5) - 2); } - // Use a tmp path. const std::string path = make_temp_path(".cube"); write_cube(path, mol, grid, field.data(), "Test cube"); @@ -360,45 +560,83 @@ TEST_CASE("write_cube round-trips header and field data", "[cube]") { for (size_t i = 0; i < field.size(); ++i) { CHECK(field_read[i] == Approx(field[i]).epsilon(1e-4).margin(1e-30)); } + in.close(); + + // Line structure: each (ix,iy) row spans ceil(nz/6) lines, full lines carry + // six 13-char fields and the row's last line carries the remainder. This + // pins the exact per-row byte count that write_cube relies on to place rows + // without a compaction pass. + { + std::ifstream lin(path); + REQUIRE(lin.is_open()); + std::string line; + for (size_t i = 0; i < 6 + mol.size(); ++i) std::getline(lin, line); + + const int64_t lines_per_row = (grid.nz + 5) / 6; + for (int64_t row = 0; row < grid.nx * grid.ny; ++row) { + for (int64_t l = 0; l < lines_per_row; ++l) { + REQUIRE(static_cast(std::getline(lin, line))); + const int64_t nvals = (l + 1 == lines_per_row) ? (grid.nz - l * 6) : 6; + CHECK(line.size() == static_cast(13 * nvals)); + } + } + CHECK_FALSE(static_cast(std::getline(lin, line))); + } std::remove(path.c_str()); } -TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { - if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. - // Spot-check the custom formatter against snprintf for a battery of values - // by writing a tiny cube file and parsing it back. This is a stronger check - // than the round-trip above since we compare the byte stream. - auto mol = make_water(); +namespace { + +/// Write `field` as a single-row cube file and return the raw data block. +std::string cube_data_block(const Molecule& mol, + const std::vector& field) { CubeGrid grid; grid.origin = {0.0, 0.0, 0.0}; grid.spacing = {0.1, 0.1, 0.1}; grid.nx = 1; grid.ny = 1; - grid.nz = 12; - - std::vector field = {0.0, -0.0, 1.23456e-10, -9.99995e-1, - 1.0e+99, -1.0e+99, 3.14159265358979, - -2.71828, 1.0, -1.0, 1e-300, - 1.234e+05}; + grid.nz = static_cast(field.size()); const std::string path = make_temp_path(".cube"); write_cube(path, mol, grid, field.data(), "fmt"); std::ifstream in(path); REQUIRE(in.is_open()); - // Skip header (2 comments + 4 grid lines + natoms atom lines). + // Skip header (2 comments + 1 natoms line + 3 axis lines + natoms atoms). std::string skip; - for (int i = 0; i < 2; ++i) std::getline(in, skip); - std::getline(in, skip); // natoms line - for (int i = 0; i < 3; ++i) std::getline(in, skip); // 3 axis lines + for (int i = 0; i < 6; ++i) std::getline(in, skip); for (size_t i = 0; i < mol.size(); ++i) std::getline(in, skip); - // Read the field-block as raw text and compare against snprintf - // line-by-line. Six values per line, then row terminator after 12 (single - // (ix, iy) row here so no early newline). - std::string data_block((std::istreambuf_iterator(in)), - std::istreambuf_iterator()); + std::string block((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + in.close(); + std::remove(path.c_str()); + return block; +} + +} // namespace + +TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // Spot-check the custom formatter against snprintf for a battery of values + // by writing a tiny cube file and parsing it back. This is a stronger check + // than the round-trip above since we compare the byte stream. + auto mol = make_water(); + + const double qnan = std::numeric_limits::quiet_NaN(); + const double inf = std::numeric_limits::infinity(); + std::vector field = {0.0, -0.0, 1.23456e-10, -9.99995e-1, + 1.0e+99, -1.0e+99, 3.14159265358979, + -2.71828, 1.0, -1.0, 1e-300, + 1.234e+05, qnan, -qnan, inf, + -inf}; + + const std::string data_block = cube_data_block(mol, field); std::ostringstream expected; for (size_t i = 0; i < field.size(); ++i) { @@ -409,13 +647,99 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { } CHECK(data_block == expected.str()); +} + +TEST_CASE("write_cube spans multiple output chunks", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // Sized to exceed write_cube's internal staging-buffer target so the chunk + // loop runs more than once, with a partial final chunk. + auto mol = make_water(); + CubeGrid grid; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {0.1, 0.1, 0.1}; + grid.nx = 2; + grid.ny = 100; + grid.nz = 4096; + const int64_t npts = grid.num_points(); + + std::vector field(static_cast(npts)); + for (int64_t i = 0; i < npts; ++i) + field[static_cast(i)] = std::sin(1e-3 * static_cast(i)); + + const std::string path = make_temp_path(".cube"); + write_cube(path, mol, grid, field.data(), "chunked"); + + std::ifstream in(path, std::ios::binary); + REQUIRE(in.is_open()); + std::string skip; + for (size_t i = 0; i < 6 + mol.size(); ++i) std::getline(in, skip); + + std::string data_block((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + in.close(); + + const int64_t bytes_per_row = grid.nz * 13 + (grid.nz + 5) / 6; + const int64_t n_rows = grid.nx * grid.ny; + REQUIRE(data_block.size() == static_cast(bytes_per_row * n_rows)); + + // Full byte comparison catches any misplacement at a chunk seam. + std::string expected; + expected.reserve(data_block.size()); + char buf[32]; + for (int64_t row = 0; row < n_rows; ++row) { + for (int64_t iz = 0; iz < grid.nz; ++iz) { + std::snprintf(buf, sizeof(buf), "%13.5E", + field[static_cast(row * grid.nz + iz)]); + expected += buf; + if ((iz + 1) % 6 == 0 || iz + 1 == grid.nz) expected += '\n'; + } + } + CHECK(data_block == expected); std::remove(path.c_str()); } +TEST_CASE("write_cube formatter rounds half-way ties within one last digit", + "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + // The hand-rolled formatter rounds exact half-way ties in the 6th + // significant digit half-away-from-zero, whereas glibc rounds the exact + // binary value half-to-even. Both are within one unit of the last printed + // digit; this pins that bound rather than byte equality. + auto mol = make_water(); + const std::vector field = {123456.5, 1.234575, -123456.5, -1.234575, + 9.9999949999999998642e-98, 0.0}; + + const std::string data_block = cube_data_block(mol, field); + REQUIRE(data_block.size() >= field.size() * 13); + + for (size_t i = 0; i < field.size(); ++i) { + const std::string tok = data_block.substr(i * 13, 13); + const double got = std::stod(tok); + const double v = field[i]; + const double last_digit = + v == 0.0 ? 1.0 + : std::pow(10.0, std::floor(std::log10(std::fabs(v))) - 5.0); + INFO("value " << v << " formatted as '" << tok << "'"); + CHECK(std::fabs(got - v) <= 0.5000001 * last_digit); + } +} + #ifdef GAUXC_HAS_HDF5 TEST_CASE("write_cube_hdf5 round-trip", "[cube]") { - if (!is_root_rank()) return; // I/O on /tmp; avoid MPI rank collisions. +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif auto mol = make_water(); auto grid = CubeGrid::from_molecule(mol, 4, 5, 6); diff --git a/tests/ut_common.hpp.in b/tests/ut_common.hpp.in index f899aeef..b18e67ce 100644 --- a/tests/ut_common.hpp.in +++ b/tests/ut_common.hpp.in @@ -18,3 +18,4 @@ #include #cmakedefine GAUXC_REF_DATA_PATH "@GAUXC_REF_DATA_PATH@" +#cmakedefine GAUXC_TEST_TMP_PATH "@GAUXC_TEST_TMP_PATH@" From 760f9188e190686c35e988c7563d50841e09d371 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Mon, 3 Aug 2026 14:25:03 -0700 Subject: [PATCH 10/31] OrbitalEvaluator: tile grid batches, fix batch sizing, expose screening 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. --- include/gauxc/orbital_evaluator.hpp | 15 +- src/orbital_evaluator.cxx | 399 ++++++++++++++++++++++------ src/orbital_evaluator_impl.hpp | 2 +- tests/orbital_evaluator_test.cxx | 98 +++++-- 4 files changed, 414 insertions(+), 100 deletions(-) diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index a510ee85..072a16ae 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -43,6 +43,13 @@ namespace detail { * thread count, so results are bit-reproducible for a fixed thread count but * may differ between thread counts by at most the basis-set shell tolerance. * + * Screening cost is governed by the shell tolerance, which the evaluator + * applies to its own private copy of the basis, so the caller's basis (and + * any SCF setup sharing it) is left untouched. Orbital error tracks the + * tolerance directly, density error goes as its square, so density tolerates + * a looser setting for the same accuracy: on a 110-atom system 1e-6 instead + * of the 1e-10 default runs ~2.5x faster on density. + * * Instances are produced by OrbitalEvaluatorFactory, which selects the * implementation for a given ExecutionSpace. */ @@ -149,9 +156,15 @@ class OrbitalEvaluatorFactory { * Currently only ExecutionSpace::Host is implemented; * anything else throws. * @param[in] basis Basis set (copied into the evaluator). + * @param[in] screening_tolerance Shell tolerance applied to the evaluator's + * own copy of `basis`, which sets the cutoff radii used + * for per-batch screening. A basis carried in from an SCF + * setup may be far tighter than cube-file precision needs + * and costs several-fold here for no benefit. */ static OrbitalEvaluator make_orbital_evaluator( ExecutionSpace ex, - BasisSet basis ); + BasisSet basis, + double screening_tolerance = detail::default_shell_tolerance ); }; // class OrbitalEvaluatorFactory diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index b3e68cc3..1b10b65e 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -34,9 +35,14 @@ namespace GauXC { namespace detail { -OrbitalEvaluatorImpl::OrbitalEvaluatorImpl( BasisSet bs ) : +OrbitalEvaluatorImpl::OrbitalEvaluatorImpl( BasisSet bs, + double screening_tolerance ) : basis( std::move(bs) ) { + // Retune our own copy: the caller's basis, which an SCF setup may share, + // is never touched. + for( auto& sh : basis ) sh.set_shell_tolerance( screening_tolerance ); + driver_owner = LocalWorkDriverFactory::make_local_work_driver( ExecutionSpace::Host, "Reference" ); host_driver = dynamic_cast( driver_owner.get() ); @@ -70,19 +76,25 @@ namespace { /** @brief Choose the number of points evaluated per batch. * * Two constraints, whichever is tighter: - * 1. Cache footprint: the AO scratch (nbf*batch doubles) should stay - * around 1 MiB per thread so it lives comfortably in L2/L3. + * 1. Cache footprint. Each thread holds `nscratch` live blocks of + * nbf*batch doubles: the AO block, the collocation kernel's own + * transpose staging, and, for density, the D*AO product. They are + * sized to share ~1 MiB, a typical per-core L2, rather than allowing + * each block that much on its own. * 2. Load balance: enough batches that each thread gets several, i.e. * batch <= npts / (4*nthreads). */ -size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads ) { - constexpr size_t kTargetAOBytesPerThread = 1024 * 1024; +size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads, + int nscratch ) { + constexpr size_t kTargetScratchBytesPerThread = 1024 * 1024; constexpr size_t kMinBatch = 128; constexpr size_t kMaxBatch = 8192; size_t batch = kMaxBatch; if( nbf > 0 ) { - batch = kTargetAOBytesPerThread / ( sizeof(double) * static_cast(nbf) ); + batch = kTargetScratchBytesPerThread / + ( sizeof(double) * static_cast(nbf) * + static_cast(nscratch) ); } if( nthreads > 0 ) { batch = std::min( batch, npts / ( 4 * static_cast(nthreads) ) ); @@ -110,36 +122,27 @@ PointBbox compute_bbox( const double* points, size_t npts ) { return b; } -/** @brief Bbox of the contiguous CubeGrid index range [p0, p0+np). - * - * Exact (not merely conservative): with iz varying fastest, a range that - * crosses an ix boundary necessarily contains both iy=0 and iy=ny-1, and a - * range that crosses an iy boundary necessarily contains both iz=0 and - * iz=nz-1. Computing this analytically avoids rescanning the 3*np - * coordinates that were just generated. - */ -PointBbox bbox_for_range( const CubeGrid& g, size_t p0, size_t np ) { +/// Bbox of the contiguous CubeGrid index range [k0, k1]. Exact: with iz +/// fastest, a range crossing an ix boundary contains both iy=0 and iy=ny-1, +/// and one crossing an iy boundary contains both iz=0 and iz=nz-1. +PointBbox bbox_for_index_range( const CubeGrid& g, int64_t k0, int64_t k1 ) { const int64_t nz = g.nz, ny = g.ny; - const int64_t k0 = static_cast(p0); - const int64_t k1 = static_cast(p0 + np) - 1; - int64_t lo_idx[3], hi_idx[3]; lo_idx[0] = k0 / (ny * nz); hi_idx[0] = k1 / (ny * nz); if( hi_idx[0] > lo_idx[0] ) { - lo_idx[1] = 0; hi_idx[1] = ny - 1; - lo_idx[2] = 0; hi_idx[2] = nz - 1; + lo_idx[1] = 0; hi_idx[1] = ny - 1; + lo_idx[2] = 0; hi_idx[2] = nz - 1; } else { lo_idx[1] = (k0 / nz) % ny; hi_idx[1] = (k1 / nz) % ny; if( hi_idx[1] > lo_idx[1] ) { - lo_idx[2] = 0; hi_idx[2] = nz - 1; + lo_idx[2] = 0; hi_idx[2] = nz - 1; } else { lo_idx[2] = k0 % nz; hi_idx[2] = k1 % nz; } } - PointBbox b; for( int k = 0; k < 3; ++k ) { const double a = g.origin[k] + g.spacing[k] * static_cast(lo_idx[k]); @@ -167,37 +170,69 @@ double dist2_center_to_bbox( const double* center, const PointBbox& bbox ) { return d2; } -/// Point source backed by a caller-supplied AoS coordinate array. +/** @brief The output indices covered by one batch. + * + * `nruns` runs of `run_len` consecutive indices starting at run_off[i]. A + * contiguous batch is the single-run case; a grid tile contributes one run + * per (ix,iy) pair. + */ +struct BatchSpan { + const size_t* run_off = nullptr; + size_t nruns = 0; + size_t run_len = 0; + size_t npts = 0; +}; + +/// Batch source over a caller-supplied AoS coordinate array. struct RawPointSource { static constexpr bool needs_scratch = false; const double* points; + size_t npts_total; + size_t batch_size; - const double* batch( size_t p0, size_t, double* ) const { - return points + 3 * p0; + size_t num_batches() const { + return ( npts_total + batch_size - 1 ) / batch_size; } - PointBbox bbox( size_t, size_t np, const double* pts ) const { - return compute_bbox( pts, np ); + size_t max_points() const { return batch_size; } + + const double* batch( size_t b, double*, std::vector& runs, + BatchSpan& span, PointBbox& bbox ) const { + const size_t p0 = b * batch_size; + const size_t np = std::min( batch_size, npts_total - p0 ); + runs.assign( 1, p0 ); + span = BatchSpan{ runs.data(), 1, np, np }; + const double* pts = points + 3 * p0; + bbox = compute_bbox( pts, np ); + return pts; } }; -/// Point source that generates CubeGrid coordinates on the fly, avoiding the -/// 3*num_points()*8 byte temporary coordinate array. -struct GridPointSource { +/// Batch source that walks a CubeGrid in contiguous index ranges, generating +/// coordinates on the fly. One run per batch, so results land in `out` +/// without a scatter. +struct GridLinearSource { static constexpr bool needs_scratch = true; const CubeGrid* grid; + size_t npts_total; + size_t batch_size; - const double* batch( size_t p0, size_t np, double* scr ) const { - // Incremental (ix,iy,iz) counters; coordinates are recomputed from the - // index (rather than accumulated) to stay bit-identical to - // CubeGrid::points_into. + size_t num_batches() const { + return ( npts_total + batch_size - 1 ) / batch_size; + } + size_t max_points() const { return batch_size; } + + const double* batch( size_t b, double* scr, std::vector& runs, + BatchSpan& span, PointBbox& bbox ) const { const CubeGrid& g = *grid; + const size_t p0 = b * batch_size; + const size_t np = std::min( batch_size, npts_total - p0 ); + const int64_t nz = g.nz, ny = g.ny; int64_t iz = static_cast(p0) % nz; int64_t iy = (static_cast(p0) / nz) % ny; int64_t ix = static_cast(p0) / (ny * nz); double x = g.origin[0] + g.spacing[0] * static_cast(ix); double y = g.origin[1] + g.spacing[1] * static_cast(iy); - for( size_t i = 0; i < np; ++i ) { scr[3 * i + 0] = x; scr[3 * i + 1] = y; @@ -212,14 +247,118 @@ struct GridPointSource { y = g.origin[1] + g.spacing[1] * static_cast(iy); } } + + runs.assign( 1, p0 ); + span = BatchSpan{ runs.data(), 1, np, np }; + bbox = bbox_for_index_range( g, static_cast(p0), + static_cast(p0 + np) - 1 ); return scr; } +}; - PointBbox bbox( size_t p0, size_t np, const double* ) const { - return bbox_for_range( *grid, p0, np ); +/** @brief Batch source that walks a CubeGrid in spatially compact tiles. + * + * A contiguous index range is a needle along z: as soon as it crosses a row + * boundary its bbox spans the whole z extent, so distant shells survive + * screening. A tile of the same point count has a far tighter bbox, which + * cuts nbe and therefore both the collocation (~nbe) and the density + * contraction (~nbe^2). The cost is that a tile's points are no longer + * contiguous in the output. + */ +struct GridTileSource { + static constexpr bool needs_scratch = true; + const CubeGrid* grid; + int64_t tx, ty, tz; ///< tile extent in grid points + int64_t ntx, nty, ntz; ///< tile counts per axis + + size_t num_batches() const { + return static_cast(ntx) * static_cast(nty) * + static_cast(ntz); + } + size_t max_points() const { + return static_cast(tx) * static_cast(ty) * + static_cast(tz); + } + + const double* batch( size_t b, double* scr, std::vector& runs, + BatchSpan& span, PointBbox& bbox ) const { + const CubeGrid& g = *grid; + const int64_t bi = static_cast(b); + const int64_t ix0 = ( bi / (nty * ntz) ) * tx; + const int64_t iy0 = ( ( bi / ntz ) % nty ) * ty; + const int64_t iz0 = ( bi % ntz ) * tz; + const int64_t ix1 = std::min( ix0 + tx, g.nx ); + const int64_t iy1 = std::min( iy0 + ty, g.ny ); + const int64_t iz1 = std::min( iz0 + tz, g.nz ); + const int64_t nzr = iz1 - iz0; + + runs.clear(); + size_t n = 0; + for( int64_t ix = ix0; ix < ix1; ++ix ) { + const double x = g.origin[0] + g.spacing[0] * static_cast(ix); + for( int64_t iy = iy0; iy < iy1; ++iy ) { + const double y = g.origin[1] + g.spacing[1] * static_cast(iy); + runs.push_back( + static_cast( (ix * g.ny + iy) * g.nz + iz0 ) ); + for( int64_t iz = iz0; iz < iz1; ++iz ) { + scr[3 * n + 0] = x; + scr[3 * n + 1] = y; + scr[3 * n + 2] = + g.origin[2] + g.spacing[2] * static_cast(iz); + ++n; + } + } + } + span = BatchSpan{ runs.data(), runs.size(), + static_cast(nzr), n }; + + const int64_t lo_idx[3] = { ix0, iy0, iz0 }; + const int64_t hi_idx[3] = { ix1 - 1, iy1 - 1, iz1 - 1 }; + for( int k = 0; k < 3; ++k ) { + const double a = + g.origin[k] + g.spacing[k] * static_cast(lo_idx[k]); + const double c = + g.origin[k] + g.spacing[k] * static_cast(hi_idx[k]); + bbox.lo[k] = std::min(a, c); + bbox.hi[k] = std::max(a, c); + } + return scr; } }; +/// Tile extents holding roughly `target` points and as close to cubic in +/// physical space as the (generally anisotropic) grid spacing allows. +GridTileSource make_tile_source( const CubeGrid& g, size_t target ) { + const int64_t n[3] = { g.nx, g.ny, g.nz }; + std::array s{}; + for( int k = 0; k < 3; ++k ) { + s[k] = std::fabs( g.spacing[k] ); + if( !(s[k] > 0.0) ) s[k] = 1.0; // degenerate axis: treat as unconstrained + } + + std::array t{ 1, 1, 1 }; + double scale = std::cbrt( static_cast(target) * s[0] * s[1] * s[2] ); + for( int attempt = 0; attempt < 4; ++attempt ) { + int64_t prod = 1; + for( int k = 0; k < 3; ++k ) { + t[k] = std::clamp( + static_cast( std::llround( scale / s[k] ) ), 1, n[k] ); + prod *= t[k]; + } + if( static_cast(prod) <= target ) break; + scale *= std::cbrt( static_cast(target) / + static_cast(prod) ); + } + + GridTileSource src; + src.grid = &g; + src.tx = t[0]; src.ty = t[1]; src.tz = t[2]; + src.ntx = (n[0] + t[0] - 1) / t[0]; + src.nty = (n[1] + t[1] - 1) / t[1]; + src.ntz = (n[2] + t[2] - 1) / t[2]; + return src; +} + /// Contraction of the AO batch against MO coefficients: /// out(np,nmo) = ao^T(np,nbe) @ C(nbe,nmo). class OrbitalContractor { @@ -233,7 +372,13 @@ class OrbitalContractor { public: - struct Scratch { std::vector C_compressed; }; + /// nbf*batch blocks this contraction keeps live, beyond the AO block. + static constexpr int scratch_blocks = 0; + + struct Scratch { + std::vector C_compressed; + std::vector staging; + }; OrbitalContractor( const detail::OrbitalEvaluatorImpl& impl, int32_t nmo, const double* C, size_t ldc, double* out, size_t ldo ) : @@ -243,16 +388,20 @@ class OrbitalContractor { scr.C_compressed.resize( static_cast(impl_.nbf_) * nmo_ ); } - void zero( size_t p0, size_t np ) const { + void zero( const BatchSpan& span ) const { for( int32_t j = 0; j < nmo_; ++j ) { - double* out_col = out_ + static_cast(j) * ldo_ + p0; - std::fill( out_col, out_col + np, 0.0 ); + double* out_col = out_ + static_cast(j) * ldo_; + for( size_t r = 0; r < span.nruns; ++r ) + std::fill( out_col + span.run_off[r], + out_col + span.run_off[r] + span.run_len, 0.0 ); } } - void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + void apply( Scratch& scr, const BatchSpan& span, int32_t nbe, const std::vector& shells, const double* ao ) const { + const size_t np = span.npts; + // Gather the rows of C for the surviving shells into a contiguous // (nbe, nmo) col-major buffer so the contraction is a dense GEMM. const BasisSetMap& basis_map = *impl_.basis_map; @@ -268,9 +417,26 @@ class OrbitalContractor { row += rng.second - rng.first; } + if( span.nruns == 1 ) { + blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, + ao, nbe, scr.C_compressed.data(), nbe, 0.0, + out_ + span.run_off[0], static_cast(ldo_) ); + return; + } + + const size_t need = np * static_cast(nmo_); + if( scr.staging.size() < need ) scr.staging.resize( need ); blas::gemm( 'T', 'N', static_cast(np), nmo_, nbe, 1.0, - ao, nbe, scr.C_compressed.data(), nbe, 0.0, out_ + p0, - static_cast(ldo_) ); + ao, nbe, scr.C_compressed.data(), nbe, 0.0, + scr.staging.data(), static_cast(np) ); + + for( int32_t j = 0; j < nmo_; ++j ) { + const double* src = scr.staging.data() + static_cast(j) * np; + double* out_col = out_ + static_cast(j) * ldo_; + for( size_t r = 0; r < span.nruns; ++r ) + std::copy( src + r * span.run_len, src + (r + 1) * span.run_len, + out_col + span.run_off[r] ); + } } }; @@ -286,32 +452,40 @@ class DensityContractor { public: + /// dm_ao is a second nbf*batch block alongside the AO block. + static constexpr int scratch_blocks = 1; + struct Scratch { std::vector dm_ao; std::vector xmat_scr; + std::vector staging; }; DensityContractor( const detail::OrbitalEvaluatorImpl& impl, const double* D, size_t ldd, double* out ) : impl_(impl), D_(D), ldd_(ldd), out_(out) {} - void init( Scratch& scr, size_t batch_size ) const { - scr.dm_ao.resize( static_cast(impl_.nbf_) * batch_size ); - } + void init( Scratch&, size_t ) const {} - void zero( size_t p0, size_t np ) const { - std::fill( out_ + p0, out_ + p0 + np, 0.0 ); + void zero( const BatchSpan& span ) const { + for( size_t r = 0; r < span.nruns; ++r ) + std::fill( out_ + span.run_off[r], + out_ + span.run_off[r] + span.run_len, 0.0 ); } - void apply( Scratch& scr, size_t p0, size_t np, int32_t nbe, + void apply( Scratch& scr, const BatchSpan& span, int32_t nbe, const std::vector& shells, const double* ao ) const { - // eval_xmat needs nbe*nbe scratch. Growing it on demand inside the - // parallel region keeps the high-water mark at the largest nbe this - // thread actually saw (not nbf) and first-touches the pages on the - // owning thread. + const size_t np = span.npts; + + // eval_xmat needs nbe*nbe scratch and writes an (nbe,np) block. Growing + // both on demand inside the parallel region keeps the high-water mark at + // the largest nbe this thread actually saw (not nbf) and first-touches + // the pages on the owning thread. const size_t scr_size = static_cast(nbe) * nbe; if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); + const size_t dm_ao_size = static_cast(nbe) * np; + if( scr.dm_ao.size() < dm_ao_size ) scr.dm_ao.resize( dm_ao_size ); const int32_t nbf = impl_.nbf_; LocalHostWorkDriver::submat_map_t submat_map; @@ -325,8 +499,19 @@ class DensityContractor { scr.dm_ao.data(), static_cast(nbe), scr.xmat_scr.data() ); // rho[p] = sum_mu ao(mu, p) * dm_ao(mu, p) + double* den = out_ + span.run_off[0]; + if( span.nruns > 1 ) { + if( scr.staging.size() < np ) scr.staging.resize( np ); + den = scr.staging.data(); + } impl_.host_driver->eval_uvvar_lda_rks( np, static_cast(nbe), ao, - scr.dm_ao.data(), static_cast(nbe), out_ + p0 ); + scr.dm_ao.data(), static_cast(nbe), den ); + + if( span.nruns > 1 ) { + for( size_t r = 0; r < span.nruns; ++r ) + std::copy( den + r * span.run_len, den + (r + 1) * span.run_len, + out_ + span.run_off[r] ); + } } }; @@ -338,40 +523,40 @@ class DensityContractor { * result is handed to `contract` for the orbital- or density-specific * reduction. All scratch is per-thread and per-call. */ -template -void batched_eval( const detail::OrbitalEvaluatorImpl& impl, size_t npts, - const PointSource& src, const Contractor& contract ) { +template +void batched_eval( const detail::OrbitalEvaluatorImpl& impl, + const BatchSource& src, const Contractor& contract ) { const BasisSet& basis = impl.basis; const auto& shell_cutoff_r2 = impl.shell_cutoff_r2; - const int32_t nbf = impl.nbf_; const int32_t nshells_total = basis.nshells(); - const size_t batch_size = choose_batch_size( nbf, npts, omp_get_max_threads() ); - const int64_t n_batches = - static_cast( (npts + batch_size - 1) / batch_size ); + const size_t max_pts = src.max_points(); + const int64_t n_batches = static_cast( src.num_batches() ); #pragma omp parallel { std::vector pt_buf; - if constexpr( PointSource::needs_scratch ) pt_buf.resize( 3 * batch_size ); - std::vector ao_buf( static_cast(nbf) * batch_size ); + if constexpr( BatchSource::needs_scratch ) pt_buf.resize( 3 * max_pts ); + std::vector run_storage; + std::vector ao_buf; std::vector screened_shells; screened_shells.reserve( nshells_total ); typename Contractor::Scratch scr; - contract.init( scr, batch_size ); + contract.init( scr, max_pts ); #pragma omp for schedule(dynamic, 1) for( int64_t b = 0; b < n_batches; ++b ) { - const size_t p0 = static_cast(b) * batch_size; - const size_t np = std::min( batch_size, npts - p0 ); - - const double* pts = src.batch( p0, np, pt_buf.data() ); + BatchSpan span; + PointBbox bbox; + const double* pts = + src.batch( static_cast(b), pt_buf.data(), run_storage, span, + bbox ); + const size_t np = span.npts; // Per-batch shell screening: keep shells whose cutoff radius reaches // any point of this batch's bounding box. - const PointBbox bbox = src.bbox( p0, np, pts ); screened_shells.clear(); int32_t nbe = 0; for( int32_t s = 0; s < nshells_total; ++s ) @@ -381,7 +566,12 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, size_t npts, } // All shells screen out -> the result is identically zero here. - if( not nbe ) { contract.zero( p0, np ); continue; } + if( not nbe ) { contract.zero( span ); continue; } + + // Only the surviving nbe rows are written, so the AO block tracks the + // largest nbe seen rather than the nbf worst case. + const size_t ao_size = static_cast(nbe) * np; + if( ao_buf.size() < ao_size ) ao_buf.resize( ao_size ); // Both contractions assume eval_collocation packs AO rows by cumulative // shell size in shell_list order. That holds for the gau2grid path; the @@ -393,13 +583,44 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, size_t npts, static_cast(nbe), pts, basis, screened_shells.data(), ao_buf.data() ); - contract.apply( scr, p0, np, nbe, screened_shells, ao_buf.data() ); + contract.apply( scr, span, nbe, screened_shells, ao_buf.data() ); } } } +/** @brief Whether spatial tiling can improve screening for this grid. + * + * A contiguous batch's bbox spans the entire z extent as soon as it crosses + * a row boundary. If that extent already lies within a shell cutoff radius, + * every shell reaching any part of the column reaches all of it, so a + * tighter box removes nothing and the scatter it costs is pure overhead. + * Tiling is therefore worthwhile only when the grid spans well beyond the + * interaction range along z. The factor of two is a margin, not a fit: at a + * ratio near one there is nothing to gain. + */ +bool should_tile( const detail::OrbitalEvaluatorImpl& impl, const CubeGrid& g ) { + if( impl.shell_cutoff_r2.empty() ) return false; + std::vector r2( impl.shell_cutoff_r2 ); + const auto mid = r2.begin() + r2.size() / 2; + std::nth_element( r2.begin(), mid, r2.end() ); + const double median_radius = std::sqrt( *mid ); + const double z_extent = + std::fabs( g.spacing[2] ) * static_cast( g.nz ); + return z_extent > 2.0 * median_radius; +} + +/// Target points per batch for a given contraction, accounting for every +/// nbf*batch block it keeps live alongside the AO block. +template +size_t batch_target( int32_t nbf, size_t npts ) { + constexpr int collocation_scratch_blocks = 1; + const int nscratch = + 1 + collocation_scratch_blocks + Contractor::scratch_blocks; + return choose_batch_size( nbf, npts, omp_get_max_threads(), nscratch ); +} + void check_orbital_args( const std::string& ctx, size_t npts, int32_t nbf, const double* C, size_t ldc, const double* out, size_t ldo ) { @@ -443,13 +664,14 @@ const BasisSet& OrbitalEvaluator::basis() const { OrbitalEvaluator OrbitalEvaluatorFactory::make_orbital_evaluator( - ExecutionSpace ex, BasisSet basis ) { + ExecutionSpace ex, BasisSet basis, double screening_tolerance ) { switch(ex) { case ExecutionSpace::Host: return OrbitalEvaluator( - std::make_unique( std::move(basis) ) + std::make_unique( std::move(basis), + screening_tolerance ) ); default: @@ -481,10 +703,27 @@ void OrbitalEvaluator::eval_orbitals( size_t npts, const double* points, check_orbital_args( "OrbitalEvaluator::eval_orbitals", npts, pimpl_->nbf_, C, ldc, out, ldo ); - batched_eval( *pimpl_, npts, RawPointSource{points}, + batched_eval( *pimpl_, + RawPointSource{ points, npts, + batch_target( pimpl_->nbf_, npts ) }, OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); } +namespace { + +/// Run a grid evaluation with whichever traversal screens better. +template +void eval_on_grid( const detail::OrbitalEvaluatorImpl& impl, + const CubeGrid& grid, size_t npts, size_t target, + const Contractor& contract ) { + if( should_tile( impl, grid ) ) + batched_eval( impl, make_tile_source( grid, target ), contract ); + else + batched_eval( impl, GridLinearSource{ &grid, npts, target }, contract ); +} + +} // namespace + void OrbitalEvaluator::eval_density( size_t npts, const double* points, const double* D, size_t ldd, double* out ) const { @@ -495,7 +734,9 @@ void OrbitalEvaluator::eval_density( size_t npts, const double* points, check_density_args( "OrbitalEvaluator::eval_density", pimpl_->nbf_, D, ldd, out ); - batched_eval( *pimpl_, npts, RawPointSource{points}, + batched_eval( *pimpl_, + RawPointSource{ points, npts, + batch_target( pimpl_->nbf_, npts ) }, DensityContractor( *pimpl_, D, ldd, out ) ); } @@ -520,7 +761,8 @@ void OrbitalEvaluator::eval_orbitals( const CubeGrid& grid, int32_t nmo, check_orbital_args( "OrbitalEvaluator::eval_orbitals(grid)", npts, pimpl_->nbf_, C, ldc, out, ldo ); - batched_eval( *pimpl_, npts, GridPointSource{&grid}, + eval_on_grid( *pimpl_, grid, npts, + batch_target( pimpl_->nbf_, npts ), OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); } @@ -531,7 +773,8 @@ void OrbitalEvaluator::eval_density( const CubeGrid& grid, const double* D, check_density_args( "OrbitalEvaluator::eval_density(grid)", pimpl_->nbf_, D, ldd, out ); - batched_eval( *pimpl_, npts, GridPointSource{&grid}, + eval_on_grid( *pimpl_, grid, npts, + batch_target( pimpl_->nbf_, npts ), DensityContractor( *pimpl_, D, ldd, out ) ); } diff --git a/src/orbital_evaluator_impl.hpp b/src/orbital_evaluator_impl.hpp index fab041ea..60fe6b38 100644 --- a/src/orbital_evaluator_impl.hpp +++ b/src/orbital_evaluator_impl.hpp @@ -35,7 +35,7 @@ class OrbitalEvaluatorImpl { std::vector shell_cutoff_r2; ///< Per-shell squared cutoff radius int32_t nbf_ = 0; - OrbitalEvaluatorImpl( BasisSet bs ); + OrbitalEvaluatorImpl( BasisSet bs, double screening_tolerance ); }; // class OrbitalEvaluatorImpl diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 65c59a75..8ba5981e 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -52,9 +52,9 @@ std::string make_temp_path(const char* suffix) { std::to_string(++counter) + suffix; } -OrbitalEvaluator make_evaluator(const BasisSet& basis) { +OrbitalEvaluator make_evaluator(const BasisSet& basis, double tol) { return OrbitalEvaluatorFactory::make_orbital_evaluator(ExecutionSpace::Host, - basis); + basis, tol); } /// AO collocation over every shell, i.e. the unscreened reference. @@ -116,7 +116,7 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", // Reference: AO collocation directly via the LocalHostWorkDriver. const auto ao_ref = reference_collocation(basis, npts, pts.data()); - auto eval = make_evaluator(basis); + auto eval = make_evaluator(basis, 1e-12); REQUIRE(eval.nbf() == nbf); SECTION("eval_orbital with one-hot coefficient reproduces single AO column") { @@ -217,7 +217,7 @@ TEST_CASE("CubeGrid eval overloads match pointer-based eval", auto basis = make_ccpvdz(mol, SphericalType(true)); for (auto& sh : basis) sh.set_shell_tolerance(1e-12); const int32_t nbf = basis.nbf(); - auto eval = make_evaluator(basis); + auto eval = make_evaluator(basis, 1e-12); // 6x7x8 fits in a single batch; 20x20x20 = 8000 points spreads over many // batches at any thread count, exercising the trailing partial batch and @@ -275,9 +275,10 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); auto basis = make_ccpvdz(mol, SphericalType(true)); - for (auto& sh : basis) sh.set_shell_tolerance(1e-10); + constexpr double shell_tol = 1e-10; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); const int32_t nbf = basis.nbf(); - auto eval = make_evaluator(basis); + auto eval = make_evaluator(basis, shell_tol); const auto C = make_random_vector(static_cast(nbf), 2024u); std::vector D(static_cast(nbf) * nbf, 0.0); @@ -319,9 +320,9 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { std::vector rho(static_cast(npts)); eval.eval_density(grid, D.data(), nbf, rho.data()); - // The grid overload screens against an analytically derived batch bbox - // while the pointer overload scans the coordinates; with screening active - // the two must still agree exactly. + // The grid overload walks spatial tiles while the pointer overload takes + // contiguous index ranges, so the two screen against different bounding + // boxes. They agree to within the shell tolerance, not bitwise. std::vector orb_pts(static_cast(npts)); eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); std::vector rho_pts(static_cast(npts)); @@ -329,8 +330,8 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { bool any_nonzero = false; for (int64_t p = 0; p < npts; ++p) { - CHECK(orb[p] == orb_pts[p]); - CHECK(rho[p] == rho_pts[p]); + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); double orb_ref = 0.0, rho_ref = 0.0; for (int32_t mu = 0; mu < nbf; ++mu) { @@ -348,14 +349,16 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { } } -TEST_CASE("OrbitalEvaluator is invariant to the OpenMP thread count", +TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", "[orbital_evaluator]") { - // Bit-exact here because the grid encloses the molecule at a 1e-12 shell - // tolerance, so nothing screens and the batch decomposition cannot change - // the arithmetic. The screened case is covered separately below. + // Regression guard: scratch must follow the thread count in force at + // evaluation time, not at construction. Batch shape is derived from the + // thread count, so the two runs screen slightly differently and agree to + // within the shell tolerance rather than bitwise. auto mol = make_water(); auto basis = make_ccpvdz(mol, SphericalType(true)); - for (auto& sh : basis) sh.set_shell_tolerance(1e-12); + constexpr double shell_tol = 1e-12; + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); const int32_t nbf = basis.nbf(); auto grid = CubeGrid::from_molecule(mol, 16, 16, 16); @@ -371,7 +374,7 @@ TEST_CASE("OrbitalEvaluator is invariant to the OpenMP thread count", // must follow the thread count in force at evaluation time. omp_set_num_threads(1); #endif - auto eval = make_evaluator(basis); + auto eval = make_evaluator(basis, shell_tol); std::vector orb_serial(static_cast(npts)); std::vector rho_serial(static_cast(npts)); @@ -392,8 +395,8 @@ TEST_CASE("OrbitalEvaluator is invariant to the OpenMP thread count", #endif for (int64_t p = 0; p < npts; ++p) { - CHECK(orb_par[p] == orb_serial[p]); - CHECK(rho_par[p] == rho_serial[p]); + CHECK(orb_par[p] == Approx(orb_serial[p]).margin(shell_tol)); + CHECK(rho_par[p] == Approx(rho_serial[p]).margin(shell_tol)); } } @@ -410,7 +413,7 @@ TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", constexpr double shell_tol = 1e-10; for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); const int32_t nbf = basis.nbf(); - auto eval = make_evaluator(basis); + auto eval = make_evaluator(basis, shell_tol); CubeGrid grid; grid.origin = {-4.0, -4.0, -4.0}; @@ -444,6 +447,61 @@ TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", } } +TEST_CASE("OrbitalEvaluator screening error scales with the shell tolerance", + "[orbital_evaluator]") { + // Pins the guidance on the class: orbital error is bounded by the shell + // tolerance, density error by its square, since dropping a shell with + // |phi| < t perturbs sum_uv D phi phi by ~t^2. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); + + constexpr double loose_tol = 1e-6; + auto basis = make_ccpvdz(mol, SphericalType(true)); + const int32_t nbf = basis.nbf(); + + std::vector radii_before; + for (const auto& sh : basis) radii_before.push_back(sh.cutoff_radius()); + + auto eval_tight = make_evaluator(basis, 1e-14); + auto eval_loose = make_evaluator(basis, loose_tol); + + // Retuning happens on each evaluator's own copy; a basis shared with an SCF + // setup must come back unchanged. + for (size_t i = 0; i < basis.size(); ++i) + CHECK(basis[i].cutoff_radius() == radii_before[i]); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -4.0}; + grid.spacing = {0.8, 2.0, 2.0}; + grid.nx = 40; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + + const auto C = make_random_vector(static_cast(nbf), 31337u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb_tight(npts), orb_loose(npts); + std::vector rho_tight(npts), rho_loose(npts); + eval_tight.eval_orbital(grid, C.data(), orb_tight.data()); + eval_loose.eval_orbital(grid, C.data(), orb_loose.data()); + eval_tight.eval_density(grid, D.data(), nbf, rho_tight.data()); + eval_loose.eval_density(grid, D.data(), nbf, rho_loose.data()); + + double orb_err = 0.0, rho_err = 0.0; + for (int64_t p = 0; p < npts; ++p) { + orb_err = std::max(orb_err, std::fabs(orb_loose[p] - orb_tight[p])); + rho_err = std::max(rho_err, std::fabs(rho_loose[p] - rho_tight[p])); + } + + INFO("orbital error " << orb_err << ", density error " << rho_err); + CHECK(orb_err > 0.0); // screening is genuinely active at the loose tolerance + CHECK(orb_err <= 10.0 * loose_tol); + CHECK(rho_err <= 100.0 * loose_tol * loose_tol); +} + TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { auto mol = make_water(); CubeGrid g = CubeGrid::from_molecule(mol, /*nx=*/16, /*ny=*/12, /*nz=*/8, From 2bbf17b886d33aa0b46b8869c524f1b971efe095 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Mon, 3 Aug 2026 14:36:38 -0700 Subject: [PATCH 11/31] Tests: cover the tiled grid traversal 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. --- tests/orbital_evaluator_test.cxx | 60 ++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 8ba5981e..10db56e4 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -349,6 +349,66 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { } } +TEST_CASE("OrbitalEvaluator tiled grid traversal matches the reference", + "[orbital_evaluator]") { + // Grids spanning well beyond a cutoff radius along z are walked in spatial + // tiles rather than contiguous index ranges. The 32 Bohr z extent here is + // several times the largest cc-pVDZ cutoff radius (13.2 Bohr), so the tiled + // path stays selected; the two centres are far enough apart that screening + // is active within it. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -6.0}; + grid.spacing = {2.0, 2.0, 0.8}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 40; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 8675u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + // The pointer overload always uses contiguous batches, so it cross-checks + // the tiled result against a different decomposition. + std::vector orb_pts(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); + + bool any_nonzero = false; + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + if (std::fabs(orb_ref) > 1e-3) any_nonzero = true; + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); + } + CHECK(any_nonzero); +} + TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", "[orbital_evaluator]") { // Regression guard: scratch must follow the thread count in force at From ce3edd802b6ddd6c6d65c3d82cb95098e901b962 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Mon, 3 Aug 2026 14:51:57 -0700 Subject: [PATCH 12/31] Cube writer: drop unreachable branches in format_e13_5, cover the rest 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. --- src/external/cube.cxx | 53 ++++++++++++++------------------ tests/orbital_evaluator_test.cxx | 40 ++++++++++++++++++------ 2 files changed, 53 insertions(+), 40 deletions(-) diff --git a/src/external/cube.cxx b/src/external/cube.cxx index 626169ee..f6a47562 100644 --- a/src/external/cube.cxx +++ b/src/external/cube.cxx @@ -27,12 +27,12 @@ namespace { /** @brief Format `v` as a fixed 13-character "%13.5E" field. * - * Layout: sign + D + '.' + 5 digits + 'E' + sign + 2 digits. snprintf in - * the inner loop dominates write time on large grids; this reproduces glibc - * "%13.5E" for every value except exact half-way ties in the 6th significant - * digit, where glibc rounds the exact binary value half-to-even while this - * routine rounds half-away-from-zero. Both agree to within one unit of the - * last printed digit. 3-digit exponents defer to snprintf. + * Layout: sign + D + '.' + 5 digits + 'E' + sign + 2 digits. snprintf in the + * inner loop dominates write time on large grids. This matches glibc + * "%13.5E" except for values lying within a few ulp of a rounding boundary + * in the 6th significant digit, where scaling the mantissa by pow(10,-exp10) + * can tip it across the boundary; the two then differ by one unit in the last + * printed digit. Values needing a 3-digit exponent defer to snprintf. * * @param[in] v Value to format. * @param[out] out Exactly 13 bytes are written (no trailing NUL). @@ -57,36 +57,29 @@ inline void format_e13_5( double v, char* out ) { return; } - // Exponent via floor(log10), with corrections for FP edge cases - // (e.g. 9.99999 rounding up across a power-of-ten boundary). int exp10 = static_cast( std::floor( std::log10(absv) ) ); - double scale = std::pow( 10.0, -exp10 ); - double mant = absv * scale; - - long long mant_int = std::llround( mant * 1e5 ); - if( mant_int >= 1000000 ) { - mant_int = 100000; - ++exp10; - } else if( mant_int < 100000 ) { - --exp10; - scale = std::pow( 10.0, -exp10 ); - mant = absv * scale; + + // Tested before scaling: pow(10,-exp10) overflows for subnormal inputs. + bool wide_exponent = ( exp10 > 99 or exp10 < -99 ); + + long long mant_int = 0; + if( not wide_exponent ) { + const double mant = absv * std::pow( 10.0, -exp10 ); + // Over this range pow() is accurate to a few ulp, so mant cannot fall + // below 1 and only the carry out of 9.999995 needs correcting. mant_int = std::llround( mant * 1e5 ); - if( mant_int >= 1000000 ) mant_int = 999999; - else if( mant_int < 100000 ) mant_int = 100000; + if( mant_int >= 1000000 ) { + mant_int = 100000; + wide_exponent = ( ++exp10 > 99 ); + } } - // 3+ digit exponents have a different field layout; defer to snprintf. - if( exp10 > 99 or exp10 < -99 ) { + // 3-digit exponents have a different field layout; defer to snprintf. + if( wide_exponent ) { char tmp[32]; + // "%13.5E" sets a minimum field width of 13, so n is never below 13. const int n = std::snprintf( tmp, sizeof(tmp), "%13.5E", v ); - if( n >= 13 ) { - std::memcpy( out, tmp + (n - 13), 13 ); - } else { - const int pad = 13 - n; - for( int i = 0; i < pad; ++i ) out[i] = ' '; - std::memcpy( out + pad, tmp, static_cast(n) ); - } + std::memcpy( out, tmp + (n - 13), 13 ); return; } diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 10db56e4..218a1386 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -748,11 +748,30 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { const double qnan = std::numeric_limits::quiet_NaN(); const double inf = std::numeric_limits::infinity(); - std::vector field = {0.0, -0.0, 1.23456e-10, -9.99995e-1, - 1.0e+99, -1.0e+99, 3.14159265358979, - -2.71828, 1.0, -1.0, 1e-300, - 1.234e+05, qnan, -qnan, inf, - -inf}; + // The last five exercise the snprintf deferral: subnormals (where scaling + // the mantissa would overflow), a 3-digit exponent either side of zero, and + // a mantissa that carries from E+99 up into a 3-digit exponent. + std::vector field = {0.0, + -0.0, + 1.23456e-10, + -9.99995e-1, + 1.0e+99, + -1.0e+99, + 3.14159265358979, + -2.71828, + 1.0, + -1.0, + 1e-300, + 1.234e+05, + qnan, + -qnan, + inf, + -inf, + std::numeric_limits::denorm_min(), + -1e-310, + 1.0e+300, + -1.0e-305, + 9.9999999e+99}; const std::string data_block = cube_data_block(mol, field); @@ -821,17 +840,18 @@ TEST_CASE("write_cube spans multiple output chunks", "[cube]") { std::remove(path.c_str()); } -TEST_CASE("write_cube formatter rounds half-way ties within one last digit", +TEST_CASE("write_cube formatter stays within one last digit at a rounding " + "boundary", "[cube]") { #ifdef GAUXC_HAS_MPI int world_rank; MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); if (world_rank) return; // File I/O; only run on root rank #endif - // The hand-rolled formatter rounds exact half-way ties in the 6th - // significant digit half-away-from-zero, whereas glibc rounds the exact - // binary value half-to-even. Both are within one unit of the last printed - // digit; this pins that bound rather than byte equality. + // For values sitting within a few ulp of a rounding boundary in the 6th + // significant digit, scaling the mantissa can tip it across the boundary, + // so the formatter and glibc may pick different last digits. Both stay + // within one unit of it; this pins that bound rather than byte equality. auto mol = make_water(); const std::vector field = {123456.5, 1.234575, -123456.5, -1.234575, 9.9999949999999998642e-98, 0.0}; From bace779954475085f03c08044f4f90807eb150f7 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Mon, 3 Aug 2026 15:05:11 -0700 Subject: [PATCH 13/31] Tests: cover the partial-row branch of the grid batch bounding box 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. --- tests/orbital_evaluator_test.cxx | 86 ++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 218a1386..1ee06868 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -12,6 +12,7 @@ #include "ut_common.hpp" #include "catch2/catch.hpp" +#include #include #include #include @@ -409,6 +410,91 @@ TEST_CASE("OrbitalEvaluator tiled grid traversal matches the reference", CHECK(any_nonzero); } +TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", + "[orbital_evaluator]") { + // With iz fastest, a contiguous batch shorter than one row is boxed by its + // own z sub-range rather than the whole axis. Three things have to line up + // for that box to matter. The row must outlast a batch, so nz exceeds the + // batch size. The grid must stay off the tiled path, so the z extent sits + // inside the tiling threshold of twice the median cutoff radius, derived + // here from the basis rather than hard-coded. And the molecule must sit at + // the far end of z, so that shrinking the box moves the face nearest the + // shells and changes what survives screening. + auto mol = make_water(); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + + std::vector radii; + radii.reserve(basis.size()); + for (const auto& sh : basis) radii.push_back(sh.cutoff_radius()); + const size_t mid = radii.size() / 2; + std::nth_element(radii.begin(), radii.begin() + mid, radii.end()); + const double median_radius = radii[mid]; + + const double z_extent = 1.8 * median_radius; + + CubeGrid grid; + grid.nx = 1; + grid.ny = 3; + grid.nz = 3000; + grid.spacing = {1.0, 1.0, z_extent / static_cast(grid.nz)}; + // Offset in x so no point lands on a nucleus; the row runs from z_extent + // below the molecule up to 1 Bohr short of it. + grid.origin = {0.3, 0.0, -(z_extent + 1.0)}; + + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 4099u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + std::vector orb_pts(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + +#ifdef _OPENMP + const int saved_threads = omp_get_max_threads(); + // Serial so the batch is npts/4 rather than thread-count dependent, which + // puts three quarters of a row in one batch on any machine. + omp_set_num_threads(1); +#endif + auto eval = make_evaluator(basis, shell_tol); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + // The pointer overload boxes the points it is handed instead of deriving a + // box from grid indices, so it checks the index arithmetic independently. + eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); + eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); +#ifdef _OPENMP + omp_set_num_threads(saved_threads); +#endif + + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao_ref[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + CHECK(orb[p] == Approx(orb_pts[p]).margin(shell_tol)); + CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); + } + + // Pins the span the batch box is being asked to resolve: the row ends on + // the molecule and starts far enough away for shells to have died off. + const size_t row_end = static_cast(grid.nz) - 1; + CHECK(rho[row_end] > 1e-2); + CHECK(rho[0] < 1e-4 * rho[row_end]); +} + TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", "[orbital_evaluator]") { // Regression guard: scratch must follow the thread count in force at From 8a53f7eea3606596931dedd3fa6955aa264625dd Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 07:52:35 -0700 Subject: [PATCH 14/31] OrbitalEvaluator: record why the C-gather bypass is not taken 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. --- src/orbital_evaluator.cxx | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 1b10b65e..309a5cd6 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -404,6 +404,9 @@ class OrbitalContractor { // Gather the rows of C for the surviving shells into a contiguous // (nbe, nmo) col-major buffer so the contraction is a dense GEMM. + // Skipping this when no shell is screened out was measured and is not + // worth it: that case is 0-5% of batches on anything larger than water, + // and it is precisely the case where there is nothing to gather. const BasisSetMap& basis_map = *impl_.basis_map; int32_t row = 0; for( int32_t ish : shells ) { From 4e407f01e04810c7a76453e6b742165e34450a19 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 08:02:02 -0700 Subject: [PATCH 15/31] OrbitalEvaluator: drop unused init parameter, note gau2grid is mandatory 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. --- src/orbital_evaluator.cxx | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 309a5cd6..a1cd4401 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -384,7 +384,7 @@ class OrbitalContractor { const double* C, size_t ldc, double* out, size_t ldo ) : impl_(impl), nmo_(nmo), C_(C), ldc_(ldc), out_(out), ldo_(ldo) {} - void init( Scratch& scr, size_t ) const { + void init( Scratch& scr ) const { scr.C_compressed.resize( static_cast(impl_.nbf_) * nmo_ ); } @@ -468,7 +468,7 @@ class DensityContractor { size_t ldd, double* out ) : impl_(impl), D_(D), ldd_(ldd), out_(out) {} - void init( Scratch&, size_t ) const {} + void init( Scratch& ) const {} void zero( const BatchSpan& span ) const { for( size_t r = 0; r < span.nruns; ++r ) @@ -546,7 +546,7 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, std::vector screened_shells; screened_shells.reserve( nshells_total ); typename Contractor::Scratch scr; - contract.init( scr, max_pts ); + contract.init( scr ); #pragma omp for schedule(dynamic, 1) for( int64_t b = 0; b < n_batches; ++b ) { @@ -580,8 +580,9 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, // shell size in shell_list order. That holds for the gau2grid path; the // non-gau2grid fallback in gau2grid_collocation.cxx instead uses the // global shell_to_first_ao offset, which is only equivalent when no - // shell is screened out. This is a pre-existing inconsistency in the - // reference driver rather than one introduced here. + // shell is screened out. That fallback is unreachable in any supported + // build (the top-level CMakeLists makes gau2grid a hard dependency), and + // the discrepancy is pre-existing rather than introduced here. impl.host_driver->eval_collocation( np, screened_shells.size(), static_cast(nbe), pts, basis, screened_shells.data(), ao_buf.data() ); From 652a93a45a04e561bc7a13518de4fa4d2c2eb2fb Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 09:34:28 -0700 Subject: [PATCH 16/31] OrbitalEvaluator: grow the gathered C buffer like every other scratch 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. --- src/orbital_evaluator.cxx | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index a1cd4401..b9813636 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -384,10 +384,6 @@ class OrbitalContractor { const double* C, size_t ldc, double* out, size_t ldo ) : impl_(impl), nmo_(nmo), C_(C), ldc_(ldc), out_(out), ldo_(ldo) {} - void init( Scratch& scr ) const { - scr.C_compressed.resize( static_cast(impl_.nbf_) * nmo_ ); - } - void zero( const BatchSpan& span ) const { for( int32_t j = 0; j < nmo_; ++j ) { double* out_col = out_ + static_cast(j) * ldo_; @@ -407,6 +403,9 @@ class OrbitalContractor { // Skipping this when no shell is screened out was measured and is not // worth it: that case is 0-5% of batches on anything larger than water, // and it is precisely the case where there is nothing to gather. + const size_t c_size = static_cast(nbe) * nmo_; + if( scr.C_compressed.size() < c_size ) scr.C_compressed.resize( c_size ); + const BasisSetMap& basis_map = *impl_.basis_map; int32_t row = 0; for( int32_t ish : shells ) { @@ -468,8 +467,6 @@ class DensityContractor { size_t ldd, double* out ) : impl_(impl), D_(D), ldd_(ldd), out_(out) {} - void init( Scratch& ) const {} - void zero( const BatchSpan& span ) const { for( size_t r = 0; r < span.nruns; ++r ) std::fill( out_ + span.run_off[r], @@ -546,7 +543,6 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, std::vector screened_shells; screened_shells.reserve( nshells_total ); typename Contractor::Scratch scr; - contract.init( scr ); #pragma omp for schedule(dynamic, 1) for( int64_t b = 0; b < n_batches; ++b ) { From ca7601b2a4737ed8d9113246d5b6e395656f3202 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 09:43:57 -0700 Subject: [PATCH 17/31] cube.hpp: include gauxc_config.hpp for the GAUXC_HAS_HDF5 guard 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. --- include/gauxc/external/cube.hpp | 1 + 1 file changed, 1 insertion(+) diff --git a/include/gauxc/external/cube.hpp b/include/gauxc/external/cube.hpp index 78d29611..dd11fc4f 100644 --- a/include/gauxc/external/cube.hpp +++ b/include/gauxc/external/cube.hpp @@ -13,6 +13,7 @@ #include +#include #include #include From f492468002c5901301d0f33048b424a6a97ed76c Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 10:03:10 -0700 Subject: [PATCH 18/31] tests: cover the multi-orbital scatter and single-plane grids 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. --- tests/orbital_evaluator_test.cxx | 103 +++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 1ee06868..5f1c9250 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -410,6 +410,109 @@ TEST_CASE("OrbitalEvaluator tiled grid traversal matches the reference", CHECK(any_nonzero); } +TEST_CASE("OrbitalEvaluator scatters multiple orbitals into a padded output", + "[orbital_evaluator]") { + // A tiled batch is not contiguous in the output, so its result is staged and + // scattered column by column. nmo > 1 with a padded ldo is the only + // combination that exercises both strides of that scatter. Padding entries + // hold sentinels: C must never be read past nbf, out never written past npts. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(1), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-4.0, -4.0, -6.0}; + grid.spacing = {2.0, 2.0, 0.8}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 40; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao_ref = reference_collocation(basis, npts, pts.data()); + + constexpr int32_t nmo = 3; + constexpr double sentinel = -12345.0; + const size_t ldc = static_cast(nbf) + 7; + const size_t ldo = static_cast(npts) + 5; + + const auto Craw = make_random_vector(static_cast(nbf) * nmo, 99u); + std::vector C(ldc * nmo, sentinel); + for (int32_t j = 0; j < nmo; ++j) + for (int32_t mu = 0; mu < nbf; ++mu) + C[j * ldc + mu] = Craw[static_cast(j) * nbf + mu]; + + std::vector out(ldo * nmo, sentinel); + eval.eval_orbitals(grid, nmo, C.data(), ldc, out.data(), ldo); + + // The pointer overload always batches contiguously, so it checks the tiled + // scatter against a different decomposition of the same grid. + std::vector out_pts(ldo * nmo, sentinel); + eval.eval_orbitals(npts, pts.data(), nmo, C.data(), ldc, out_pts.data(), ldo); + + for (int32_t j = 0; j < nmo; ++j) { + for (int64_t p = 0; p < npts; ++p) { + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) + ref += Craw[static_cast(j) * nbf + mu] * + ao_ref[static_cast(p) * nbf + mu]; + const size_t k = static_cast(j) * ldo + static_cast(p); + CHECK(out[k] == Approx(ref).margin(1e-9)); + CHECK(out[k] == Approx(out_pts[k]).margin(shell_tol)); + } + for (size_t p = static_cast(npts); p < ldo; ++p) { + CHECK(out[static_cast(j) * ldo + p] == sentinel); + CHECK(out_pts[static_cast(j) * ldo + p] == sentinel); + } + } +} + +TEST_CASE("OrbitalEvaluator handles single-plane grids", "[orbital_evaluator]") { + // An axis of one point gets zero spacing from from_molecule, which the + // tile-shaping code must treat as unconstrained rather than divide by. + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + const auto C = make_random_vector(static_cast(nbf), 4242u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + const std::vector grids = {CubeGrid::from_molecule(mol, 1, 9, 9), + CubeGrid::from_molecule(mol, 9, 1, 9), + CubeGrid::from_molecule(mol, 9, 9, 1)}; + + for (const auto& grid : grids) { + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao = reference_collocation(basis, npts, pts.data()); + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + for (int64_t p = 0; p < npts; ++p) { + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + const double a = ao[static_cast(p) * nbf + mu]; + orb_ref += C[static_cast(mu)] * a; + rho_ref += a * a; + } + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + } + } +} + TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", "[orbital_evaluator]") { // With iz fastest, a contiguous batch shorter than one row is boxed by its From e747520690ddebe56af5a47dcfa6183ba4e68b7d Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 10:16:17 -0700 Subject: [PATCH 19/31] tests: cover the rest of the public API surface 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. --- tests/orbital_evaluator_test.cxx | 154 +++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 5f1c9250..d90c1d2f 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -212,6 +212,108 @@ TEST_CASE("OrbitalEvaluator / Water cc-pVDZ matches eval_collocation", } } +TEST_CASE("OrbitalEvaluator public API surface", "[orbital_evaluator]") { + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + + SECTION("basis() exposes the evaluator's own retuned copy") { + auto tight = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : tight) sh.set_shell_tolerance(1e-12); + auto eval = make_evaluator(tight, 1e-4); + + REQUIRE(eval.basis().nbf() == nbf); + REQUIRE(eval.basis().size() == tight.size()); + + // The caller's basis keeps the tolerance it was given, so the evaluator's + // looser one must show up as shorter cutoff radii on its copy alone. + bool any_shorter = false; + for (size_t i = 0; i < tight.size(); ++i) { + CHECK(eval.basis()[i].cutoff_radius() <= tight[i].cutoff_radius()); + if (eval.basis()[i].cutoff_radius() < tight[i].cutoff_radius()) + any_shorter = true; + } + CHECK(any_shorter); + } + + SECTION("move construction and assignment carry the implementation") { + const int64_t npts = 64; + const auto pts = make_random_points(npts, 7u); + const auto C = make_random_vector(static_cast(nbf), 11u); + + auto eval = make_evaluator(basis, shell_tol); + std::vector ref(static_cast(npts)); + eval.eval_orbital(npts, pts.data(), C.data(), ref.data()); + + OrbitalEvaluator moved(std::move(eval)); + REQUIRE(moved.nbf() == nbf); + std::vector out(static_cast(npts)); + moved.eval_orbital(npts, pts.data(), C.data(), out.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(out[p] == ref[p]); + + // Assigned over an evaluator built with a looser tolerance, so the result + // would shift if the target's original implementation survived. + auto target = make_evaluator(basis, 1e-3); + target = std::move(moved); + REQUIRE(target.nbf() == nbf); + std::vector out2(static_cast(npts)); + target.eval_orbital(npts, pts.data(), C.data(), out2.data()); + for (int64_t p = 0; p < npts; ++p) CHECK(out2[p] == ref[p]); + } + + SECTION("a non-Host execution space is rejected") { + CHECK_THROWS(OrbitalEvaluatorFactory::make_orbital_evaluator( + ExecutionSpace::Device, basis)); + } + + SECTION("null pointers and undersized leading dimensions throw") { + auto eval = make_evaluator(basis, shell_tol); + const int64_t npts = 8; + const auto pts = make_random_points(npts, 3u); + std::vector C(static_cast(nbf), 1.0); + std::vector out(static_cast(npts), 0.0); + std::vector D(static_cast(nbf) * nbf, 0.0); + + CHECK_THROWS(eval.eval_orbital(npts, nullptr, C.data(), out.data())); + CHECK_THROWS(eval.eval_orbital(npts, pts.data(), nullptr, out.data())); + CHECK_THROWS(eval.eval_orbital(npts, pts.data(), C.data(), nullptr)); + CHECK_THROWS(eval.eval_density(npts, nullptr, D.data(), nbf, out.data())); + CHECK_THROWS(eval.eval_density(npts, pts.data(), nullptr, nbf, out.data())); + CHECK_THROWS(eval.eval_density(npts, pts.data(), D.data(), nbf, nullptr)); + CHECK_THROWS( + eval.eval_density(npts, pts.data(), D.data(), nbf - 1, out.data())); + + const auto grid = CubeGrid::from_molecule(mol, 3, 3, 3); + std::vector gout(static_cast(grid.num_points()), 0.0); + CHECK_THROWS(eval.eval_orbital(grid, nullptr, gout.data())); + CHECK_THROWS(eval.eval_density(grid, D.data(), nbf - 1, gout.data())); + } + + SECTION("empty work is a no-op rather than an error") { + auto eval = make_evaluator(basis, shell_tol); + std::vector C(static_cast(nbf), 1.0); + std::vector D(static_cast(nbf) * nbf, 0.0); + constexpr double untouched = -7.0; + std::vector out(4, untouched); + const auto pts = make_random_points(4, 5u); + + CHECK_NOTHROW( + eval.eval_orbitals(0, nullptr, 1, C.data(), nbf, out.data(), 0)); + CHECK_NOTHROW(eval.eval_density(0, nullptr, D.data(), nbf, out.data())); + CHECK_NOTHROW( + eval.eval_orbitals(4, pts.data(), 0, C.data(), nbf, out.data(), 4)); + + CubeGrid empty; + empty.nx = 0; + CHECK_NOTHROW(eval.eval_orbital(empty, C.data(), out.data())); + CHECK_NOTHROW(eval.eval_density(empty, D.data(), nbf, out.data())); + + for (double v : out) CHECK(v == untouched); + } +} + TEST_CASE("CubeGrid eval overloads match pointer-based eval", "[orbital_evaluator]") { auto mol = make_water(); @@ -784,6 +886,58 @@ TEST_CASE("CubeGrid construction and grid-points layout", "[cube]") { CHECK(pts[3 * 1 + 2] == Approx(g.origin[2] + g.spacing[2])); } +TEST_CASE("CubeGrid rejects degenerate specifications", "[cube]") { + auto mol = make_water(); + CHECK_THROWS(CubeGrid::from_molecule(Molecule{}, 4, 4, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 0, 4, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 4, 0, 4)); + CHECK_THROWS(CubeGrid::from_molecule(mol, 4, 4, 0)); + + const auto grid = CubeGrid::from_molecule(mol, 3, 4, 5); + const auto pts = grid.points(); + REQUIRE(pts.size() == static_cast(grid.num_points()) * 3); + + std::vector buf(pts.size(), -1.0); + grid.points_into(buf.data()); + for (size_t i = 0; i < pts.size(); ++i) CHECK(buf[i] == pts[i]); +} + +TEST_CASE("write_cube rejects bad input and defaults its comment", "[cube]") { +#ifdef GAUXC_HAS_MPI + int world_rank; + MPI_Comm_rank(MPI_COMM_WORLD, &world_rank); + if (world_rank) return; // File I/O; only run on root rank +#endif + auto mol = make_water(); + const auto grid = CubeGrid::from_molecule(mol, 2, 2, 2); + std::vector field(static_cast(grid.num_points()), 1.0); + const auto path = make_temp_path(".cube"); + + CHECK_THROWS(write_cube(path, mol, grid, nullptr)); + + CubeGrid empty; + empty.nx = 0; + CHECK_THROWS(write_cube(path, mol, empty, field.data())); + CHECK_THROWS( + write_cube("/nonexistent-gauxc-dir/out.cube", mol, grid, field.data())); + +#ifdef GAUXC_HAS_HDF5 + const auto h5 = make_temp_path(".h5"); + CHECK_THROWS(write_cube_hdf5(h5, mol, grid, nullptr)); + CHECK_THROWS(write_cube_hdf5(h5, mol, empty, field.data())); +#endif + + // An omitted comment falls back to a fixed first line. + write_cube(path, mol, grid, field.data()); + std::ifstream in(path); + REQUIRE(in.good()); + std::string line; + std::getline(in, line); + CHECK(line == "GauXC cube file"); + std::getline(in, line); + CHECK(line == "Generated by GauXC"); +} + TEST_CASE("write_cube round-trips header and field data", "[cube]") { #ifdef GAUXC_HAS_MPI int world_rank; From de68decf5583bf147cdd9770802b4ca1e54d098a Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 10:38:53 -0700 Subject: [PATCH 20/31] tests: exercise the density path with a general D and a padded ldd 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. --- tests/orbital_evaluator_test.cxx | 67 ++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index d90c1d2f..146a251d 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -615,6 +615,73 @@ TEST_CASE("OrbitalEvaluator handles single-plane grids", "[orbital_evaluator]") } } +TEST_CASE("OrbitalEvaluator density with a general D and padded ldd", + "[orbital_evaluator]") { + // Every other density test uses a diagonal D and ldd == nbf, and neither + // can see a whole class of defect. Permuting an identity on both sides + // leaves it unchanged, so a scrambled compressed-submatrix map is invisible + // to it; and ldd == nbf makes a leading dimension dropped in favour of nbf + // a no-op. Screening has to be active for the compressed submatrix to + // differ from D at all, hence the two well-separated centres. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 20.0); + + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-3.0, -3.0, -4.0}; + grid.spacing = {2.0, 2.0, 1.0}; + grid.nx = 4; + grid.ny = 4; + grid.nz = 28; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao = reference_collocation(basis, npts, pts.data()); + + constexpr double sentinel = -4321.0; + const size_t ldd = static_cast(nbf) + 6; + + const auto rnd = make_random_vector(static_cast(nbf) * nbf, 31337u); + std::vector Dsq(static_cast(nbf) * nbf); + for (int32_t i = 0; i < nbf; ++i) + for (int32_t j = 0; j < nbf; ++j) + Dsq[static_cast(j) * nbf + i] = + rnd[static_cast(j) * nbf + i] + + rnd[static_cast(i) * nbf + j]; + + // Rows past nbf are poisoned, so an ldd mistaken for nbf reads them. + std::vector D(ldd * nbf, sentinel); + for (int32_t j = 0; j < nbf; ++j) + for (int32_t i = 0; i < nbf; ++i) + D[static_cast(j) * ldd + i] = Dsq[static_cast(j) * nbf + i]; + + std::vector rho_grid(static_cast(npts)); + std::vector rho_pts(static_cast(npts)); + eval.eval_density(grid, D.data(), ldd, rho_grid.data()); + eval.eval_density(npts, pts.data(), D.data(), ldd, rho_pts.data()); + + double max_ref = 0.0; + for (int64_t p = 0; p < npts; ++p) { + const double* a = ao.data() + static_cast(p) * nbf; + double ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + double acc = 0.0; + for (int32_t nu = 0; nu < nbf; ++nu) + acc += Dsq[static_cast(nu) * nbf + mu] * a[nu]; + ref += acc * a[mu]; + } + max_ref = std::max(max_ref, std::fabs(ref)); + CHECK(rho_grid[p] == Approx(ref).margin(1e-8)); + CHECK(rho_pts[p] == Approx(ref).margin(1e-8)); + } + CHECK(max_ref > 1e-2); +} + TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", "[orbital_evaluator]") { // With iz fastest, a contiguous batch shorter than one row is boxed by its From e3e607c6bbb805f1d07556b3854214e017cf6c7b Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 11:01:25 -0700 Subject: [PATCH 21/31] OrbitalEvaluator: correct the documented cross-thread deviation 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. --- include/gauxc/orbital_evaluator.hpp | 5 ++++- tests/orbital_evaluator_test.cxx | 8 +++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index 072a16ae..15cfc8b1 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -41,7 +41,10 @@ namespace detail { * Points are evaluated in batches and each batch is screened against the * per-shell cutoff radii. The batch size is derived in part from the OpenMP * thread count, so results are bit-reproducible for a fixed thread count but - * may differ between thread counts by at most the basis-set shell tolerance. + * may differ between thread counts. Each neglected shell contributes at most + * the shell tolerance, so that difference is of the order of the tolerance + * rather than bounded by it: it grows with the number of shells dropped, and + * reaches roughly twice the tolerance on benzene in cc-pVDZ. * * Screening cost is governed by the shell tolerance, which the evaluator * applies to its own private copy of the basis, so the caller's basis (and diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 146a251d..7b545d01 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -822,8 +822,10 @@ TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", "[orbital_evaluator]") { // Batch size is derived from the thread count, so when screening is active // different thread counts screen against different batch bounding boxes and - // the results are not bit-identical. The discrepancy is bounded by the shell - // tolerance, which is what this pins down. + // the results are not bit-identical. Each dropped shell contributes at most + // the shell tolerance and the errors accumulate, so the bound scales with + // how many can be dropped; nbf is the generous ceiling on that. Still eight + // orders below the field itself, so it remains a real constraint. Molecule mol; mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); @@ -861,7 +863,7 @@ TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", #endif for (int64_t p = 0; p < npts; ++p) { - CHECK(rho_par[p] == Approx(rho_serial[p]).margin(shell_tol)); + CHECK(rho_par[p] == Approx(rho_serial[p]).margin(nbf * shell_tol)); } } From 0591bfe0e55a3fe25b8bfe172e8298a8fc232a2b Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 12:12:02 -0700 Subject: [PATCH 22/31] CubeGrid: note what an isotropic margin does to a flat molecule 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. --- include/gauxc/cube_grid.hpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/include/gauxc/cube_grid.hpp b/include/gauxc/cube_grid.hpp index 6b01f112..58511876 100644 --- a/include/gauxc/cube_grid.hpp +++ b/include/gauxc/cube_grid.hpp @@ -46,6 +46,12 @@ struct CubeGrid { * coincide with the extended bounding-box corners (PySCF cubegen * convention). * + * The margin is the same on all three axes, so a flat molecule gets a box + * that is thin in the direction it is flat in. An orbital that reaches out + * that way, such as the pi system of an aromatic ring, can still be large + * where the box ends and will look cut off when plotted. Use a bigger + * margin, or set origin, spacing and point counts yourself, if that matters. + * * @param mol Molecule whose atomic centres define the bounding box. * @param nx,ny,nz Number of grid points along each axis. * @param margin Margin (Bohr) added on each side. Default 3.0 matches From 2b71ef9a89f1fd70b7cec157f7d0ac2384f30358 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Tue, 4 Aug 2026 15:01:42 -0700 Subject: [PATCH 23/31] tests: pin the box a multi-row batch draws around itself 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. --- tests/orbital_evaluator_test.cxx | 77 ++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 7b545d01..11ab0eff 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -682,6 +682,83 @@ TEST_CASE("OrbitalEvaluator density with a general D and padded ldd", CHECK(max_ref > 1e-2); } +TEST_CASE("OrbitalEvaluator boxes a multi-row batch out to its far face", + "[orbital_evaluator]") { + // A batch spanning more than one row is boxed by the whole extent of the + // faster axes. Computing that box one grid step short would screen out a + // shell sitting on the far face, and those points would lose it entirely. + // + // One step has to be decisive for this to be visible at all. A shell only + // dropped by a one-step shrink lies between cutoff-step and cutoff of the + // box, where by construction it contributes about the shell tolerance, so a + // fine grid hides the error inside the tolerance it is measured against. + // The spacing along the axis under test is therefore set to 1.5x the largest + // cutoff radius, measured from the basis rather than assumed, which puts a + // shell on the far face either fully in or fully out. + constexpr double shell_tol = 1e-10; + + double radius = 0.0; + { + Molecule probe; + probe.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + auto pb = make_ccpvdz(probe, SphericalType(true)); + for (auto& sh : pb) sh.set_shell_tolerance(shell_tol); + for (const auto& sh : pb) radius = std::max(radius, sh.cutoff_radius()); + } + REQUIRE(radius > 1.0); + const double step = 1.5 * radius; + + // Only the y axis is worth testing this way. Making the z shrink decisive + // would need a z spacing above the cutoff radius, and a grid that coarse in + // z has a z extent past the tiling threshold, so it takes the tiled path and + // never builds this box at all. With nz == 1 the shrunk corner falls below + // the low corner and the min/max that follows widens the box instead of + // narrowing it. Either way a z shrink here cannot lose a shell. + Molecule mol; + mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); + mol.emplace_back(AtomicNumber(8), 0.0, 3.0 * step, 0.0); + + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + + CubeGrid grid; + grid.origin = {-7.0, 0.0, 0.0}; + grid.spacing = {2.0, step, 2.0}; + grid.nx = 8; + grid.ny = 4; + grid.nz = 4; + const int64_t npts = grid.num_points(); + const auto pts = grid.points(); + const auto ao = reference_collocation(basis, npts, pts.data()); + + const auto C = make_random_vector(static_cast(nbf), 606u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + // The far atom has to register on the grid, or nothing is being proved about + // whether its shells survived screening. + double max_rho = 0.0; + for (int64_t p = 0; p < npts; ++p) { + const double* a = ao.data() + static_cast(p) * nbf; + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + orb_ref += C[static_cast(mu)] * a[mu]; + rho_ref += a[mu] * a[mu]; + } + max_rho = std::max(max_rho, rho_ref); + CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); + CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); + } + CHECK(max_rho > 1e-2); +} + TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", "[orbital_evaluator]") { // With iz fastest, a contiguous batch shorter than one row is boxed by its From 8acc8988039a43c8e31d0f602c27b86d8fce21c3 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:30:27 -0700 Subject: [PATCH 24/31] OrbitalEvaluator: gather the density submatrix only when it is read 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. --- src/orbital_evaluator.cxx | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index b9813636..d43a6d3b 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -478,20 +478,23 @@ class DensityContractor { const size_t np = span.npts; - // eval_xmat needs nbe*nbe scratch and writes an (nbe,np) block. Growing - // both on demand inside the parallel region keeps the high-water mark at - // the largest nbe this thread actually saw (not nbf) and first-touches - // the pages on the owning thread. - const size_t scr_size = static_cast(nbe) * nbe; - if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); - const size_t dm_ao_size = static_cast(nbe) * np; - if( scr.dm_ao.size() < dm_ao_size ) scr.dm_ao.resize( dm_ao_size ); - const int32_t nbf = impl_.nbf_; LocalHostWorkDriver::submat_map_t submat_map; std::tie( submat_map, std::ignore ) = gen_compressed_submat_map( *impl_.basis_map, shells, nbf, nbf ); + // eval_xmat gathers D into scratch only when the surviving shells span + // more than one contiguous AO block; a single block is read in place from + // D itself, so nbe*nbe is never touched. Growing on demand inside the + // parallel region keeps the high-water mark at the largest nbe this thread + // actually saw (not nbf) and first-touches the pages on the owning thread. + if( submat_map.size() > 1 ) { + const size_t scr_size = static_cast(nbe) * nbe; + if( scr.xmat_scr.size() < scr_size ) scr.xmat_scr.resize( scr_size ); + } + const size_t dm_ao_size = static_cast(nbe) * np; + if( scr.dm_ao.size() < dm_ao_size ) scr.dm_ao.resize( dm_ao_size ); + // dm_ao = D_compressed @ ao impl_.host_driver->eval_xmat( np, static_cast(nbf), static_cast(nbe), submat_map, 1.0, D_, ldd_, From 725d09745386d56c35be390ccc6cb12bb4a50bd6 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:30:51 -0700 Subject: [PATCH 25/31] OrbitalEvaluator: reject an ldd beyond the range BLAS is handed 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. --- src/orbital_evaluator.cxx | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index d43a6d3b..c0359fb2 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -644,6 +644,9 @@ void check_density_args( const std::string& ctx, int32_t nbf, const double* D, GAUXC_GENERIC_EXCEPTION( ctx + ": null pointer argument." ); if( ldd < static_cast(nbf) ) GAUXC_GENERIC_EXCEPTION( ctx + ": ldd must be >= nbf()." ); + // ldd reaches blas::gemm through eval_xmat, which narrows it to int32_t. + if( ldd > static_cast(std::numeric_limits::max()) ) + GAUXC_GENERIC_EXCEPTION( ctx + ": ldd exceeds BLAS int range." ); } } // namespace From 237e6ccd9f10fd96b08fb36036325758c2e9df60 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:33:28 -0700 Subject: [PATCH 26/31] OrbitalEvaluator: always walk grids in spatial tiles 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. --- include/gauxc/orbital_evaluator.hpp | 6 ++ src/orbital_evaluator.cxx | 157 ++++++---------------------- tests/orbital_evaluator_test.cxx | 141 ++++++++++++------------- 3 files changed, 106 insertions(+), 198 deletions(-) diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index 15cfc8b1..ab6fd035 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -46,6 +46,12 @@ namespace detail { * rather than bounded by it: it grows with the number of shells dropped, and * reaches roughly twice the tolerance on benzene in cc-pVDZ. * + * For the same reason the CubeGrid overloads are not bitwise equal to the + * pointer overloads given the same points. A CubeGrid is walked in spatially + * compact tiles, which screen better, whereas a caller-supplied point array + * is batched in the order it arrives; the two decompositions drop different + * shells and so agree to the shell tolerance rather than bitwise. + * * Screening cost is governed by the shell tolerance, which the evaluator * applies to its own private copy of the basis, so the caller's basis (and * any SCF setup sharing it) is left untouched. Orbital error tracks the diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index c0359fb2..b2f8f3d4 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -122,37 +122,6 @@ PointBbox compute_bbox( const double* points, size_t npts ) { return b; } -/// Bbox of the contiguous CubeGrid index range [k0, k1]. Exact: with iz -/// fastest, a range crossing an ix boundary contains both iy=0 and iy=ny-1, -/// and one crossing an iy boundary contains both iz=0 and iz=nz-1. -PointBbox bbox_for_index_range( const CubeGrid& g, int64_t k0, int64_t k1 ) { - const int64_t nz = g.nz, ny = g.ny; - int64_t lo_idx[3], hi_idx[3]; - lo_idx[0] = k0 / (ny * nz); - hi_idx[0] = k1 / (ny * nz); - if( hi_idx[0] > lo_idx[0] ) { - lo_idx[1] = 0; hi_idx[1] = ny - 1; - lo_idx[2] = 0; hi_idx[2] = nz - 1; - } else { - lo_idx[1] = (k0 / nz) % ny; - hi_idx[1] = (k1 / nz) % ny; - if( hi_idx[1] > lo_idx[1] ) { - lo_idx[2] = 0; hi_idx[2] = nz - 1; - } else { - lo_idx[2] = k0 % nz; - hi_idx[2] = k1 % nz; - } - } - PointBbox b; - for( int k = 0; k < 3; ++k ) { - const double a = g.origin[k] + g.spacing[k] * static_cast(lo_idx[k]); - const double c = g.origin[k] + g.spacing[k] * static_cast(hi_idx[k]); - b.lo[k] = std::min(a, c); - b.hi[k] = std::max(a, c); - } - return b; -} - /// Squared distance from `center` to the nearest point of the bbox. Zero if /// the center lies inside. double dist2_center_to_bbox( const double* center, const PointBbox& bbox ) { @@ -207,63 +176,22 @@ struct RawPointSource { } }; -/// Batch source that walks a CubeGrid in contiguous index ranges, generating -/// coordinates on the fly. One run per batch, so results land in `out` -/// without a scatter. -struct GridLinearSource { - static constexpr bool needs_scratch = true; - const CubeGrid* grid; - size_t npts_total; - size_t batch_size; - - size_t num_batches() const { - return ( npts_total + batch_size - 1 ) / batch_size; - } - size_t max_points() const { return batch_size; } - - const double* batch( size_t b, double* scr, std::vector& runs, - BatchSpan& span, PointBbox& bbox ) const { - const CubeGrid& g = *grid; - const size_t p0 = b * batch_size; - const size_t np = std::min( batch_size, npts_total - p0 ); - - const int64_t nz = g.nz, ny = g.ny; - int64_t iz = static_cast(p0) % nz; - int64_t iy = (static_cast(p0) / nz) % ny; - int64_t ix = static_cast(p0) / (ny * nz); - double x = g.origin[0] + g.spacing[0] * static_cast(ix); - double y = g.origin[1] + g.spacing[1] * static_cast(iy); - for( size_t i = 0; i < np; ++i ) { - scr[3 * i + 0] = x; - scr[3 * i + 1] = y; - scr[3 * i + 2] = g.origin[2] + g.spacing[2] * static_cast(iz); - if( ++iz == nz ) { - iz = 0; - if( ++iy == ny ) { - iy = 0; - ++ix; - x = g.origin[0] + g.spacing[0] * static_cast(ix); - } - y = g.origin[1] + g.spacing[1] * static_cast(iy); - } - } - - runs.assign( 1, p0 ); - span = BatchSpan{ runs.data(), 1, np, np }; - bbox = bbox_for_index_range( g, static_cast(p0), - static_cast(p0 + np) - 1 ); - return scr; - } -}; - /** @brief Batch source that walks a CubeGrid in spatially compact tiles. * - * A contiguous index range is a needle along z: as soon as it crosses a row - * boundary its bbox spans the whole z extent, so distant shells survive - * screening. A tile of the same point count has a far tighter bbox, which + * A contiguous index range is a needle along z, and as soon as it crosses a + * row boundary its bbox spans the whole z extent, so distant shells survive + * screening. A tile of the same point count draws a far tighter box, which * cuts nbe and therefore both the collocation (~nbe) and the density * contraction (~nbe^2). The cost is that a tile's points are no longer - * contiguous in the output. + * contiguous in the output and have to be scattered. + * + * Tiling is the only grid traversal because a contiguous one was never + * measured to win. Counting nbe*np and nbe^2*np over every batch (a + * deterministic proxy for collocation and GEMM cost) across benzene and + * taxol, grids of 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, tiling never did + * more work: from a 2-3% saving on benzene at 128^3 to a 53% saving on taxol + * at 200^3. Paired single-threaded timings agreed in direction throughout. */ struct GridTileSource { static constexpr bool needs_scratch = true; @@ -338,8 +266,9 @@ GridTileSource make_tile_source( const CubeGrid& g, size_t target ) { std::array t{ 1, 1, 1 }; double scale = std::cbrt( static_cast(target) * s[0] * s[1] * s[2] ); + int64_t prod = 1; for( int attempt = 0; attempt < 4; ++attempt ) { - int64_t prod = 1; + prod = 1; for( int k = 0; k < 3; ++k ) { t[k] = std::clamp( static_cast( std::llround( scale / s[k] ) ), 1, n[k] ); @@ -350,6 +279,18 @@ GridTileSource make_tile_source( const CubeGrid& g, size_t target ) { static_cast(prod) ); } + // Each rescale shrinks the overshoot as r -> r^(2/3), so a grid anisotropic + // enough to start far above target is still above it after four. Halving the + // longest axis converges regardless, at the cost of a less cubic tile. + while( static_cast(prod) > target ) { + const int kmax = static_cast( + std::max_element( t.begin(), t.end() ) - t.begin() ); + if( t[kmax] == 1 ) break; // nothing left to halve + prod /= t[kmax]; + t[kmax] = ( t[kmax] + 1 ) / 2; + prod *= t[kmax]; + } + GridTileSource src; src.grid = &g; src.tx = t[0]; src.ty = t[1]; src.tz = t[2]; @@ -593,27 +534,6 @@ void batched_eval( const detail::OrbitalEvaluatorImpl& impl, } -/** @brief Whether spatial tiling can improve screening for this grid. - * - * A contiguous batch's bbox spans the entire z extent as soon as it crosses - * a row boundary. If that extent already lies within a shell cutoff radius, - * every shell reaching any part of the column reaches all of it, so a - * tighter box removes nothing and the scatter it costs is pure overhead. - * Tiling is therefore worthwhile only when the grid spans well beyond the - * interaction range along z. The factor of two is a margin, not a fit: at a - * ratio near one there is nothing to gain. - */ -bool should_tile( const detail::OrbitalEvaluatorImpl& impl, const CubeGrid& g ) { - if( impl.shell_cutoff_r2.empty() ) return false; - std::vector r2( impl.shell_cutoff_r2 ); - const auto mid = r2.begin() + r2.size() / 2; - std::nth_element( r2.begin(), mid, r2.end() ); - const double median_radius = std::sqrt( *mid ); - const double z_extent = - std::fabs( g.spacing[2] ) * static_cast( g.nz ); - return z_extent > 2.0 * median_radius; -} - /// Target points per batch for a given contraction, accounting for every /// nbf*batch block it keeps live alongside the AO block. template @@ -715,21 +635,6 @@ void OrbitalEvaluator::eval_orbitals( size_t npts, const double* points, OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); } -namespace { - -/// Run a grid evaluation with whichever traversal screens better. -template -void eval_on_grid( const detail::OrbitalEvaluatorImpl& impl, - const CubeGrid& grid, size_t npts, size_t target, - const Contractor& contract ) { - if( should_tile( impl, grid ) ) - batched_eval( impl, make_tile_source( grid, target ), contract ); - else - batched_eval( impl, GridLinearSource{ &grid, npts, target }, contract ); -} - -} // namespace - void OrbitalEvaluator::eval_density( size_t npts, const double* points, const double* D, size_t ldd, double* out ) const { @@ -767,8 +672,9 @@ void OrbitalEvaluator::eval_orbitals( const CubeGrid& grid, int32_t nmo, check_orbital_args( "OrbitalEvaluator::eval_orbitals(grid)", npts, pimpl_->nbf_, C, ldc, out, ldo ); - eval_on_grid( *pimpl_, grid, npts, - batch_target( pimpl_->nbf_, npts ), + batched_eval( *pimpl_, + make_tile_source( grid, + batch_target( pimpl_->nbf_, npts ) ), OrbitalContractor( *pimpl_, nmo, C, ldc, out, ldo ) ); } @@ -779,8 +685,9 @@ void OrbitalEvaluator::eval_density( const CubeGrid& grid, const double* D, check_density_args( "OrbitalEvaluator::eval_density(grid)", pimpl_->nbf_, D, ldd, out ); - eval_on_grid( *pimpl_, grid, npts, - batch_target( pimpl_->nbf_, npts ), + batched_eval( *pimpl_, + make_tile_source( grid, + batch_target( pimpl_->nbf_, npts ) ), DensityContractor( *pimpl_, D, ldd, out ) ); } diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 11ab0eff..f2e10e42 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -454,11 +454,11 @@ TEST_CASE("OrbitalEvaluator shell screening", "[orbital_evaluator]") { TEST_CASE("OrbitalEvaluator tiled grid traversal matches the reference", "[orbital_evaluator]") { - // Grids spanning well beyond a cutoff radius along z are walked in spatial - // tiles rather than contiguous index ranges. The 32 Bohr z extent here is - // several times the largest cc-pVDZ cutoff radius (13.2 Bohr), so the tiled - // path stays selected; the two centres are far enough apart that screening - // is active within it. + // Grids are walked in spatial tiles rather than contiguous index ranges, so + // a tile's points are scattered into the output. The two centres here are + // far enough apart that screening is active, and the 32 Bohr z extent is + // several times the largest cc-pVDZ cutoff radius (13.2 Bohr), so tiles at + // opposite ends of the grid see different shell sets. Molecule mol; mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); mol.emplace_back(AtomicNumber(1), 0.0, 0.0, 20.0); @@ -682,19 +682,25 @@ TEST_CASE("OrbitalEvaluator density with a general D and padded ldd", CHECK(max_ref > 1e-2); } -TEST_CASE("OrbitalEvaluator boxes a multi-row batch out to its far face", +TEST_CASE("OrbitalEvaluator boxes a batch out to its far face", "[orbital_evaluator]") { - // A batch spanning more than one row is boxed by the whole extent of the - // faster axes. Computing that box one grid step short would screen out a - // shell sitting on the far face, and those points would lose it entirely. + // A batch is screened against the box drawn round its own points. Building + // that box one grid step short on any axis would drop a shell sitting on the + // far face, and the points there would lose it entirely. // // One step has to be decisive for this to be visible at all. A shell only // dropped by a one-step shrink lies between cutoff-step and cutoff of the // box, where by construction it contributes about the shell tolerance, so a // fine grid hides the error inside the tolerance it is measured against. - // The spacing along the axis under test is therefore set to 1.5x the largest - // cutoff radius, measured from the basis rather than assumed, which puts a - // shell on the far face either fully in or fully out. + // The spacing is therefore set to 1.5x the largest cutoff radius, measured + // from the basis rather than assumed, which puts a shell on a tile face + // either fully in or fully out. + // + // Which grid point lands on a tile face depends on how the grid is + // subdivided, which is not observable from here and moves with the batch + // size, so the atom is swept over every grid point instead. Six points an + // axis puts it on a low face, a far face and an interior point of every + // axis under any tiling the batch size can produce. constexpr double shell_tol = 1e-10; double radius = 0.0; @@ -708,67 +714,65 @@ TEST_CASE("OrbitalEvaluator boxes a multi-row batch out to its far face", REQUIRE(radius > 1.0); const double step = 1.5 * radius; - // Only the y axis is worth testing this way. Making the z shrink decisive - // would need a z spacing above the cutoff radius, and a grid that coarse in - // z has a z extent past the tiling threshold, so it takes the tiled path and - // never builds this box at all. With nz == 1 the shrunk corner falls below - // the low corner and the min/max that follows widens the box instead of - // narrowing it. Either way a z shrink here cannot lose a shell. - Molecule mol; - mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); - mol.emplace_back(AtomicNumber(8), 0.0, 3.0 * step, 0.0); - - auto basis = make_ccpvdz(mol, SphericalType(true)); - for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); - const int32_t nbf = basis.nbf(); - auto eval = make_evaluator(basis, shell_tol); - CubeGrid grid; - grid.origin = {-7.0, 0.0, 0.0}; - grid.spacing = {2.0, step, 2.0}; - grid.nx = 8; - grid.ny = 4; - grid.nz = 4; + grid.origin = {0.0, 0.0, 0.0}; + grid.spacing = {step, step, step}; + grid.nx = 6; + grid.ny = 6; + grid.nz = 6; const int64_t npts = grid.num_points(); const auto pts = grid.points(); - const auto ao = reference_collocation(basis, npts, pts.data()); - const auto C = make_random_vector(static_cast(nbf), 606u); - std::vector D(static_cast(nbf) * nbf, 0.0); - for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + double worst_orb = 0.0, worst_rho = 0.0, max_rho = 0.0; + for (int64_t site = 0; site < npts; ++site) { + Molecule mol; + mol.emplace_back(AtomicNumber(8), pts[3 * site + 0], pts[3 * site + 1], + pts[3 * site + 2]); - std::vector orb(static_cast(npts)); - std::vector rho(static_cast(npts)); - eval.eval_orbital(grid, C.data(), orb.data()); - eval.eval_density(grid, D.data(), nbf, rho.data()); + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + const auto ao = reference_collocation(basis, npts, pts.data()); - // The far atom has to register on the grid, or nothing is being proved about - // whether its shells survived screening. - double max_rho = 0.0; - for (int64_t p = 0; p < npts; ++p) { - const double* a = ao.data() + static_cast(p) * nbf; - double orb_ref = 0.0, rho_ref = 0.0; - for (int32_t mu = 0; mu < nbf; ++mu) { - orb_ref += C[static_cast(mu)] * a[mu]; - rho_ref += a[mu] * a[mu]; + const auto C = make_random_vector(static_cast(nbf), 606u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + std::vector orb(static_cast(npts)); + std::vector rho(static_cast(npts)); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + + for (int64_t p = 0; p < npts; ++p) { + const double* a = ao.data() + static_cast(p) * nbf; + double orb_ref = 0.0, rho_ref = 0.0; + for (int32_t mu = 0; mu < nbf; ++mu) { + orb_ref += C[static_cast(mu)] * a[mu]; + rho_ref += a[mu] * a[mu]; + } + max_rho = std::max(max_rho, rho_ref); + worst_orb = std::max(worst_orb, std::fabs(orb[p] - orb_ref)); + worst_rho = std::max(worst_rho, std::fabs(rho[p] - rho_ref)); } - max_rho = std::max(max_rho, rho_ref); - CHECK(orb[p] == Approx(orb_ref).margin(1e-9)); - CHECK(rho[p] == Approx(rho_ref).margin(1e-9)); } + + CHECK(worst_orb < 1e-9); + CHECK(worst_rho < 1e-9); + // The atom has to register on the grid, or nothing is being proved about + // whether its shells survived screening. CHECK(max_rho > 1e-2); } -TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", +TEST_CASE("OrbitalEvaluator screens a tile spanning part of a long axis", "[orbital_evaluator]") { - // With iz fastest, a contiguous batch shorter than one row is boxed by its - // own z sub-range rather than the whole axis. Three things have to line up - // for that box to matter. The row must outlast a batch, so nz exceeds the - // batch size. The grid must stay off the tiled path, so the z extent sits - // inside the tiling threshold of twice the median cutoff radius, derived - // here from the basis rather than hard-coded. And the molecule must sit at - // the far end of z, so that shrinking the box moves the face nearest the - // shells and changes what survives screening. + // A tile covers a sub-range of each axis, so on a grid far longer in z than + // the tile is, the box has to close around that sub-range rather than the + // whole axis. Three things have to line up for that to matter. The axis must + // outlast a tile, so nz is far above any tile extent. The grid must span + // well past a cutoff radius, taken here from the basis rather than + // hard-coded. And the molecule must sit at the far end of z, so that a box + // reaching further than its own points changes what survives screening. auto mol = make_water(); constexpr double shell_tol = 1e-10; @@ -783,7 +787,7 @@ TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", std::nth_element(radii.begin(), radii.begin() + mid, radii.end()); const double median_radius = radii[mid]; - const double z_extent = 1.8 * median_radius; + const double z_extent = 4.0 * median_radius; CubeGrid grid; grid.nx = 1; @@ -807,22 +811,13 @@ TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", std::vector orb_pts(static_cast(npts)); std::vector rho_pts(static_cast(npts)); -#ifdef _OPENMP - const int saved_threads = omp_get_max_threads(); - // Serial so the batch is npts/4 rather than thread-count dependent, which - // puts three quarters of a row in one batch on any machine. - omp_set_num_threads(1); -#endif auto eval = make_evaluator(basis, shell_tol); eval.eval_orbital(grid, C.data(), orb.data()); eval.eval_density(grid, D.data(), nbf, rho.data()); // The pointer overload boxes the points it is handed instead of deriving a - // box from grid indices, so it checks the index arithmetic independently. + // box from tile indices, so it checks the index arithmetic independently. eval.eval_orbital(npts, pts.data(), C.data(), orb_pts.data()); eval.eval_density(npts, pts.data(), D.data(), nbf, rho_pts.data()); -#ifdef _OPENMP - omp_set_num_threads(saved_threads); -#endif for (int64_t p = 0; p < npts; ++p) { double orb_ref = 0.0, rho_ref = 0.0; @@ -837,7 +832,7 @@ TEST_CASE("OrbitalEvaluator screens a batch lying inside one grid row", CHECK(rho[p] == Approx(rho_pts[p]).margin(shell_tol)); } - // Pins the span the batch box is being asked to resolve: the row ends on + // Pins the span the tile box is being asked to resolve: the row ends on // the molecule and starts far enough away for shells to have died off. const size_t row_end = static_cast(grid.nz) - 1; CHECK(rho[row_end] > 1e-2); From 145f98c52ccafddcfa589a97b8f287cb64f42994 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:34:05 -0700 Subject: [PATCH 27/31] OrbitalEvaluator: say what the batch budget actually bounds 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. --- src/orbital_evaluator.cxx | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index b2f8f3d4..5a701152 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -76,13 +76,18 @@ namespace { /** @brief Choose the number of points evaluated per batch. * * Two constraints, whichever is tighter: - * 1. Cache footprint. Each thread holds `nscratch` live blocks of - * nbf*batch doubles: the AO block, the collocation kernel's own - * transpose staging, and, for density, the D*AO product. They are - * sized to share ~1 MiB, a typical per-core L2, rather than allowing - * each block that much on its own. + * 1. Cache footprint of the blocks that scale with the batch. Each thread + * holds `nscratch` live blocks of nbf*batch doubles: the AO block, the + * collocation kernel's own transpose staging, and, for density, the + * D*AO product. They are sized to share ~1 MiB, a typical per-core L2, + * rather than allowing each block that much on its own. * 2. Load balance: enough batches that each thread gets several, i.e. * batch <= npts / (4*nthreads). + * + * Only the batch-proportional blocks are budgeted here. The gathers that + * depend on nbe alone -- D compressed to nbe*nbe, C compressed to nbe*nmo -- + * are unaffected by the batch size, so no choice made here bounds them; on a + * large basis with weak screening the nbe*nbe gather dominates this budget. */ size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads, int nscratch ) { From 986415cd66df77ba3ad8ffae66f18ee4a3c30f84 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:35:27 -0700 Subject: [PATCH 28/31] OrbitalEvaluator: derive the batch decomposition without the thread count 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. --- include/gauxc/orbital_evaluator.hpp | 22 ++++----- src/orbital_evaluator.cxx | 28 +++++------ tests/orbital_evaluator_test.cxx | 73 +++++++++++++++++++---------- 3 files changed, 72 insertions(+), 51 deletions(-) diff --git a/include/gauxc/orbital_evaluator.hpp b/include/gauxc/orbital_evaluator.hpp index ab6fd035..ebecb2a8 100644 --- a/include/gauxc/orbital_evaluator.hpp +++ b/include/gauxc/orbital_evaluator.hpp @@ -39,18 +39,18 @@ namespace detail { * `gauxc/external/cube.hpp`. * * Points are evaluated in batches and each batch is screened against the - * per-shell cutoff radii. The batch size is derived in part from the OpenMP - * thread count, so results are bit-reproducible for a fixed thread count but - * may differ between thread counts. Each neglected shell contributes at most - * the shell tolerance, so that difference is of the order of the tolerance - * rather than bounded by it: it grows with the number of shells dropped, and - * reaches roughly twice the tolerance on benzene in cc-pVDZ. + * per-shell cutoff radii. The batch decomposition is derived from the point + * count and the basis alone, never from the thread count, so results are + * bitwise reproducible from run to run and across thread counts. * - * For the same reason the CubeGrid overloads are not bitwise equal to the - * pointer overloads given the same points. A CubeGrid is walked in spatially - * compact tiles, which screen better, whereas a caller-supplied point array - * is batched in the order it arrives; the two decompositions drop different - * shells and so agree to the shell tolerance rather than bitwise. + * The CubeGrid overloads are not bitwise equal to the pointer overloads given + * the same points, however. A CubeGrid is walked in spatially compact tiles, + * which screen better, whereas a caller-supplied point array is batched in + * the order it arrives, so the two drop different shells. Each neglected + * shell contributes at most the shell tolerance, which makes the difference + * of the order of the tolerance rather than bounded by it: it grows with the + * number of shells dropped, and reaches roughly twice the tolerance on + * benzene in cc-pVDZ. * * Screening cost is governed by the shell tolerance, which the evaluator * applies to its own private copy of the basis, so the caller's basis (and diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 5a701152..3f151745 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -19,12 +19,6 @@ #include #include -#ifdef _OPENMP -#include -#else -inline int omp_get_max_threads() { return 1; } -#endif - #include #include @@ -81,19 +75,27 @@ namespace { * collocation kernel's own transpose staging, and, for density, the * D*AO product. They are sized to share ~1 MiB, a typical per-core L2, * rather than allowing each block that much on its own. - * 2. Load balance: enough batches that each thread gets several, i.e. - * batch <= npts / (4*nthreads). + * 2. Load balance: at least kMinBatches batches, so that every thread of a + * wide machine gets several. + * + * The second bound is a fixed count rather than a multiple of the thread + * count. Batch shape decides which shells screen out, so taking the thread + * count here would make the numbers depend on it. 1024 covers 256 threads at + * four batches each, and the batches it forces on a narrower machine cost + * little: against a run with the bound lifted entirely, benzene/cc-pVDZ at + * 64^3 serial was ~4% slower on orbitals and ~10% faster on density, and + * water was faster on both at 1 and 16 threads. * * Only the batch-proportional blocks are budgeted here. The gathers that * depend on nbe alone -- D compressed to nbe*nbe, C compressed to nbe*nmo -- * are unaffected by the batch size, so no choice made here bounds them; on a * large basis with weak screening the nbe*nbe gather dominates this budget. */ -size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads, - int nscratch ) { +size_t choose_batch_size( int32_t nbf, size_t npts, int nscratch ) { constexpr size_t kTargetScratchBytesPerThread = 1024 * 1024; constexpr size_t kMinBatch = 128; constexpr size_t kMaxBatch = 8192; + constexpr size_t kMinBatches = 1024; size_t batch = kMaxBatch; if( nbf > 0 ) { @@ -101,9 +103,7 @@ size_t choose_batch_size( int32_t nbf, size_t npts, int nthreads, ( sizeof(double) * static_cast(nbf) * static_cast(nscratch) ); } - if( nthreads > 0 ) { - batch = std::min( batch, npts / ( 4 * static_cast(nthreads) ) ); - } + batch = std::min( batch, npts / kMinBatches ); return std::clamp( batch, kMinBatch, kMaxBatch ); } @@ -546,7 +546,7 @@ size_t batch_target( int32_t nbf, size_t npts ) { constexpr int collocation_scratch_blocks = 1; const int nscratch = 1 + collocation_scratch_blocks + Contractor::scratch_blocks; - return choose_batch_size( nbf, npts, omp_get_max_threads(), nscratch ); + return choose_batch_size( nbf, npts, nscratch ); } void check_orbital_args( const std::string& ctx, size_t npts, int32_t nbf, diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index f2e10e42..017360b0 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -842,9 +842,9 @@ TEST_CASE("OrbitalEvaluator screens a tile spanning part of a long axis", TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", "[orbital_evaluator]") { // Regression guard: scratch must follow the thread count in force at - // evaluation time, not at construction. Batch shape is derived from the - // thread count, so the two runs screen slightly differently and agree to - // within the shell tolerance rather than bitwise. + // evaluation time, not at construction. Nothing about the batch + // decomposition depends on the thread count, so the two runs screen + // identically and have to agree bitwise. auto mol = make_water(); auto basis = make_ccpvdz(mol, SphericalType(true)); constexpr double shell_tol = 1e-12; @@ -885,19 +885,22 @@ TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", #endif for (int64_t p = 0; p < npts; ++p) { - CHECK(orb_par[p] == Approx(orb_serial[p]).margin(shell_tol)); - CHECK(rho_par[p] == Approx(rho_serial[p]).margin(shell_tol)); + CHECK(orb_par[p] == orb_serial[p]); + CHECK(rho_par[p] == rho_serial[p]); } } -TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", +TEST_CASE("OrbitalEvaluator results do not depend on the thread count", "[orbital_evaluator]") { - // Batch size is derived from the thread count, so when screening is active - // different thread counts screen against different batch bounding boxes and - // the results are not bit-identical. Each dropped shell contributes at most - // the shell tolerance and the errors accumulate, so the bound scales with - // how many can be dropped; nbf is the generous ceiling on that. Still eight - // orders below the field itself, so it remains a real constraint. + // The batch decomposition is derived from the point count and the basis + // alone, so every thread count screens against the same set of bounding + // boxes and has to return the same bits. Deriving it from the thread count + // instead is the easy mistake, and two things have to hold for that mistake + // to be visible. Screening must be active, hence two centres 20 Bohr apart + // with plenty of empty grid around them. And the batch size must be free to + // move: below kMinBatches * kMinBatch points the minimum-batch-count bound + // rounds away and the batch floor pins the batch whatever the thread count, + // so the grid is deliberately larger than that. Molecule mol; mol.emplace_back(AtomicNumber(8), 0.0, 0.0, 0.0); mol.emplace_back(AtomicNumber(1), 20.0, 0.0, 0.0); @@ -909,34 +912,52 @@ TEST_CASE("OrbitalEvaluator thread-count dependence is bounded by screening", CubeGrid grid; grid.origin = {-4.0, -4.0, -4.0}; - grid.spacing = {0.8, 2.0, 2.0}; - grid.nx = 40; - grid.ny = 4; - grid.nz = 4; - const int64_t npts = grid.num_points(); - + grid.spacing = {0.2, 0.5, 0.5}; + grid.nx = 160; + grid.ny = 32; + grid.nz = 32; + const auto npts = static_cast(grid.num_points()); + REQUIRE(npts > 65536); + + const auto C = make_random_vector(static_cast(nbf), 1618u); std::vector D(static_cast(nbf) * nbf, 0.0); for (int32_t i = 0; i < nbf; ++i) D[i * nbf + i] = 1.0; - std::vector rho_serial(static_cast(npts)); - std::vector rho_par(static_cast(npts)); + std::vector orb_ref(npts), rho_ref(npts); + std::vector orb(npts), rho(npts); #ifdef _OPENMP const int saved_threads = omp_get_max_threads(); omp_set_num_threads(1); #endif - eval.eval_density(grid, D.data(), nbf, rho_serial.data()); + eval.eval_orbital(grid, C.data(), orb_ref.data()); + eval.eval_density(grid, D.data(), nbf, rho_ref.data()); + + // Screening has to be doing something, or the invariance is vacuous. + double max_rho = 0.0; + for (size_t p = 0; p < npts; ++p) max_rho = std::max(max_rho, rho_ref[p]); + REQUIRE(max_rho > 1e-2); + + size_t orb_diff = 0, rho_diff = 0; + for (int nthreads : {2, 3, 5, 8, 13, 16, 32}) { #ifdef _OPENMP - omp_set_num_threads(saved_threads > 1 ? saved_threads : 2); + omp_set_num_threads(nthreads); +#else + (void)nthreads; #endif - eval.eval_density(grid, D.data(), nbf, rho_par.data()); + eval.eval_orbital(grid, C.data(), orb.data()); + eval.eval_density(grid, D.data(), nbf, rho.data()); + for (size_t p = 0; p < npts; ++p) { + if (orb[p] != orb_ref[p]) ++orb_diff; + if (rho[p] != rho_ref[p]) ++rho_diff; + } + } #ifdef _OPENMP omp_set_num_threads(saved_threads); #endif - for (int64_t p = 0; p < npts; ++p) { - CHECK(rho_par[p] == Approx(rho_serial[p]).margin(nbf * shell_tol)); - } + CHECK(orb_diff == 0); + CHECK(rho_diff == 0); } TEST_CASE("OrbitalEvaluator screening error scales with the shell tolerance", From 83ff1cb5710608ad414a81950861298b44a97745 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:36:15 -0700 Subject: [PATCH 29/31] tests: pin that a const evaluator is safe to call concurrently 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. --- tests/orbital_evaluator_test.cxx | 69 ++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 017360b0..605c9707 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -890,6 +891,74 @@ TEST_CASE("OrbitalEvaluator survives a thread-count change after construction", } } +TEST_CASE("OrbitalEvaluator is safe to invoke concurrently through const", + "[orbital_evaluator]") { + // The evaluator holds no mutable state, so a const invocation has to be safe + // both from several threads at once and from inside an existing parallel + // region. Neither mode crashed when scratch was shared between them; they + // silently returned wrong numbers, which is why they are pinned here rather + // than left to the thread-count test to catch. + // + // Nothing here changes the thread count, so every run derives the same batch + // shape and screens identically. Agreement is therefore bitwise: any + // difference at all means state leaked between calls. + auto mol = make_water(); + constexpr double shell_tol = 1e-10; + auto basis = make_ccpvdz(mol, SphericalType(true)); + for (auto& sh : basis) sh.set_shell_tolerance(shell_tol); + const int32_t nbf = basis.nbf(); + auto eval = make_evaluator(basis, shell_tol); + const OrbitalEvaluator& ceval = eval; + + auto grid = CubeGrid::from_molecule(mol, 12, 12, 12); + const auto npts = static_cast(grid.num_points()); + + const auto C = make_random_vector(static_cast(nbf), 271828u); + std::vector D(static_cast(nbf) * nbf, 0.0); + for (int32_t i = 0; i < nbf; ++i) D[static_cast(i) * nbf + i] = 1.0; + + auto run = [&](double* orb, double* rho) { + ceval.eval_orbital(grid, C.data(), orb); + ceval.eval_density(grid, D.data(), nbf, rho); + }; + + std::vector orb_ref(npts), rho_ref(npts); + run(orb_ref.data(), rho_ref.data()); + double max_ref = 0.0; + for (size_t p = 0; p < npts; ++p) max_ref = std::max(max_ref, rho_ref[p]); + REQUIRE(max_ref > 1e-2); + + constexpr int nrunners = 3; + std::vector orb(npts * nrunners), rho(npts * nrunners); + + SECTION("concurrent std::threads on the same evaluator") { + std::vector pool; + for (int t = 0; t < nrunners; ++t) + pool.emplace_back([&, t] { + run(orb.data() + t * npts, rho.data() + t * npts); + }); + for (auto& th : pool) th.join(); + } + +#ifdef _OPENMP + SECTION("evaluation from inside an enclosing parallel region") { +#pragma omp parallel for num_threads(nrunners) schedule(static, 1) + for (int t = 0; t < nrunners; ++t) + run(orb.data() + static_cast(t) * npts, + rho.data() + static_cast(t) * npts); + } +#endif + + size_t orb_diff = 0, rho_diff = 0; + for (size_t t = 0; t < nrunners; ++t) + for (size_t p = 0; p < npts; ++p) { + if (orb[t * npts + p] != orb_ref[p]) ++orb_diff; + if (rho[t * npts + p] != rho_ref[p]) ++rho_diff; + } + CHECK(orb_diff == 0); + CHECK(rho_diff == 0); +} + TEST_CASE("OrbitalEvaluator results do not depend on the thread count", "[orbital_evaluator]") { // The batch decomposition is derived from the point count and the basis From 8384c71ff152b3342951365d6b948614da2eeaec Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Thu, 6 Aug 2026 08:37:01 -0700 Subject: [PATCH 30/31] OrbitalEvaluator: lower the batch floor to the measured optimum 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. --- src/orbital_evaluator.cxx | 40 +++++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/src/orbital_evaluator.cxx b/src/orbital_evaluator.cxx index 3f151745..c67ff5c2 100644 --- a/src/orbital_evaluator.cxx +++ b/src/orbital_evaluator.cxx @@ -69,22 +69,30 @@ namespace { /** @brief Choose the number of points evaluated per batch. * - * Two constraints, whichever is tighter: - * 1. Cache footprint of the blocks that scale with the batch. Each thread - * holds `nscratch` live blocks of nbf*batch doubles: the AO block, the - * collocation kernel's own transpose staging, and, for density, the - * D*AO product. They are sized to share ~1 MiB, a typical per-core L2, - * rather than allowing each block that much on its own. - * 2. Load balance: at least kMinBatches batches, so that every thread of a - * wide machine gets several. + * Three bounds, listed by which one binds as the basis grows. * - * The second bound is a fixed count rather than a multiple of the thread - * count. Batch shape decides which shells screen out, so taking the thread - * count here would make the numbers depend on it. 1024 covers 256 threads at - * four batches each, and the batches it forces on a narrower machine cost - * little: against a run with the bound lifted entirely, benzene/cc-pVDZ at - * 64^3 serial was ~4% slower on orbitals and ~10% faster on density, and - * water was faster on both at 1 and 16 threads. + * 1. Cache footprint, on small and medium bases. Each thread holds `nscratch` + * live blocks of nbf*batch doubles -- the AO block, the collocation + * kernel's own transpose staging, and, for density, the D*AO product -- + * sized to share ~1 MiB, a typical per-core L2, rather than allowing each + * block that much on its own. + * + * 2. A minimum batch count, so every thread of a wide machine gets several. + * Fixed rather than a multiple of the thread count: batch shape decides + * which shells screen out, so taking the thread count here would make the + * numbers depend on it. 1024 covers 256 threads at four batches each and + * costs a few percent either way on a narrower one. + * + * 3. A floor, which is what binds on large bases. The cache bound keeps + * shrinking with nbf -- six points at ubiquitin/cc-pVDZ -- while the + * measured optimum does not: a batch that thin cannot amortise the + * per-batch nbe*nbe gather and leaves the GEMM too skinny to run well. + * 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), with 128 + * costing ~10% and 512 costing 2.5x. Above nbf ~340 this floor, not the + * cache bound, sets the batch, so the ~1 MiB target is not met there. + * It does not defeat bound 2: at 64 points a batch, any grid of 65536 + * points or more already yields the 1024 batches that bound asks for. * * Only the batch-proportional blocks are budgeted here. The gathers that * depend on nbe alone -- D compressed to nbe*nbe, C compressed to nbe*nmo -- @@ -93,7 +101,7 @@ namespace { */ size_t choose_batch_size( int32_t nbf, size_t npts, int nscratch ) { constexpr size_t kTargetScratchBytesPerThread = 1024 * 1024; - constexpr size_t kMinBatch = 128; + constexpr size_t kMinBatch = 64; constexpr size_t kMaxBatch = 8192; constexpr size_t kMinBatches = 1024; From 201b5710a45ce279035b3934b5ade959aefd9405 Mon Sep 17 00:00:00 2001 From: Conrad Johnston Date: Fri, 7 Aug 2026 08:20:16 -0700 Subject: [PATCH 31/31] tests: stop comparing a NaN's sign against the platform snprintf C99 leaves the sign of a printed NaN implementation-defined: glibc emits it, BSD libc does not. The formatter deliberately emits it, matching glibc, so byte-comparing that one field against snprintf passed on Linux and failed the macOS job. The other twenty values keep the byte-exact comparison; the NaN spelling is now pinned directly. Infinities are unaffected, since there the sign is required. --- tests/orbital_evaluator_test.cxx | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/orbital_evaluator_test.cxx b/tests/orbital_evaluator_test.cxx index 605c9707..e0a09383 100644 --- a/tests/orbital_evaluator_test.cxx +++ b/tests/orbital_evaluator_test.cxx @@ -1352,7 +1352,14 @@ TEST_CASE("write_cube agrees with snprintf %13.5E formatting", "[cube]") { std::ostringstream expected; for (size_t i = 0; i < field.size(); ++i) { char buf[32]; - std::snprintf(buf, sizeof(buf), "%13.5E", field[i]); + // C99 leaves the sign of a printed NaN implementation-defined: glibc emits + // it, BSD libc does not, so snprintf is not a portable reference for that + // one value. The formatter deliberately emits it, which is pinned here. + if (std::isnan(field[i])) + std::snprintf(buf, sizeof(buf), "%13s", + std::signbit(field[i]) ? "-NAN" : "NAN"); + else + std::snprintf(buf, sizeof(buf), "%13.5E", field[i]); expected << buf; if ((i + 1) % 6 == 0 || (i + 1) == field.size()) expected << '\n'; }