feat: SIMD-accelerated CPU Verlet (cluster-pair search + packed candidate blocks)

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])
 

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

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)

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)

vesin/src/cluster.hpp

@@ -0,0 +1,175 @@
+#ifndef VESIN_CLUSTER_HPP
+#define VESIN_CLUSTER_HPP
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#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<Cluster> clusters;
+
+    // Grid dimensions (number of cells in each direction)
+    std::array<int32_t, 3> 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<int32_t> 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<CellShift> 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<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<ClusterPairCandidate> 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<ClusterPairCandidate>& 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