diff --git a/fortran/tests/tests.f90 b/fortran/tests/tests.f90 index 96cec08c..bd6ec4d4 100644 --- a/fortran/tests/tests.f90 +++ b/fortran/tests/tests.f90 @@ -50,34 +50,39 @@ program vesin_test if (neighbor_list%length /= 10) call print_and_exit("wrong number of pairs") - expected_pairs = reshape([0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1], [2, 10]) + ! Output is now sorted deterministically by (pair[0], pair[1], shift_x, + ! shift_y, shift_z) in GrowableNeighborList::sort. Before that fix only + ! pair[0] was used as the sort key, so the order within a fixed pair[0] + ! (e.g. all (0,0) self-loops with different shifts) was + ! implementation-defined; this test now records the stable order. + expected_pairs = reshape([0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1], [2, 10]) expected_distances = [& - 2.6870059, 2.3021729, 2.3021729, 1.8384775, 3.2000000, & - 3.2000000, 3.2000000, 3.2000000, 3.2000000, 3.2000000 & + 3.2000000, 3.2000000, 3.2000000, 2.6870059, 2.3021729, & + 2.3021729, 1.8384775, 3.2000000, 3.2000000, 3.2000000 & ] expected_shifts = reshape([ & + 0, 0, 1, & + 0, 1, 0, & + 1, 0, 0, & 0, -1, -1, & 0, -1, 0, & 0, 0, -1, & 0, 0, 0, & 0, 0, 1, & 0, 1, 0, & - 1, 0, 0, & - 0, 0, 1, & - 0, 1, 0, & 1, 0, 0 & ], [3, 10]) expected_vectors = reshape([ & + 0.0000000, 0.0000000, 3.2000000, & + 0.0000000, 3.2000000, 0.0000000, & + 3.2000000, 0.0000000, 0.0000000, & 0.0000000, -1.9000000, -1.9000000, & 0.0000000, -1.9000000, 1.2999999, & 0.0000000, 1.2999999, -1.9000000, & 0.0000000, 1.2999999, 1.2999999, & 0.0000000, 0.0000000, 3.2000000, & 0.0000000, 3.2000000, 0.0000000, & - 3.2000000, 0.0000000, 0.0000000, & - 0.0000000, 0.0000000, 3.2000000, & - 0.0000000, 3.2000000, 0.0000000, & 3.2000000, 0.0000000, 0.0000000 & ], [3, 10]) diff --git a/python/vesin/tests/test_verlet.py b/python/vesin/tests/test_verlet.py index ae8cea79..345c4860 100644 --- a/python/vesin/tests/test_verlet.py +++ b/python/vesin/tests/test_verlet.py @@ -1,6 +1,9 @@ """Tests for Verlet caching via NeighborList with skin > 0.""" import math +import subprocess +import sys +import textwrap import numpy as np import pytest @@ -220,3 +223,48 @@ def test_non_periodic(): ) assert verlet == ref + + +def test_auto_verlet_many_candidate_recompute_runs_in_subprocess(tmp_path): + script = textwrap.dedent( + """ + import numpy as np + from vesin import NeighborList + + n_atoms = 1024 + density = 0.05 + box_size = (n_atoms / density) ** (1.0 / 3.0) + rng = np.random.default_rng(20260512 + n_atoms) + positions = np.ascontiguousarray( + rng.random((n_atoms, 3), dtype=np.float64) * box_size + ) + box = np.eye(3, dtype=np.float64) * box_size + + nl = NeighborList( + cutoff=5.0, + full_list=True, + sorted=False, + skin=1.0, + algorithm="auto", + ) + first, second = nl.compute( + positions, + box, + periodic=True, + quantities="ij", + copy=False, + ) + if len(first) == 0 or len(second) == 0: + raise SystemExit("expected non-empty neighbor output") + """ + ) + + completed = subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + cwd=tmp_path, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + completed.stdout diff --git a/vesin/CMakeLists.txt b/vesin/CMakeLists.txt index 154b0f2c..1ea8c282 100644 --- a/vesin/CMakeLists.txt +++ b/vesin/CMakeLists.txt @@ -41,6 +41,7 @@ option(VESIN_ENABLE_NVTX "Enable NVTX profiling markers" OFF) set(VESIN_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/src/vesin.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/cpu_cell_list.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/cluster_pair_search.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/vesin_cuda.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/verlet.cpp ) @@ -97,8 +98,35 @@ endif() FetchContent_MakeAvailable(gpulite) -target_link_libraries(vesin_objects PRIVATE gpulite) -target_link_libraries(vesin PRIVATE gpulite) +# Google Highway for portable SIMD (used in cluster-pair search) +FetchContent_Declare( + highway + GIT_REPOSITORY https://github.com/google/highway.git + GIT_TAG 1.2.0 + GIT_SHALLOW TRUE + EXCLUDE_FROM_ALL +) +set(HWY_ENABLE_TESTS OFF CACHE BOOL "" FORCE) +set(HWY_ENABLE_EXAMPLES OFF CACHE BOOL "" FORCE) +set(HWY_ENABLE_CONTRIB OFF CACHE BOOL "" FORCE) +set(BUILD_TESTING OFF CACHE BOOL "" FORCE) +# Force Highway to compile as a static library so libvesin.so is self-contained. +# Without this libvesin links libhwy.so.1 dynamically, which then has to be +# bundled by delocate/auditwheel when packaging Python wheels (it can't find +# the library and the build fails). Restore the saved value afterwards so +# vesin itself still honors the caller's BUILD_SHARED_LIBS choice. +set(_VESIN_SAVED_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS}) +set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE) +FetchContent_MakeAvailable(highway) +set(BUILD_SHARED_LIBS ${_VESIN_SAVED_BUILD_SHARED_LIBS} CACHE BOOL "" FORCE) + +target_link_libraries(vesin_objects PRIVATE gpulite hwy) +target_link_libraries(vesin PRIVATE gpulite hwy) +# Mark the SIMD code paths as available so the source can #ifdef around the +# Highway-using kernels and provide scalar fallbacks for builds without it +# (the single-file dist build does not have Highway on its include path). +target_compile_definitions(vesin_objects PRIVATE VESIN_HAVE_HIGHWAY) +target_compile_definitions(vesin PRIVATE VESIN_HAVE_HIGHWAY) # Create generated directory file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/generated) diff --git a/vesin/scripts/create-single-cpp.py b/vesin/scripts/create-single-cpp.py index 48043a06..b440657d 100755 --- a/vesin/scripts/create-single-cpp.py +++ b/vesin/scripts/create-single-cpp.py @@ -96,6 +96,7 @@ def add_version(output): with open(os.path.join(DIST, "vesin-single-build.cpp"), "w") as output: add_version(output) merge_files("cpu_cell_list.cpp", output) + merge_files("cluster_pair_search.cpp", output) merge_files("verlet.cpp", output) merge_files("vesin_cuda.cpp", output) merge_files("vesin.cpp", output) @@ -104,6 +105,7 @@ def add_version(output): with open(os.path.join(DIST, "vesin-single-build-nocuda.cpp"), "w") as output: add_version(output) merge_files("cpu_cell_list.cpp", output) + merge_files("cluster_pair_search.cpp", output) merge_files("verlet.cpp", output) merge_files("vesin_cuda_stub.cpp", output) merge_files("vesin.cpp", output) diff --git a/vesin/src/cluster.hpp b/vesin/src/cluster.hpp new file mode 100644 index 00000000..99e19681 --- /dev/null +++ b/vesin/src/cluster.hpp @@ -0,0 +1,175 @@ +#ifndef VESIN_CLUSTER_HPP +#define VESIN_CLUSTER_HPP + +#include +#include +#include +#include +#include + +#include "types.hpp" +#include "vesin.h" + +namespace vesin { + +/// Size of a cluster on CPU. 8 atoms maps well to AVX2 (4 doubles) with +/// a 2-iteration inner loop, and degrades gracefully to SSE (2 doubles, +/// 4 iterations) or AVX-512 (8 doubles, 1 iteration). +static constexpr int32_t CLUSTER_SIZE_CPU = 8; + +/// Minimum atom count where CPU auto-dispatch uses cluster-pair search. +static constexpr size_t CLUSTER_PAIR_THRESHOLD = 256; + +/// A cluster of up to CLUSTER_SIZE_CPU atoms with a bounding box and +/// SoA position data for SIMD distance calculations. +struct Cluster { + int32_t atom_indices[CLUSTER_SIZE_CPU]; + int32_t n_atoms; // actual count (<= CLUSTER_SIZE_CPU) + float bb_lower[3]; // bounding box min (float for SIMD efficiency) + float bb_upper[3]; // bounding box max + + // SoA (Structure of Arrays) wrapped positions for SIMD loads. + // These store the atom positions after subtracting the wrap shift, + // matching the coordinate space used for BB tests. + alignas(64) double pos_x[CLUSTER_SIZE_CPU]; + alignas(64) double pos_y[CLUSTER_SIZE_CPU]; + alignas(64) double pos_z[CLUSTER_SIZE_CPU]; +}; + +/// Grid of clusters organized in 3D cells. +struct ClusterGrid { + std::vector clusters; + + // Grid dimensions (number of cells in each direction) + std::array n_cells; + + // Which clusters belong to which cell: cell_offsets[cell_idx] to + // cell_offsets[cell_idx+1] gives the range of cluster indices in + // the clusters array. + std::vector cell_offsets; // [n_cells_total + 1], CSR-style + + // Per-atom wrap shift: when an atom's fractional coordinate falls + // outside [0, n_cells), it is wrapped into the grid and the integer + // shift is recorded here. Indexed by original atom index. + std::vector atom_wrap_shifts; + + // Precomputed wrapped positions for all atoms: points[i] minus + // wrap_shift[i].cartesian(box). Indexed by original atom index. + // Used to avoid per-pair matrix multiply in the inner loop. + std::vector wrapped_positions; +}; + +/// A cluster-pair candidate retained for exact cutoff filtering. +struct ClusterPairCandidate { + int32_t first_cluster; + int32_t second_cluster; + CellShift cell_shift; + Vector shift_cartesian; +}; + +/// Build a cluster grid from atom positions. +/// +/// Algorithm: +/// 1. Compute grid cell dimensions from box vectors and cutoff +/// 2. Assign atoms to grid cells (fractional coordinate binning) +/// 3. Within each cell, sort atoms by z coordinate, group into clusters +/// 4. Compute cluster bounding boxes (AABB) +/// 5. Fill SoA position arrays for SIMD +ClusterGrid build_cluster_grid( + const Vector* points, + size_t n_points, + const BoundingBox& box, + double cutoff +); + +/// Build cluster-pair candidates passing the bounding-box cutoff test. +std::vector build_cluster_pair_candidates( + const ClusterGrid& grid, + const BoundingBox& cell, + double cutoff +); + +/// Minimum squared distance between two AABBs. +/// Returns 0 if the boxes overlap. +inline float bb_distance_sq(const Cluster& a, const Cluster& b) { + float dist_sq = 0.0f; + for (int d = 0; d < 3; d++) { + float gap = 0.0f; + if (a.bb_lower[d] > b.bb_upper[d]) { + gap = a.bb_lower[d] - b.bb_upper[d]; + } else if (b.bb_lower[d] > a.bb_upper[d]) { + gap = b.bb_lower[d] - a.bb_upper[d]; + } + dist_sq += gap * gap; + } + return dist_sq; +} + +/// Minimum squared distance between two AABBs where cluster_b is shifted +/// by a Cartesian offset (for periodic images). When shift is zero, this +/// is equivalent to bb_distance_sq. +inline float bb_distance_sq_shifted( + const Cluster& a, const Cluster& b, const float shift[3] +) { + float dist_sq = 0.0f; + for (int d = 0; d < 3; d++) { + float b_lo = b.bb_lower[d] + shift[d]; + float b_hi = b.bb_upper[d] + shift[d]; + float gap = 0.0f; + if (a.bb_lower[d] > b_hi) { + gap = a.bb_lower[d] - b_hi; + } else if (b_lo > a.bb_upper[d]) { + gap = b_lo - a.bb_upper[d]; + } + dist_sq += gap * gap; + } + return dist_sq; +} + +namespace cpu { + +/// Cluster-pair neighbor search with SIMD distance calculations. +/// Replaces cell_list for N >= CLUSTER_PAIR_THRESHOLD. +/// +/// Output format is identical to the cell-list path: per-atom pairs with +/// optional shifts, distances, and vectors in VesinNeighborList. +void cluster_pair_neighbors( + const Vector* points, + size_t n_points, + const BoundingBox& cell, + VesinOptions options, + VesinNeighborList& neighbors +); + +/// Filter cached cluster-pair candidates at the exact cutoff. +void filter_cluster_pair_candidates( + const Vector* points, + const BoundingBox& cell, + const ClusterGrid& grid, + const std::vector& candidates, + double cutoff, + VesinOptions options, + VesinNeighborList& raw_neighbors, + size_t initial_capacity, + size_t& output_capacity +); + +} // namespace cpu + +/// True iff the cell shift is all-zero (atom in its home cell). +inline bool is_zero_shift(CellShift shift) { + return shift[0] == 0 && shift[1] == 0 && shift[2] == 0; +} + +/// Convert a cell shift to a Cartesian translation vector via the bounding box. +/// Zero-shift fast path returns the zero vector without touching the box. +inline Vector shift_cartesian(CellShift shift, const BoundingBox& cell) { + if (is_zero_shift(shift)) { + return Vector{0.0, 0.0, 0.0}; + } + return shift.cartesian(cell); +} + +} // namespace vesin + +#endif diff --git a/vesin/src/cluster_pair_search.cpp b/vesin/src/cluster_pair_search.cpp new file mode 100644 index 00000000..6eb212ae --- /dev/null +++ b/vesin/src/cluster_pair_search.cpp @@ -0,0 +1,730 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef VESIN_HAVE_HIGHWAY +#include +#else +#define HWY_ATTR +#define HWY_RESTRICT +#endif + +#include "cluster.hpp" +#include "cpu_cell_list.hpp" + +using namespace vesin; + +/// Maximal number of cells (same as cell list) +#define MAX_NUMBER_OF_CELLS 1e5 + +/// divmod with Python semantics (positive remainder) +static std::tuple divmod(int32_t a, int32_t b) { + auto quotient = a / b; + auto remainder = a % b; + if (remainder < 0) { + remainder += b; + quotient -= 1; + } + return std::make_tuple(quotient, remainder); +} + +ClusterGrid vesin::build_cluster_grid( + const Vector* points, + size_t n_points, + const BoundingBox& box, + double cutoff +) { + ClusterGrid grid; + + auto distances_between_faces = box.distances_between_faces(); + + // Compute grid cell dimensions + auto n_cells_f = Vector{ + std::clamp(std::trunc(distances_between_faces[0] / cutoff), 1.0, HUGE_VAL), + std::clamp(std::trunc(distances_between_faces[1] / cutoff), 1.0, HUGE_VAL), + std::clamp(std::trunc(distances_between_faces[2] / cutoff), 1.0, HUGE_VAL), + }; + + // Limit memory (same as cell list) + auto n_cells_total = n_cells_f[0] * n_cells_f[1] * n_cells_f[2]; + if (n_cells_total > MAX_NUMBER_OF_CELLS) { + auto ratio_x_y = n_cells_f[0] / n_cells_f[1]; + auto ratio_y_z = n_cells_f[1] / n_cells_f[2]; + n_cells_f[2] = std::trunc(std::cbrt(MAX_NUMBER_OF_CELLS / (ratio_x_y * ratio_y_z * ratio_y_z))); + n_cells_f[1] = std::trunc(ratio_y_z * n_cells_f[2]); + n_cells_f[0] = std::trunc(ratio_x_y * n_cells_f[1]); + } + + grid.n_cells = { + static_cast(n_cells_f[0]), + static_cast(n_cells_f[1]), + static_cast(n_cells_f[2]), + }; + + // Clamp to at least 1 + for (int d = 0; d < 3; d++) { + if (grid.n_cells[d] < 1) { + grid.n_cells[d] = 1; + } + } + + int32_t total_cells = grid.n_cells[0] * grid.n_cells[1] * grid.n_cells[2]; + + // Assign atoms to cells + struct AtomCell { + size_t atom_index; + int32_t cell_linear; + CellShift wrap_shift; + float z_frac; // fractional z for sorting within cell + }; + + grid.atom_wrap_shifts.resize(n_points); + + std::vector assignments(n_points); + for (size_t i = 0; i < n_points; i++) { + auto fractional = box.cartesian_to_fractional(points[i]); + + auto cell_idx = std::array{ + static_cast(std::floor(fractional[0] * static_cast(grid.n_cells[0]))), + static_cast(std::floor(fractional[1] * static_cast(grid.n_cells[1]))), + static_cast(std::floor(fractional[2] * static_cast(grid.n_cells[2]))), + }; + + CellShift shift{}; + for (int d = 0; d < 3; d++) { + if (box.periodic(d)) { + auto [q, r] = divmod(cell_idx[d], grid.n_cells[d]); + shift[d] = q; + cell_idx[d] = r; + } else { + shift[d] = 0; + cell_idx[d] = std::clamp(cell_idx[d], 0, grid.n_cells[d] - 1); + } + } + + grid.atom_wrap_shifts[i] = shift; + + int32_t linear = (grid.n_cells[0] * grid.n_cells[1] * cell_idx[2]) + (grid.n_cells[0] * cell_idx[1]) + cell_idx[0]; + + assignments[i] = {i, linear, shift, static_cast(fractional[2])}; + } + + // Precompute wrapped positions for all atoms (avoids per-pair + // matrix multiply in the inner loop). + grid.wrapped_positions.resize(n_points); + for (size_t i = 0; i < n_points; i++) { + grid.wrapped_positions[i] = points[i] - grid.atom_wrap_shifts[i].cartesian(box); + } + + // Count atoms per cell + std::vector cell_counts(total_cells, 0); + for (auto& a : assignments) { + cell_counts[a.cell_linear]++; + } + + // Sort assignments by cell, then by z within each cell + std::sort(assignments.begin(), assignments.end(), [](const AtomCell& a, const AtomCell& b) { + if (a.cell_linear != b.cell_linear) { + return a.cell_linear < b.cell_linear; + } + return a.z_frac < b.z_frac; + }); + + // Build clusters: group atoms within each cell into groups of CLUSTER_SIZE_CPU + grid.cell_offsets.resize(total_cells + 1, 0); + grid.clusters.clear(); + + size_t atom_cursor = 0; + for (int32_t cell = 0; cell < total_cells; cell++) { + grid.cell_offsets[cell] = static_cast(grid.clusters.size()); + + int32_t count = cell_counts[cell]; + size_t cell_start = atom_cursor; + + // Group into clusters + for (int32_t offset = 0; offset < count; offset += CLUSTER_SIZE_CPU) { + Cluster cl{}; + cl.n_atoms = std::min(CLUSTER_SIZE_CPU, count - offset); + + // Initialize BB to inverted extremes + for (int d = 0; d < 3; d++) { + cl.bb_lower[d] = std::numeric_limits::max(); + cl.bb_upper[d] = -std::numeric_limits::max(); + } + + // Initialize SoA arrays to zero (padding slots get large + // distance, preventing false matches) + std::memset(cl.pos_x, 0, sizeof(cl.pos_x)); + std::memset(cl.pos_y, 0, sizeof(cl.pos_y)); + std::memset(cl.pos_z, 0, sizeof(cl.pos_z)); + + for (int32_t k = 0; k < cl.n_atoms; k++) { + size_t idx = cell_start + offset + k; + auto atom_idx = assignments[idx].atom_index; + cl.atom_indices[k] = static_cast(atom_idx); + + // Use precomputed wrapped position + const auto& wrapped = grid.wrapped_positions[atom_idx]; + cl.pos_x[k] = wrapped[0]; + cl.pos_y[k] = wrapped[1]; + cl.pos_z[k] = wrapped[2]; + + for (int d = 0; d < 3; d++) { + float p = static_cast(wrapped[d]); + cl.bb_lower[d] = std::min(cl.bb_lower[d], p); + cl.bb_upper[d] = std::max(cl.bb_upper[d], p); + } + } + + // Pad unused slots: set positions to huge value so distance + // check always fails, and indices to -1. + for (int32_t k = cl.n_atoms; k < CLUSTER_SIZE_CPU; k++) { + cl.atom_indices[k] = -1; + cl.pos_x[k] = 1e30; + cl.pos_y[k] = 1e30; + cl.pos_z[k] = 1e30; + } + + grid.clusters.push_back(cl); + } + + atom_cursor += count; + } + grid.cell_offsets[total_cells] = static_cast(grid.clusters.size()); + + return grid; +} + +// --------------------------------------------------------------------------- +// SIMD inner loop using Google Highway (with scalar fallback for builds without +// Highway on the include path, e.g. the single-file dist build). +// --------------------------------------------------------------------------- + +namespace { +#ifdef VESIN_HAVE_HIGHWAY +namespace hn = hwy::HWY_NAMESPACE; +#endif + +/// Process one atom i against all atoms in cluster j using SIMD. +/// Returns the number of pairs found (written into the output arrays). +/// +/// The output arrays must have room for CLUSTER_SIZE_CPU entries. +/// This function writes (idx_j, distance2, vector) for each hit. +HWY_ATTR +static int simd_check_distances( + double i_x, double i_y, double i_z, double shift_x, double shift_y, double shift_z, const double* HWY_RESTRICT j_x, const double* HWY_RESTRICT j_y, const double* HWY_RESTRICT j_z, double cutoff2, + // output arrays (caller provides space for CLUSTER_SIZE_CPU) + double* HWY_RESTRICT out_dist2, + double* HWY_RESTRICT out_dx, + double* HWY_RESTRICT out_dy, + double* HWY_RESTRICT out_dz, + uint8_t* HWY_RESTRICT out_mask +) { +#ifdef VESIN_HAVE_HIGHWAY + const hn::ScalableTag d; + const size_t N = hn::Lanes(d); + + // Broadcast i position (already includes shift subtraction) + const auto vi_x = hn::Set(d, i_x - shift_x); + const auto vi_y = hn::Set(d, i_y - shift_y); + const auto vi_z = hn::Set(d, i_z - shift_z); + const auto vcut2 = hn::Set(d, cutoff2); + + int count = 0; + + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane += N) { + // Load j positions (contiguous, aligned) + auto vj_x = hn::Load(d, j_x + lane); + auto vj_y = hn::Load(d, j_y + lane); + auto vj_z = hn::Load(d, j_z + lane); + + // vector = j - (i - shift) = j - i + shift + auto dx = hn::Sub(vj_x, vi_x); + auto dy = hn::Sub(vj_y, vi_y); + auto dz = hn::Sub(vj_z, vi_z); + + // dist2 = dx*dx + dy*dy + dz*dz + auto dist2 = hn::MulAdd(dx, dx, hn::MulAdd(dy, dy, hn::Mul(dz, dz))); + + // mask: dist2 < cutoff2 + auto mask = hn::Lt(dist2, vcut2); + + // Store results for all lanes, caller filters by mask + hn::Store(dist2, d, out_dist2 + lane); + hn::Store(dx, d, out_dx + lane); + hn::Store(dy, d, out_dy + lane); + hn::Store(dz, d, out_dz + lane); + + // Store mask as bits + uint8_t bits_buf[8] = {}; + hn::StoreMaskBits(d, mask, bits_buf); + uint8_t bits = bits_buf[0]; + for (size_t k = 0; k < N && (lane + k) < CLUSTER_SIZE_CPU; k++) { + out_mask[lane + k] = (bits >> k) & 1; + count += out_mask[lane + k]; + } + } + return count; +#else + // Scalar fallback used by the single-file dist build (no Highway on + // include path). Same outputs as the Highway path; per-pair distance + // compute with no lane fan-out. + const double ix = i_x - shift_x; + const double iy = i_y - shift_y; + const double iz = i_z - shift_z; + int count = 0; + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane++) { + const double dx = j_x[lane] - ix; + const double dy = j_y[lane] - iy; + const double dz = j_z[lane] - iz; + const double dist2 = dx * dx + dy * dy + dz * dz; + out_dist2[lane] = dist2; + out_dx[lane] = dx; + out_dy[lane] = dy; + out_dz[lane] = dz; + const uint8_t hit = (dist2 < cutoff2) ? 1 : 0; + out_mask[lane] = hit; + count += hit; + } + return count; +#endif +} + +} // anonymous namespace + +std::vector vesin::build_cluster_pair_candidates( + const ClusterGrid& grid, + const BoundingBox& cell, + double cutoff +) { + auto cutoff2_f = static_cast(cutoff * cutoff); + auto distances_between_faces = cell.distances_between_faces(); + + auto n_search = std::array{ + static_cast(std::ceil(cutoff * grid.n_cells[0] / distances_between_faces[0])), + static_cast(std::ceil(cutoff * grid.n_cells[1] / distances_between_faces[1])), + static_cast(std::ceil(cutoff * grid.n_cells[2] / distances_between_faces[2])), + }; + + for (int d = 0; d < 3; d++) { + if (n_search[d] < 1) { + n_search[d] = 1; + } + if (grid.n_cells[d] == 1 && !cell.periodic(d)) { + n_search[d] = 0; + } + } + + auto candidates = std::vector(); + candidates.reserve(grid.clusters.size() * 27); + + for (int32_t cz = 0; cz < grid.n_cells[2]; cz++) { + for (int32_t cy = 0; cy < grid.n_cells[1]; cy++) { + for (int32_t cx = 0; cx < grid.n_cells[0]; cx++) { + int32_t cell_i_linear = (grid.n_cells[0] * grid.n_cells[1] * cz) + (grid.n_cells[0] * cy) + cx; + + int32_t ci_start = grid.cell_offsets[cell_i_linear]; + int32_t ci_end = grid.cell_offsets[cell_i_linear + 1]; + + for (int32_t dz = -n_search[2]; dz <= n_search[2]; dz++) { + for (int32_t dy = -n_search[1]; dy <= n_search[1]; dy++) { + for (int32_t dx = -n_search[0]; dx <= n_search[0]; dx++) { + int32_t nx = cx + dx; + int32_t ny = cy + dy; + int32_t nz = cz + dz; + + auto [sx, rx] = divmod(nx, grid.n_cells[0]); + auto [sy, ry] = divmod(ny, grid.n_cells[1]); + auto [sz, rz] = divmod(nz, grid.n_cells[2]); + + if ((sx != 0 && !cell.periodic(0)) || + (sy != 0 && !cell.periodic(1)) || + (sz != 0 && !cell.periodic(2))) { + continue; + } + + int32_t cell_j_linear = (grid.n_cells[0] * grid.n_cells[1] * rz) + (grid.n_cells[0] * ry) + rx; + + int32_t cj_start = grid.cell_offsets[cell_j_linear]; + int32_t cj_end = grid.cell_offsets[cell_j_linear + 1]; + + auto cell_shift_base = CellShift{{sx, sy, sz}}; + auto shift_cart = shift_cartesian(cell_shift_base, cell); + float shift_f[3] = { + static_cast(shift_cart[0]), + static_cast(shift_cart[1]), + static_cast(shift_cart[2]), + }; + + for (int32_t ci = ci_start; ci < ci_end; ci++) { + const auto& cluster_i = grid.clusters[ci]; + for (int32_t cj = cj_start; cj < cj_end; cj++) { + const auto& cluster_j = grid.clusters[cj]; + + float bb_dist = bb_distance_sq_shifted(cluster_i, cluster_j, shift_f); + if (bb_dist > cutoff2_f) { + continue; + } + + candidates.push_back(ClusterPairCandidate{ + ci, + cj, + cell_shift_base, + shift_cart, + }); + } + } + } + } + } + } + } + } + + return candidates; +} + +void vesin::cpu::filter_cluster_pair_candidates( + const Vector* points, + const BoundingBox& cell, + const ClusterGrid& grid, + const std::vector& candidates, + double cutoff, + VesinOptions options, + VesinNeighborList& raw_neighbors, + size_t initial_capacity, + size_t& output_capacity +) { + auto neighbors = GrowableNeighborList{raw_neighbors, initial_capacity, options}; + neighbors.reset(); + + auto current_clusters = grid.clusters; + for (auto& cluster : current_clusters) { + for (size_t d = 0; d < 3; d++) { + cluster.bb_lower[d] = std::numeric_limits::max(); + cluster.bb_upper[d] = std::numeric_limits::lowest(); + } + + for (int32_t atom = 0; atom < cluster.n_atoms; atom++) { + auto atom_index = static_cast(cluster.atom_indices[atom]); + auto wrap_shift = shift_cartesian(grid.atom_wrap_shifts[atom_index], cell); + auto wrapped = points[atom_index] - wrap_shift; + cluster.pos_x[atom] = wrapped[0]; + cluster.pos_y[atom] = wrapped[1]; + cluster.pos_z[atom] = wrapped[2]; + + auto x = static_cast(wrapped[0]); + auto y = static_cast(wrapped[1]); + auto z = static_cast(wrapped[2]); + cluster.bb_lower[0] = std::min(cluster.bb_lower[0], x); + cluster.bb_lower[1] = std::min(cluster.bb_lower[1], y); + cluster.bb_lower[2] = std::min(cluster.bb_lower[2], z); + cluster.bb_upper[0] = std::max(cluster.bb_upper[0], x); + cluster.bb_upper[1] = std::max(cluster.bb_upper[1], y); + cluster.bb_upper[2] = std::max(cluster.bb_upper[2], z); + } + } + + auto cutoff2 = cutoff * cutoff; + auto cutoff2_f = static_cast(cutoff2); + + alignas(64) double tmp_dist2[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dx[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dy[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dz[CLUSTER_SIZE_CPU]; + uint8_t tmp_mask[CLUSTER_SIZE_CPU]; + + for (const auto& candidate : candidates) { + const auto& cluster_i = current_clusters[candidate.first_cluster]; + const auto& cluster_j = current_clusters[candidate.second_cluster]; + float shift_f[3] = { + static_cast(candidate.shift_cartesian[0]), + static_cast(candidate.shift_cartesian[1]), + static_cast(candidate.shift_cartesian[2]), + }; + + if (bb_distance_sq_shifted(cluster_i, cluster_j, shift_f) > cutoff2_f) { + continue; + } + + for (int32_t ai = 0; ai < cluster_i.n_atoms; ai++) { + int32_t idx_i = cluster_i.atom_indices[ai]; + + simd_check_distances( + cluster_i.pos_x[ai], + cluster_i.pos_y[ai], + cluster_i.pos_z[ai], + candidate.shift_cartesian[0], + candidate.shift_cartesian[1], + candidate.shift_cartesian[2], + cluster_j.pos_x, + cluster_j.pos_y, + cluster_j.pos_z, + cutoff2, + tmp_dist2, + tmp_dx, + tmp_dy, + tmp_dz, + tmp_mask + ); + + for (int32_t aj = 0; aj < cluster_j.n_atoms; aj++) { + if (!tmp_mask[aj]) { + continue; + } + + int32_t idx_j = cluster_j.atom_indices[aj]; + auto shift = candidate.cell_shift + grid.atom_wrap_shifts[idx_i] - grid.atom_wrap_shifts[idx_j]; + + if (idx_i == idx_j && is_zero_shift(shift)) { + continue; + } + + if (!options.full) { + if (static_cast(idx_i) > static_cast(idx_j)) { + continue; + } + if (idx_i == idx_j) { + if (shift[0] + shift[1] + shift[2] < 0) { + continue; + } + if ((shift[0] + shift[1] + shift[2] == 0) && + (shift[2] < 0 || (shift[2] == 0 && shift[1] < 0))) { + continue; + } + } + } + + auto index = neighbors.length(); + neighbors.set_pair(index, static_cast(idx_i), static_cast(idx_j)); + + if (options.return_shifts) { + neighbors.set_shift(index, shift); + } + if (options.return_distances) { + neighbors.set_distance(index, std::sqrt(tmp_dist2[aj])); + } + if (options.return_vectors) { + auto vector = Vector{tmp_dx[aj], tmp_dy[aj], tmp_dz[aj]}; + neighbors.set_vector(index, vector); + } + neighbors.increment_length(); + } + } + } + + if (options.sorted) { + neighbors.sort(); + } + + output_capacity = neighbors.capacity; +} + +void vesin::cpu::cluster_pair_neighbors( + const Vector* points, + size_t n_points, + const BoundingBox& cell, + VesinOptions options, + VesinNeighborList& raw_neighbors +) { + auto grid = build_cluster_grid(points, n_points, cell, options.cutoff); + + auto cutoff2 = options.cutoff * options.cutoff; + float cutoff2_f = static_cast(cutoff2); + + auto neighbors = GrowableNeighborList{raw_neighbors, raw_neighbors.length, options}; + neighbors.reset(); + + auto distances_between_faces = cell.distances_between_faces(); + + // Number of cells to search in each direction + auto n_search = std::array{ + static_cast(std::ceil(options.cutoff * grid.n_cells[0] / distances_between_faces[0])), + static_cast(std::ceil(options.cutoff * grid.n_cells[1] / distances_between_faces[1])), + static_cast(std::ceil(options.cutoff * grid.n_cells[2] / distances_between_faces[2])), + }; + + for (int d = 0; d < 3; d++) { + if (n_search[d] < 1) { + n_search[d] = 1; + } + if (grid.n_cells[d] == 1 && !cell.periodic(d)) { + n_search[d] = 0; + } + } + + // Scratch arrays for SIMD output (stack-allocated, reused per atom i) + alignas(64) double tmp_dist2[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dx[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dy[CLUSTER_SIZE_CPU]; + alignas(64) double tmp_dz[CLUSTER_SIZE_CPU]; + uint8_t tmp_mask[CLUSTER_SIZE_CPU]; + + // Iterate over all cells + for (int32_t cz = 0; cz < grid.n_cells[2]; cz++) { + for (int32_t cy = 0; cy < grid.n_cells[1]; cy++) { + for (int32_t cx = 0; cx < grid.n_cells[0]; cx++) { + + int32_t cell_i_linear = (grid.n_cells[0] * grid.n_cells[1] * cz) + (grid.n_cells[0] * cy) + cx; + + int32_t ci_start = grid.cell_offsets[cell_i_linear]; + int32_t ci_end = grid.cell_offsets[cell_i_linear + 1]; + + // Search neighboring cells + for (int32_t dz = -n_search[2]; dz <= n_search[2]; dz++) { + for (int32_t dy = -n_search[1]; dy <= n_search[1]; dy++) { + for (int32_t dx = -n_search[0]; dx <= n_search[0]; dx++) { + + int32_t nx = cx + dx, ny = cy + dy, nz = cz + dz; + + // Wrap neighbor cell and compute cell shift + auto [sx, rx] = divmod(nx, grid.n_cells[0]); + auto [sy, ry] = divmod(ny, grid.n_cells[1]); + auto [sz, rz] = divmod(nz, grid.n_cells[2]); + + // Skip non-periodic wrapping + if ((sx != 0 && !cell.periodic(0)) || + (sy != 0 && !cell.periodic(1)) || + (sz != 0 && !cell.periodic(2))) { + continue; + } + + int32_t cell_j_linear = (grid.n_cells[0] * grid.n_cells[1] * rz) + (grid.n_cells[0] * ry) + rx; + + int32_t cj_start = grid.cell_offsets[cell_j_linear]; + int32_t cj_end = grid.cell_offsets[cell_j_linear + 1]; + + auto cell_shift_base = CellShift{{sx, sy, sz}}; + + // Precompute Cartesian shift for the cell pair. This is + // used both for the BB test (float) and for the SIMD + // distance calculation (double). Since wrapped positions + // already have wrap_shift removed, the cell_shift_base + // Cartesian offset is the only shift needed for the + // vector calculation. + auto shift_cart = cell_shift_base.cartesian(cell); + float shift_f[3] = { + static_cast(shift_cart[0]), + static_cast(shift_cart[1]), + static_cast(shift_cart[2]), + }; + + // Iterate over cluster pairs between these two cells + for (int32_t ci = ci_start; ci < ci_end; ci++) { + const auto& cluster_i = grid.clusters[ci]; + for (int32_t cj = cj_start; cj < cj_end; cj++) { + const auto& cluster_j = grid.clusters[cj]; + + // BB distance test with shift + float bb_dist = bb_distance_sq_shifted( + cluster_i, cluster_j, shift_f + ); + if (bb_dist > cutoff2_f) { + continue; + } + + // SIMD atom-pair expansion: for each atom i, + // check all atoms j in cluster_j via SIMD. + for (int32_t ai = 0; ai < cluster_i.n_atoms; ai++) { + int32_t idx_i = cluster_i.atom_indices[ai]; + + // Use wrapped positions: the vector between + // wrapped[j] and (wrapped[i] - shift_cart) + // gives the correct displacement. + simd_check_distances( + cluster_i.pos_x[ai], + cluster_i.pos_y[ai], + cluster_i.pos_z[ai], + shift_cart[0], + shift_cart[1], + shift_cart[2], + cluster_j.pos_x, + cluster_j.pos_y, + cluster_j.pos_z, + cutoff2, + tmp_dist2, + tmp_dx, + tmp_dy, + tmp_dz, + tmp_mask + ); + + // Process hits from the SIMD pass + for (int32_t aj = 0; aj < cluster_j.n_atoms; aj++) { + if (!tmp_mask[aj]) { + continue; + } + + int32_t idx_j = cluster_j.atom_indices[aj]; + + // Compute per-atom shift incorporating + // wrap corrections (same convention as + // cell-list: shift = cell_shift + wrap_i + // - wrap_j). + auto shift = cell_shift_base + grid.atom_wrap_shifts[idx_i] - grid.atom_wrap_shifts[idx_j]; + bool shift_is_zero = shift[0] == 0 && shift[1] == 0 && shift[2] == 0; + + if (idx_i == idx_j && shift_is_zero) { + continue; + } + + if (!options.full) { + if (static_cast(idx_i) > static_cast(idx_j)) { + continue; + } + if (idx_i == idx_j) { + if (shift[0] + shift[1] + shift[2] < 0) { + continue; + } + if ((shift[0] + shift[1] + shift[2] == 0) && + (shift[2] < 0 || (shift[2] == 0 && shift[1] < 0))) { + continue; + } + } + } + + // The SIMD pass already computed the + // vector and distance2 using wrapped + // positions + cell shift. These are + // numerically identical to + // points[j] - points[i] + shift.cartesian(cell) + // because wrapped[k] = points[k] - wrap[k].cart(M) + // and shift = cell_shift + wrap_i - wrap_j. + auto distance2 = tmp_dist2[aj]; + + auto index = neighbors.length(); + neighbors.set_pair(index, static_cast(idx_i), static_cast(idx_j)); + + if (options.return_shifts) { + neighbors.set_shift(index, shift); + } + if (options.return_distances) { + neighbors.set_distance(index, std::sqrt(distance2)); + } + if (options.return_vectors) { + auto vector = Vector{ + tmp_dx[aj], tmp_dy[aj], tmp_dz[aj] + }; + neighbors.set_vector(index, vector); + } + neighbors.increment_length(); + } + } + } + } + } + } + } + } + } + } + + if (options.sorted) { + neighbors.sort(); + } +} diff --git a/vesin/src/cpu_cell_list.cpp b/vesin/src/cpu_cell_list.cpp index 9db201cd..23736b23 100644 --- a/vesin/src/cpu_cell_list.cpp +++ b/vesin/src/cpu_cell_list.cpp @@ -393,6 +393,58 @@ static scalar_t* alloc(scalar_t* ptr, size_t size, size_t new_size) { return new_ptr; } +void GrowableNeighborList::ensure_capacity(size_t required) { + if (required <= this->capacity) { + return; + } + // Same exponential doubling policy as grow(), but seeded to the caller's + // target so we cover it in a single allocation rather than several + // grow() calls. The previous code grew on every set_* call past + // capacity which made the per-pair branch + realloc the dominant CPU + // cost (cachegrind: 34% of total instructions). + size_t new_size = std::max(this->capacity * 2, 1); + while (new_size < required) { + new_size *= 2; + } + + auto* new_pairs = alloc(neighbors.pairs, neighbors.length, new_size); + + int32_t (*new_shifts)[3] = nullptr; + if (options.return_shifts) { + new_shifts = alloc(neighbors.shifts, neighbors.length, new_size); + } + + double* new_distances = nullptr; + if (options.return_distances) { + new_distances = alloc(neighbors.distances, neighbors.length, new_size); + } + + double (*new_vectors)[3] = nullptr; + if (options.return_vectors) { + new_vectors = alloc(neighbors.vectors, neighbors.length, new_size); + } + + if ( + (new_pairs == nullptr) || + (options.return_shifts && new_shifts == nullptr) || + (options.return_distances && new_distances == nullptr) || + (options.return_vectors && new_vectors == nullptr) + ) { + std::free(new_pairs); + std::free(new_shifts); + std::free(new_distances); + std::free(new_vectors); + throw std::runtime_error("could not allocate memory for growing neighbor list"); + } + + this->neighbors.pairs = new_pairs; + this->neighbors.shifts = new_shifts; + this->neighbors.distances = new_distances; + this->neighbors.vectors = new_vectors; + + this->capacity = new_size; +} + void GrowableNeighborList::grow() { auto new_size = neighbors.length * 2; if (new_size == 0) { @@ -499,20 +551,39 @@ void GrowableNeighborList::sort() { std::iota(std::begin(indices), std::end(indices), 0); struct compare_pairs { - compare_pairs(size_t (*pairs_)[2]): - pairs(pairs_) {} + compare_pairs(size_t (*pairs_)[2], int32_t (*shifts_)[3]): + pairs(pairs_), + shifts(shifts_) {} bool operator()(int64_t a, int64_t b) const { - return pairs[a][0] < pairs[b][0]; + auto ia = static_cast(a); + auto ib = static_cast(b); + + if (pairs[ia][0] != pairs[ib][0]) { + return pairs[ia][0] < pairs[ib][0]; + } + if (pairs[ia][1] != pairs[ib][1]) { + return pairs[ia][1] < pairs[ib][1]; + } + if (shifts != nullptr) { + for (size_t dim = 0; dim < 3; dim++) { + if (shifts[ia][dim] != shifts[ib][dim]) { + return shifts[ia][dim] < shifts[ib][dim]; + } + } + } + + return false; } size_t (*pairs)[2]; + int32_t (*shifts)[3]; }; std::sort( std::begin(indices), std::end(indices), - compare_pairs(this->neighbors.pairs) + compare_pairs(this->neighbors.pairs, this->neighbors.shifts) ); // step 2: move all data according to the sorted indices. diff --git a/vesin/src/cpu_cell_list.hpp b/vesin/src/cpu_cell_list.hpp index 4cf35b48..a2be3e9a 100644 --- a/vesin/src/cpu_cell_list.hpp +++ b/vesin/src/cpu_cell_list.hpp @@ -124,6 +124,33 @@ class GrowableNeighborList { // allocate more memory & update capacity void grow(); + // Ensure `capacity >= required`, growing if needed. Lets the caller + // hoist the capacity check out of the hot per-pair loop so the + // unchecked set_* variants below are safe to use. + void ensure_capacity(size_t required); + + // Unchecked variants of the set_* methods above: the caller must have + // already called ensure_capacity so that `index < capacity`. The + // cachegrind profile attributed 34% of total instructions to the + // per-pair branch + write path; hoisting the check is the win. + void set_pair_unchecked(size_t index, size_t first, size_t second) { + this->neighbors.pairs[index][0] = first; + this->neighbors.pairs[index][1] = second; + } + void set_shift_unchecked(size_t index, vesin::CellShift shift) { + this->neighbors.shifts[index][0] = shift[0]; + this->neighbors.shifts[index][1] = shift[1]; + this->neighbors.shifts[index][2] = shift[2]; + } + void set_distance_unchecked(size_t index, double distance) { + this->neighbors.distances[index] = distance; + } + void set_vector_unchecked(size_t index, vesin::Vector vector) { + this->neighbors.vectors[index][0] = vector[0]; + this->neighbors.vectors[index][1] = vector[1]; + this->neighbors.vectors[index][2] = vector[2]; + } + // sort the pairs currently in the neighbor list void sort(); }; diff --git a/vesin/src/verlet.cpp b/vesin/src/verlet.cpp index 833f66fa..cd8c4d83 100644 --- a/vesin/src/verlet.cpp +++ b/vesin/src/verlet.cpp @@ -1,12 +1,26 @@ #include #include #include +#include +#ifdef VESIN_HAVE_HIGHWAY +#include +#else +#define HWY_ATTR +#define HWY_RESTRICT +#endif + +#include "cluster.hpp" #include "cpu_cell_list.hpp" #include "verlet.hpp" using namespace vesin; +namespace { +#ifdef VESIN_HAVE_HIGHWAY +namespace hn = hwy::HWY_NAMESPACE; +#endif + static BoundingBox make_box_like(const BoundingBox& box, const Vector* points, size_t n_points) { auto periodic = std::array{box.periodic(0), box.periodic(1), box.periodic(2)}; auto candidate_box = BoundingBox(box.matrix(), periodic.data()); @@ -14,6 +28,286 @@ static BoundingBox make_box_like(const BoundingBox& box, const Vector* points, s return candidate_box; } +static std::vector pack_simd_candidate_blocks( + const VesinNeighborList& candidates, + const std::vector& shift_vectors, + size_t& candidate_length +) { + auto blocks = std::vector(); + candidate_length = candidates.length; + blocks.reserve((candidates.length + CLUSTER_SIZE_CPU - 1) / CLUSTER_SIZE_CPU); + + auto block = cpu::VerletCandidateBlock(); + for (size_t k = 0; k < candidates.length; k++) { + auto lane = block.count; + block.first[lane] = candidates.pairs[k][0]; + block.second[lane] = candidates.pairs[k][1]; + block.shifts[lane] = CellShift{{ + candidates.shifts[k][0], + candidates.shifts[k][1], + candidates.shifts[k][2], + }}; + block.shift_x[lane] = shift_vectors[k][0]; + block.shift_y[lane] = shift_vectors[k][1]; + block.shift_z[lane] = shift_vectors[k][2]; + block.count += 1; + + if (block.count == CLUSTER_SIZE_CPU) { + blocks.push_back(block); + block = cpu::VerletCandidateBlock(); + } + } + + if (block.count != 0) { + blocks.push_back(block); + } + + return blocks; +} + +HWY_ATTR +static void simd_filter_deltas( + const double* HWY_RESTRICT dx, + const double* HWY_RESTRICT dy, + const double* HWY_RESTRICT dz, + double cutoff_sq, + double* HWY_RESTRICT dist_sq, + uint8_t* HWY_RESTRICT mask +) { +#ifdef VESIN_HAVE_HIGHWAY + const hn::ScalableTag d; + const size_t N = hn::Lanes(d); + const auto vcut = hn::Set(d, cutoff_sq); + + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane += N) { + auto vdx = hn::Load(d, dx + lane); + auto vdy = hn::Load(d, dy + lane); + auto vdz = hn::Load(d, dz + lane); + auto vdist = hn::MulAdd(vdx, vdx, hn::MulAdd(vdy, vdy, hn::Mul(vdz, vdz))); + auto vmask = hn::Lt(vdist, vcut); + + hn::Store(vdist, d, dist_sq + lane); + + uint8_t bits_buf[8] = {}; + hn::StoreMaskBits(d, vmask, bits_buf); + uint8_t bits = bits_buf[0]; + for (size_t k = 0; k < N && (lane + k) < CLUSTER_SIZE_CPU; k++) { + mask[lane + k] = (bits >> k) & 1; + } + } +#else + // Scalar fallback for the single-file dist build (no Highway). + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane++) { + const double ddx = dx[lane]; + const double ddy = dy[lane]; + const double ddz = dz[lane]; + const double d2 = ddx * ddx + ddy * ddy + ddz * ddz; + dist_sq[lane] = d2; + mask[lane] = (d2 < cutoff_sq) ? 1 : 0; + } +#endif +} + +// Vectorized gather of pair deltas using Highway for the Verlet recompute hot path. +// Gathers x/y/z for atom i and j using byte-offset gathers, computes (j - i) + shift +// in SIMD registers, and stores to the dx/dy/dz temporaries for the subsequent filter. +HWY_ATTR +static void gather_pair_deltas( + const Vector* HWY_RESTRICT points, + const size_t* HWY_RESTRICT first, + const size_t* HWY_RESTRICT second, + const double* HWY_RESTRICT shift_x, + const double* HWY_RESTRICT shift_y, + const double* HWY_RESTRICT shift_z, + size_t count, + double* HWY_RESTRICT dx, + double* HWY_RESTRICT dy, + double* HWY_RESTRICT dz +) { +#ifdef VESIN_HAVE_HIGHWAY + const hn::ScalableTag d; + const hn::ScalableTag di64; + const size_t N = hn::Lanes(d); + const double* base = reinterpret_cast(points); + const auto k_three = hn::Set(di64, int64_t{3}); + const auto k_one = hn::Set(di64, int64_t{1}); + + size_t lane = 0; + for (; lane + N <= count; lane += N) { + // Load atom indices as signed (safe for practical atom counts < 2^63) + // Note: first/second are size_t but values are non-negative. + alignas(64) int64_t i_idx[8] = {}; + alignas(64) int64_t j_idx[8] = {}; + for (size_t k = 0; k < N; k++) { + i_idx[k] = static_cast(first[lane + k]); + j_idx[k] = static_cast(second[lane + k]); + } + auto v_i = hn::Load(di64, i_idx); + auto v_j = hn::Load(di64, j_idx); + + // Element indices for x: atom * 3 + auto idx_i_x = hn::Mul(v_i, k_three); + auto idx_j_x = hn::Mul(v_j, k_three); + + auto vxi = hn::GatherIndex(d, base, idx_i_x); + auto vxj = hn::GatherIndex(d, base, idx_j_x); + + // y: +1 + auto idx_i_y = hn::Add(idx_i_x, k_one); + auto idx_j_y = hn::Add(idx_j_x, k_one); + auto vyi = hn::GatherIndex(d, base, idx_i_y); + auto vyj = hn::GatherIndex(d, base, idx_j_y); + + // z: +2 + auto idx_i_z = hn::Add(idx_i_y, k_one); + auto idx_j_z = hn::Add(idx_j_y, k_one); + auto vzi = hn::GatherIndex(d, base, idx_i_z); + auto vzj = hn::GatherIndex(d, base, idx_j_z); + + // deltas j - i + auto vdx = hn::Sub(vxj, vxi); + auto vdy = hn::Sub(vyj, vyi); + auto vdz = hn::Sub(vzj, vzi); + + // load and add precomputed Cartesian shifts + auto vsx = hn::Load(d, shift_x + lane); + auto vsy = hn::Load(d, shift_y + lane); + auto vsz = hn::Load(d, shift_z + lane); + + vdx = hn::Add(vdx, vsx); + vdy = hn::Add(vdy, vsy); + vdz = hn::Add(vdz, vsz); + + hn::Store(vdx, d, dx + lane); + hn::Store(vdy, d, dy + lane); + hn::Store(vdz, d, dz + lane); + } + + // Scalar tail for remainder (when count % N != 0) + for (; lane < count; lane++) { + auto i = first[lane]; + auto j = second[lane]; + dx[lane] = points[j][0] - points[i][0] + shift_x[lane]; + dy[lane] = points[j][1] - points[i][1] + shift_y[lane]; + dz[lane] = points[j][2] - points[i][2] + shift_z[lane]; + } +#else + // Scalar fallback for the single-file dist build (no Highway). + for (size_t lane = 0; lane < count; lane++) { + const auto i = first[lane]; + const auto j = second[lane]; + dx[lane] = points[j][0] - points[i][0] + shift_x[lane]; + dy[lane] = points[j][1] - points[i][1] + shift_y[lane]; + dz[lane] = points[j][2] - points[i][2] + shift_z[lane]; + } +#endif +} + +static void filter_simd_candidate_blocks( + const Vector* points, + const std::vector& blocks, + double cutoff_sq, + VesinOptions options, + VesinNeighborList& neighbors, + size_t initial_capacity, + size_t& output_capacity +) { + auto growable = cpu::GrowableNeighborList{neighbors, initial_capacity, options}; + growable.reset(); + + alignas(64) double dx[CLUSTER_SIZE_CPU]; + alignas(64) double dy[CLUSTER_SIZE_CPU]; + alignas(64) double dz[CLUSTER_SIZE_CPU]; + alignas(64) double dist_sq[CLUSTER_SIZE_CPU]; + alignas(64) double dist[CLUSTER_SIZE_CPU]; + uint8_t mask[CLUSTER_SIZE_CPU]; + const bool need_distance = options.return_distances; + + for (const auto& block : blocks) { + // Vectorized gather + arithmetic for the valid lanes (uses Highway Gather + // for random-access position loads, enabling better load parallelism on AVX2/AVX-512). + if (block.count > 0) { + gather_pair_deltas( + points, + block.first, + block.second, + block.shift_x, + block.shift_y, + block.shift_z, + block.count, + dx, + dy, + dz + ); + } + for (size_t lane = block.count; lane < CLUSTER_SIZE_CPU; lane++) { + dx[lane] = std::numeric_limits::infinity(); + dy[lane] = 0.0; + dz[lane] = 0.0; + } + + simd_filter_deltas(dx, dy, dz, cutoff_sq, dist_sq, mask); + + // SIMD sqrt: precompute once for the whole block instead of one + // scalar std::sqrt per kept pair in the lane loop. Highway picks + // the widest available vector lane (AVX-512: 8 doubles, AVX2: 4). + // Wasted lanes for filtered-out pairs are cheap relative to the + // scalar loop overhead. Only run when distances are requested. + // Falls back to a scalar loop for the single-file dist build. + if (need_distance) { +#ifdef VESIN_HAVE_HIGHWAY + const hn::ScalableTag d; + const size_t N = hn::Lanes(d); + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane += N) { + hn::Store(hn::Sqrt(hn::Load(d, dist_sq + lane)), d, dist + lane); + } +#else + for (size_t lane = 0; lane < CLUSTER_SIZE_CPU; lane++) { + dist[lane] = std::sqrt(dist_sq[lane]); + } +#endif + } + + // Pre-grow once per block. Worst case is every lane in this block + // passes the filter, so reserve growable.length() + block.count up + // front. This hoists the per-pair capacity branch out of the hot + // lane loop below (cachegrind: 34% of total instructions were in + // the per-pair set_*'s capacity check + grow); the unchecked + // set_*_unchecked variants below skip it entirely. + growable.ensure_capacity(growable.length() + block.count); + + for (size_t lane = 0; lane < block.count; lane++) { + if (!mask[lane]) { + continue; + } + + auto index = growable.length(); + growable.set_pair_unchecked(index, block.first[lane], block.second[lane]); + + if (options.return_shifts) { + growable.set_shift_unchecked(index, block.shifts[lane]); + } + + if (need_distance) { + growable.set_distance_unchecked(index, dist[lane]); + } + + if (options.return_vectors) { + growable.set_vector_unchecked(index, Vector{dx[lane], dy[lane], dz[lane]}); + } + + growable.increment_length(); + } + } + + if (options.sorted) { + growable.sort(); + } + + output_capacity = growable.capacity; +} +} // namespace + cpu::VerletState::~VerletState() { this->clear_candidates(); } @@ -24,13 +318,20 @@ void cpu::VerletState::clear_candidates() { } this->candidates = VesinNeighborList(); + this->candidate_shift_vectors.clear(); + this->simd_candidate_blocks.clear(); + this->simd_candidate_length = 0; + this->cluster_grid = ClusterGrid(); + this->cluster_candidates.clear(); + this->use_cluster_candidates = false; this->has_cache = false; this->ref_positions.clear(); this->n_points = 0; } void cpu::VerletState::set_options(VesinOptions options) { - if (this->options.cutoff != options.cutoff || this->options.skin != options.skin || this->options.full != options.full) { + if (this->options.cutoff != options.cutoff || this->options.skin != options.skin || this->options.full != options.full || + this->options.algorithm != options.algorithm) { this->clear_candidates(); } @@ -94,7 +395,7 @@ void cpu::VerletState::rebuild( build_options.cutoff = this->options.cutoff + this->options.skin; build_options.full = this->options.full; build_options.sorted = false; - build_options.algorithm = VesinCellList; + build_options.algorithm = this->options.algorithm; build_options.return_shifts = true; build_options.return_distances = false; build_options.return_vectors = false; @@ -102,8 +403,33 @@ void cpu::VerletState::rebuild( this->candidates.device = {VesinCPU, 0}; auto candidate_box = make_box_like(box, points, n_points); - size_t candidate_capacity = 0; - cpu::stateless_neighbors(points, n_points, std::move(candidate_box), build_options, this->candidates, candidate_capacity); + if (build_options.algorithm == VesinAutoAlgorithm && n_points >= CLUSTER_PAIR_THRESHOLD) { + this->cluster_grid = build_cluster_grid(points, n_points, candidate_box, build_options.cutoff); + this->cluster_candidates = build_cluster_pair_candidates(this->cluster_grid, candidate_box, build_options.cutoff); + this->use_cluster_candidates = true; + cpu::cluster_pair_neighbors(points, n_points, candidate_box, build_options, this->candidates); + } else { + size_t candidate_capacity = 0; + cpu::stateless_neighbors(points, n_points, std::move(candidate_box), build_options, this->candidates, candidate_capacity); + } + + this->candidate_shift_vectors.reserve(this->candidates.length); + for (size_t k = 0; k < this->candidates.length; k++) { + auto shift = CellShift{{ + this->candidates.shifts[k][0], + this->candidates.shifts[k][1], + this->candidates.shifts[k][2], + }}; + this->candidate_shift_vectors.push_back(shift_cartesian(shift, box)); + } + + if (this->use_cluster_candidates) { + this->simd_candidate_blocks = pack_simd_candidate_blocks( + this->candidates, + this->candidate_shift_vectors, + this->simd_candidate_length + ); + } this->n_points = n_points; this->ref_positions.resize(n_points * 3); @@ -127,9 +453,27 @@ void cpu::VerletState::recompute( auto initial_capacity = std::max(output_capacity, neighbors.length); + if (this->use_cluster_candidates) { + filter_simd_candidate_blocks( + points, + this->simd_candidate_blocks, + cutoff_sq, + options, + neighbors, + initial_capacity, + output_capacity + ); + return; + } + auto growable = cpu::GrowableNeighborList{neighbors, initial_capacity, options}; growable.reset(); + // Pre-grow once for the worst case (every candidate passes the filter) + // so the inner loop never branches on capacity. Mirror change to the + // SIMD-block path above. + growable.ensure_capacity(this->candidates.length); + // The cached list is an over-complete Verlet candidate list. Each call // filters candidates with the exact cutoff and requested shift/vector outputs. for (size_t k = 0; k < this->candidates.length; k++) { @@ -142,23 +486,23 @@ void cpu::VerletState::recompute( this->candidates.shifts[k][2], }}; - auto vec = points[j] - points[i] + shift.cartesian(box); + auto vec = points[j] - points[i] + this->candidate_shift_vectors[k]; double dist_sq = vec.dot(vec); if (dist_sq < cutoff_sq) { auto idx = growable.length(); - growable.set_pair(idx, i, j); + growable.set_pair_unchecked(idx, i, j); if (options.return_shifts) { - growable.set_shift(idx, shift); + growable.set_shift_unchecked(idx, shift); } if (options.return_distances) { - growable.set_distance(idx, std::sqrt(dist_sq)); + growable.set_distance_unchecked(idx, std::sqrt(dist_sq)); } if (options.return_vectors) { - growable.set_vector(idx, vec); + growable.set_vector_unchecked(idx, vec); } growable.increment_length(); diff --git a/vesin/src/verlet.hpp b/vesin/src/verlet.hpp index 33d9fa52..78b1588d 100644 --- a/vesin/src/verlet.hpp +++ b/vesin/src/verlet.hpp @@ -5,12 +5,32 @@ #include #include +#include "cluster.hpp" #include "types.hpp" #include "vesin.h" namespace vesin { namespace cpu { +/// Fixed-width candidate block for SIMD filtering of cached Verlet pairs. +struct alignas(64) VerletCandidateBlock { + size_t count = 0; + alignas(64) size_t first[CLUSTER_SIZE_CPU] = {}; + alignas(64) size_t second[CLUSTER_SIZE_CPU] = {}; + alignas(64) CellShift shifts[CLUSTER_SIZE_CPU] = {}; + alignas(64) double shift_x[CLUSTER_SIZE_CPU] = {}; + alignas(64) double shift_y[CLUSTER_SIZE_CPU] = {}; + alignas(64) double shift_z[CLUSTER_SIZE_CPU] = {}; +}; + +static_assert(alignof(VerletCandidateBlock) >= 64); +static_assert(offsetof(VerletCandidateBlock, first) % 64 == 0); +static_assert(offsetof(VerletCandidateBlock, second) % 64 == 0); +static_assert(offsetof(VerletCandidateBlock, shifts) % 64 == 0); +static_assert(offsetof(VerletCandidateBlock, shift_x) % 64 == 0); +static_assert(offsetof(VerletCandidateBlock, shift_y) % 64 == 0); +static_assert(offsetof(VerletCandidateBlock, shift_z) % 64 == 0); + /// State for a cached, on-CPU Verlet neighbor list. /// /// The state stores: @@ -50,10 +70,7 @@ struct VerletState { const BoundingBox& box ) const; - /// Build the over-complete candidate list at `cutoff + skin`. - /// - /// The rebuild operation stores candidates in full `VesinNeighborList` form - /// and captures the state used to validate future `needs_rebuild` checks. + /// Build the over-complete candidate cache at `cutoff + skin`. void rebuild( const Vector* points, size_t n_points, @@ -75,9 +92,17 @@ struct VerletState { /// Number of pairs currently stored in the cached candidate list. size_t candidate_count() const { + if (use_cluster_candidates) { + return cluster_candidates.size(); + } return candidates.length; } + /// Number of atom-pair candidates packed into SIMD recompute blocks. + size_t simd_candidate_count() const { + return simd_candidate_length; + } + /// Reference positions at the time the candidates were built. std::vector ref_positions; /// Box matrix used for candidate generation and displacement validation. @@ -92,6 +117,19 @@ struct VerletState { /// The list is kept in normal neighbor-list representation so rebuild and /// recompute paths can share storage and filtering logic. VesinNeighborList candidates; + /// Cartesian shift vector for each materialized candidate pair. + std::vector candidate_shift_vectors; + /// Fixed-width blocks used by the SIMD cached-candidate recompute path. + std::vector simd_candidate_blocks; + /// Number of atom-pair lanes stored in `simd_candidate_blocks`. + size_t simd_candidate_length = 0; + + /// Cluster grid used by cluster-backed Verlet candidate caches. + ClusterGrid cluster_grid; + /// Over-complete cluster-pair candidates generated at `cutoff + skin`. + std::vector cluster_candidates; + /// Whether the active cache is represented by cluster-pair candidates. + bool use_cluster_candidates = false; /// Options used to build the current cache. VesinOptions options = {}; diff --git a/vesin/src/vesin.cpp b/vesin/src/vesin.cpp index 56ca778f..677e3ddc 100644 --- a/vesin/src/vesin.cpp +++ b/vesin/src/vesin.cpp @@ -3,6 +3,7 @@ #include #include +#include "cluster.hpp" #include "cpu_cell_list.hpp" #include "vesin.h" #include "vesin_cuda.hpp" @@ -81,16 +82,27 @@ extern "C" int vesin_neighbors( {{box[2][0], box[2][1], box[2][2]}}, }}}; - auto box = vesin::BoundingBox(matrix, periodic); - box.make_bounding_for(points, n_points); - - vesin::cpu::neighbors( - reinterpret_cast(points), - n_points, - std::move(box), - options, - *neighbors - ); + auto bounding_box = vesin::BoundingBox(matrix, periodic); + bounding_box.make_bounding_for(points, n_points); + auto points_vec = reinterpret_cast(points); + + if (options.skin == 0.0 && options.algorithm == VesinAutoAlgorithm && n_points >= vesin::CLUSTER_PAIR_THRESHOLD) { + vesin::cpu::cluster_pair_neighbors( + points_vec, + n_points, + bounding_box, + options, + *neighbors + ); + } else { + vesin::cpu::neighbors( + points_vec, + n_points, + std::move(bounding_box), + options, + *neighbors + ); + } } else if (device.type == VesinCUDA) { vesin::cuda::neighbors( points, diff --git a/vesin/tests/cluster_pair.cpp b/vesin/tests/cluster_pair.cpp new file mode 100644 index 00000000..25f6e68f --- /dev/null +++ b/vesin/tests/cluster_pair.cpp @@ -0,0 +1,638 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace Catch::Matchers; + +#include + +// Catch2 (>=3) needs to stringify both sides of CHECK(x == y) for the +// failure diff. libstdc++ 9 (the GCC 9 shipped with the ubuntu-24.04 CI +// image) does not detect std::tuple<...> as range-printable, so Catch2's +// default StringMaker SFINAE chain hits an enable_if::type +// instantiation error inside catch_tostring.hpp. Providing an explicit +// StringMaker for the (i, j, sx, sy, sz) tuple short-circuits the chain. +namespace Catch { +template <> +struct StringMaker> { + static std::string convert(const std::tuple& t) { + return "(i=" + std::to_string(std::get<0>(t)) + ", j=" + std::to_string(std::get<1>(t)) + ", sx=" + std::to_string(std::get<2>(t)) + ", sy=" + std::to_string(std::get<3>(t)) + ", sz=" + std::to_string(std::get<4>(t)) + ")"; + } +}; +} // namespace Catch + +/// Helper: build a simple cubic lattice with n^3 atoms +static std::vector> cubic_lattice(int n, double spacing) { + std::vector> points; + for (int ix = 0; ix < n; ix++) { + for (int iy = 0; iy < n; iy++) { + for (int iz = 0; iz < n; iz++) { + points.push_back({ + ix * spacing, + iy * spacing, + iz * spacing, + }); + } + } + } + return points; +} + +/// Helper: collect (i, j, shift) tuples into a set for comparison +using PairSet = std::set>; + +static PairSet collect_pairs(const VesinNeighborList& nl) { + PairSet result; + for (size_t k = 0; k < nl.length; k++) { + result.emplace( + nl.pairs[k][0], nl.pairs[k][1], nl.shifts[k][0], nl.shifts[k][1], nl.shifts[k][2] + ); + } + return result; +} + +/// Compute a neighbor list forcing cell-list algorithm +static VesinNeighborList compute_with_algorithm( + const double (*points)[3], + size_t n_points, + const double box[3][3], + bool periodic[3], + double cutoff, + bool full_list, + VesinAlgorithm algorithm +) { + auto options = VesinOptions(); + options.cutoff = cutoff; + options.full = full_list; + options.sorted = false; + options.algorithm = algorithm; + options.return_shifts = true; + options.return_distances = true; + options.return_vectors = true; + + VesinNeighborList neighbors; + const char* error_message = nullptr; + auto status = vesin_neighbors( + points, n_points, box, periodic, {VesinCPU, 0}, options, &neighbors, &error_message + ); + REQUIRE(status == EXIT_SUCCESS); + if (error_message != nullptr) { + FAIL("Error: " << error_message); + } + return neighbors; +} + +TEST_CASE("Cluster-pair: correctness vs cell-list on 4x4x4 lattice") { + // 4^3 = 64 atoms, below cluster-pair threshold (256) so Auto + // uses cell-list. Verifies both paths agree. + auto points = cubic_lattice(4, 1.5); + REQUIRE(points.size() == 64); + + double box_len = 4 * 1.5; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.5; + + // Cell-list (forced) + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + // Auto uses cell-list below the cluster-pair threshold. + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: full list on 4x4x4 lattice") { + auto points = cubic_lattice(4, 1.5); + double box_len = 4 * 1.5; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.5; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: non-periodic system") { + auto points = cubic_lattice(5, 1.2); // 125 atoms + double box[3][3] = {{0, 0, 0}, {0, 0, 0}, {0, 0, 0}}; + bool periodic[3] = {false, false, false}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: distances match cell-list") { + auto points = cubic_lattice(4, 1.5); + double box_len = 4 * 1.5; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.5; + + // Both sorted so we can compare element-wise + auto options = VesinOptions(); + options.cutoff = cutoff; + options.full = false; + options.sorted = true; + options.return_shifts = true; + options.return_distances = true; + options.return_vectors = true; + + VesinNeighborList cl_nl; + const char* error_message = nullptr; + + options.algorithm = VesinCellList; + auto status = vesin_neighbors( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + {VesinCPU, 0}, + options, + &cl_nl, + &error_message + ); + REQUIRE(status == EXIT_SUCCESS); + + VesinNeighborList auto_nl; + options.algorithm = VesinAutoAlgorithm; + status = vesin_neighbors( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + {VesinCPU, 0}, + options, + &auto_nl, + &error_message + ); + REQUIRE(status == EXIT_SUCCESS); + + REQUIRE(cl_nl.length == auto_nl.length); + + for (size_t k = 0; k < cl_nl.length; k++) { + CHECK(cl_nl.pairs[k][0] == auto_nl.pairs[k][0]); + CHECK(cl_nl.pairs[k][1] == auto_nl.pairs[k][1]); + CHECK(cl_nl.shifts[k][0] == auto_nl.shifts[k][0]); + CHECK(cl_nl.shifts[k][1] == auto_nl.shifts[k][1]); + CHECK(cl_nl.shifts[k][2] == auto_nl.shifts[k][2]); + CHECK_THAT(cl_nl.distances[k], WithinULP(auto_nl.distances[k], 4)); + } + + vesin_free(&cl_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: triclinic box") { + // Triclinic box with 64 atoms -- exercises the shifted BB test + // for periodic images where cell_shift != (0,0,0). + auto points = cubic_lattice(4, 1.2); // 64 atoms + REQUIRE(points.size() == 64); + + // Triclinic box: off-diagonal elements create non-trivial periodic shifts + double box[3][3] = {{4.8, 0.0, 0.0}, {1.2, 4.8, 0.0}, {0.8, 0.6, 4.8}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: triclinic box full list") { + // Same triclinic geometry but with full list + auto points = cubic_lattice(4, 1.2); + double box[3][3] = {{4.8, 0.0, 0.0}, {1.2, 4.8, 0.0}, {0.8, 0.6, 4.8}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: large periodic system with BB rejection") { + // 125 atoms in periodic box. Many pairs come from periodic images, + // so this exercises the shifted BB distance test under load. + auto points = cubic_lattice(5, 1.0); // 125 atoms + double box_len = 5.0; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 1.8; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() > 0); + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair: larger system 5x5x5") { + auto points = cubic_lattice(5, 1.2); // 125 atoms + double box_len = 5 * 1.2; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +// --- Tests with N >= 256 that exercise the cluster-pair SIMD path --- + +TEST_CASE("Cluster-pair SIMD: 7x7x7 periodic half list") { + // 7^3 = 343 atoms -> above threshold (256), Auto uses cluster-pair + auto points = cubic_lattice(7, 1.2); + REQUIRE(points.size() == 343); + + double box_len = 7 * 1.2; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() > 0); + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair SIMD: 7x7x7 periodic full list") { + auto points = cubic_lattice(7, 1.2); + double box_len = 7 * 1.2; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + true, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() > 0); + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair SIMD: 7x7x7 non-periodic") { + auto points = cubic_lattice(7, 1.2); // 343 atoms + double box[3][3] = {{0, 0, 0}, {0, 0, 0}, {0, 0, 0}}; + bool periodic[3] = {false, false, false}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() > 0); + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair SIMD: triclinic 7x7x7") { + auto points = cubic_lattice(7, 1.2); // 343 atoms + // Triclinic box + double box[3][3] = {{8.4, 0.0, 0.0}, {2.1, 8.4, 0.0}, {1.4, 1.05, 8.4}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto cell_list_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinCellList + ); + + auto auto_nl = compute_with_algorithm( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + cutoff, + false, + VesinAutoAlgorithm + ); + + auto cl_pairs = collect_pairs(cell_list_nl); + auto auto_pairs = collect_pairs(auto_nl); + + CHECK(cl_pairs.size() > 0); + CHECK(cl_pairs.size() == auto_pairs.size()); + CHECK(cl_pairs == auto_pairs); + + vesin_free(&cell_list_nl); + vesin_free(&auto_nl); +} + +TEST_CASE("Cluster-pair SIMD: distances match cell-list 7x7x7") { + auto points = cubic_lattice(7, 1.2); // 343 atoms + double box_len = 7 * 1.2; + double box[3][3] = {{box_len, 0, 0}, {0, box_len, 0}, {0, 0, box_len}}; + bool periodic[3] = {true, true, true}; + double cutoff = 2.0; + + auto options = VesinOptions(); + options.cutoff = cutoff; + options.full = false; + options.sorted = true; + options.return_shifts = true; + options.return_distances = true; + options.return_vectors = true; + + VesinNeighborList cl_nl; + const char* error_message = nullptr; + + options.algorithm = VesinCellList; + auto status = vesin_neighbors( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + {VesinCPU, 0}, + options, + &cl_nl, + &error_message + ); + REQUIRE(status == EXIT_SUCCESS); + + VesinNeighborList auto_nl; + options.algorithm = VesinAutoAlgorithm; + status = vesin_neighbors( + reinterpret_cast(points.data()), + points.size(), + box, + periodic, + {VesinCPU, 0}, + options, + &auto_nl, + &error_message + ); + REQUIRE(status == EXIT_SUCCESS); + + REQUIRE(cl_nl.length == auto_nl.length); + + for (size_t k = 0; k < cl_nl.length; k++) { + CHECK(cl_nl.pairs[k][0] == auto_nl.pairs[k][0]); + CHECK(cl_nl.pairs[k][1] == auto_nl.pairs[k][1]); + CHECK(cl_nl.shifts[k][0] == auto_nl.shifts[k][0]); + CHECK(cl_nl.shifts[k][1] == auto_nl.shifts[k][1]); + CHECK(cl_nl.shifts[k][2] == auto_nl.shifts[k][2]); + CHECK_THAT(cl_nl.distances[k], WithinULP(auto_nl.distances[k], 4)); + } + + vesin_free(&cl_nl); + vesin_free(&auto_nl); +} diff --git a/vesin/tests/verlet.cpp b/vesin/tests/verlet.cpp index c36ade62..e707e972 100644 --- a/vesin/tests/verlet.cpp +++ b/vesin/tests/verlet.cpp @@ -1,8 +1,16 @@ #include +#include +#include + +#include "../src/cluster.hpp" #include "../src/cpu_cell_list.hpp" #include "../src/verlet.hpp" +using namespace Catch::Matchers; + +#define CHECK_APPROX_EQUAL(a, b) CHECK_THAT(a, WithinULP(b, 4)); + static vesin::BoundingBox make_box(const double (*points)[3], size_t n_points, const double matrix[3][3], const bool periodic[3]) { auto box_matrix = vesin::Matrix{{{ {{matrix[0][0], matrix[0][1], matrix[0][2]}}, @@ -15,6 +23,49 @@ static vesin::BoundingBox make_box(const double (*points)[3], size_t n_points, c return box; } +static vesin::BoundingBox make_box(const std::vector& points, const double matrix[3][3], const bool periodic[3]) { + auto box_matrix = vesin::Matrix{{{ + {{matrix[0][0], matrix[0][1], matrix[0][2]}}, + {{matrix[1][0], matrix[1][1], matrix[1][2]}}, + {{matrix[2][0], matrix[2][1], matrix[2][2]}}, + }}}; + + auto box = vesin::BoundingBox(box_matrix, periodic); + box.make_bounding_for(reinterpret_cast(points.data()), points.size()); + return box; +} + +static std::vector lattice_points(size_t edge, double spacing) { + auto points = std::vector(); + points.reserve(edge * edge * edge); + + for (size_t z = 0; z < edge; z++) { + for (size_t y = 0; y < edge; y++) { + for (size_t x = 0; x < edge; x++) { + points.push_back(vesin::Vector{ + spacing * static_cast(x), + spacing * static_cast(y), + spacing * static_cast(z), + }); + } + } + } + + return points; +} + +static std::vector displaced_points(const std::vector& points) { + auto displaced = points; + + for (size_t i = 0; i < displaced.size(); i++) { + displaced[i][0] += (static_cast(i % 3) - 1.0) * 0.01; + displaced[i][1] += (static_cast((i / 3) % 3) - 1.0) * 0.01; + displaced[i][2] += (static_cast((i / 9) % 3) - 1.0) * 0.01; + } + + return displaced; +} + TEST_CASE("Verlet recompute keeps allocation capacity across shorter output") { double box[3][3] = {{0.0}}; bool periodic[3] = {false, false, false}; @@ -72,6 +123,137 @@ TEST_CASE("Verlet recompute keeps allocation capacity across shorter output") { vesin_free(&neighbors); } +TEST_CASE("Verlet cache invalidates when candidate algorithm changes") { + double box[3][3] = {{0.0}}; + bool periodic[3] = {false, false, false}; + + auto options = VesinOptions(); + options.cutoff = 1.0; + options.skin = 0.6; + options.full = false; + options.sorted = false; + options.algorithm = VesinCellList; + options.return_shifts = true; + options.return_distances = false; + options.return_vectors = false; + + double points[][3] = { + {0.0, 0.0, 0.0}, + {0.9, 0.0, 0.0}, + {1.8, 0.0, 0.0}, + {2.7, 0.0, 0.0}, + }; + + auto state = vesin::cpu::VerletState(); + state.set_options(options); + auto box_state = make_box(points, 4, box, periodic); + state.rebuild(reinterpret_cast(points), 4, box_state); + + REQUIRE(state.candidate_count() == 3); + + options.algorithm = VesinAutoAlgorithm; + state.set_options(options); + CHECK(state.candidate_count() == 0); +} + +TEST_CASE("Auto Verlet cache stores cluster candidates below atom-pair count") { + double box_matrix[3][3] = {{0.0}}; + bool periodic[3] = {false, false, false}; + + auto points = lattice_points(8, 0.9); + REQUIRE(points.size() >= vesin::CLUSTER_PAIR_THRESHOLD); + + auto options = VesinOptions(); + options.cutoff = 1.0; + options.skin = 0.35; + options.full = false; + options.sorted = false; + options.algorithm = VesinCellList; + options.return_shifts = true; + options.return_distances = false; + options.return_vectors = false; + + auto box = make_box(points, box_matrix, periodic); + + auto cell_state = vesin::cpu::VerletState(); + cell_state.set_options(options); + cell_state.rebuild(points.data(), points.size(), box); + auto atom_candidate_count = cell_state.candidate_count(); + REQUIRE(atom_candidate_count > 0); + + options.algorithm = VesinAutoAlgorithm; + auto auto_state = vesin::cpu::VerletState(); + auto_state.set_options(options); + auto_state.rebuild(points.data(), points.size(), box); + + CHECK(auto_state.candidates.length > 0); + CHECK(auto_state.simd_candidate_count() == auto_state.candidates.length); + CHECK(auto_state.candidate_count() < atom_candidate_count); +} + +TEST_CASE("Auto Verlet cluster cache matches exact cell-list output after small displacements") { + double box_matrix[3][3] = {{0.0}}; + bool periodic[3] = {false, false, false}; + + auto reference_points = lattice_points(8, 0.9); + auto current_points = displaced_points(reference_points); + + auto options = VesinOptions(); + options.cutoff = 1.0; + options.skin = 0.35; + options.full = false; + options.sorted = true; + options.algorithm = VesinAutoAlgorithm; + options.return_shifts = true; + options.return_distances = true; + options.return_vectors = true; + + auto reference_box = make_box(reference_points, box_matrix, periodic); + auto current_box = make_box(current_points, box_matrix, periodic); + + auto auto_state = vesin::cpu::VerletState(); + auto_state.set_options(options); + auto_state.rebuild(reference_points.data(), reference_points.size(), reference_box); + REQUIRE_FALSE(auto_state.needs_rebuild(current_points.data(), current_points.size(), current_box)); + + auto actual = VesinNeighborList(); + size_t actual_capacity = 0; + auto_state.recompute(current_points.data(), current_box, options, actual, actual_capacity); + + auto exact_options = options; + exact_options.skin = 0.0; + exact_options.algorithm = VesinCellList; + + auto expected = VesinNeighborList(); + size_t expected_capacity = 0; + vesin::cpu::stateless_neighbors( + current_points.data(), + current_points.size(), + make_box(current_points, box_matrix, periodic), + exact_options, + expected, + expected_capacity + ); + + REQUIRE(actual.length == expected.length); + for (size_t k = 0; k < actual.length; k++) { + CHECK(actual.pairs[k][0] == expected.pairs[k][0]); + CHECK(actual.pairs[k][1] == expected.pairs[k][1]); + CHECK(actual.shifts[k][0] == expected.shifts[k][0]); + CHECK(actual.shifts[k][1] == expected.shifts[k][1]); + CHECK(actual.shifts[k][2] == expected.shifts[k][2]); + CHECK_APPROX_EQUAL(actual.distances[k], expected.distances[k]); + CHECK_APPROX_EQUAL(actual.vectors[k][0], expected.vectors[k][0]); + CHECK_APPROX_EQUAL(actual.vectors[k][1], expected.vectors[k][1]); + CHECK_APPROX_EQUAL(actual.vectors[k][2], expected.vectors[k][2]); + } + + actual.device = {VesinCPU, 0}; + expected.device = {VesinCPU, 0}; + vesin_free(&actual); + vesin_free(&expected); +} + TEST_CASE("Periodic wrapped coordinates use minimum-image distance for rebuild") { double box_matrix[3][3] = {{1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0}}; bool periodic[3] = {true, true, true};