diff --git a/vesin/CMakeLists.txt b/vesin/CMakeLists.txt index b17ac288..9be7efee 100644 --- a/vesin/CMakeLists.txt +++ b/vesin/CMakeLists.txt @@ -63,6 +63,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 ) @@ -103,8 +104,22 @@ 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) +FetchContent_MakeAvailable(highway) + +target_link_libraries(vesin_objects PRIVATE gpulite hwy) +target_link_libraries(vesin PRIVATE gpulite hwy) # Create generated directory file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/generated) diff --git a/vesin/src/cluster.hpp b/vesin/src/cluster.hpp new file mode 100644 index 00000000..aa704068 --- /dev/null +++ b/vesin/src/cluster.hpp @@ -0,0 +1,128 @@ +#ifndef VESIN_CLUSTER_HPP +#define VESIN_CLUSTER_HPP + +#include +#include +#include +#include + +#include "vesin.h" +#include "types.hpp" + +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; + +/// 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(cell_matrix). Indexed by original atom + // index. Used to avoid per-pair matrix multiply in the inner loop. + std::vector wrapped_positions; +}; + +/// 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, + BoundingBox box, + 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, + BoundingBox cell, + VesinOptions options, + VesinNeighborList& neighbors +); + +} // namespace cpu +} // 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..64a39960 --- /dev/null +++ b/vesin/src/cluster_pair_search.cpp @@ -0,0 +1,464 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "hwy/highway.h" + +#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, + 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 + }; + + auto cell_matrix = box.matrix(); + 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(cell_matrix); + } + + // 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 +// --------------------------------------------------------------------------- + +namespace { +namespace hn = hwy::HWY_NAMESPACE; + +/// 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 +) { + 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; +} +} // anonymous namespace + +void vesin::cpu::cluster_pair_neighbors( + const Vector* points, + size_t n_points, + BoundingBox cell, + VesinOptions options, + VesinNeighborList& raw_neighbors +) { + auto grid = build_cluster_grid(points, n_points, cell, options.cutoff); + + auto cell_matrix = cell.matrix(); + 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_matrix); + 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_matrix) + // 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/vesin.cpp b/vesin/src/vesin.cpp index 13789a57..72579ea7 100644 --- a/vesin/src/vesin.cpp +++ b/vesin/src/vesin.cpp @@ -2,10 +2,17 @@ #include #include +#include "cluster.hpp" #include "cpu_cell_list.hpp" #include "vesin.h" #include "vesin_cuda.hpp" +/// Threshold for switching from cell-list to cluster-pair search. +/// This is an internal auto-dispatch parameter, not exposed in the C API +/// (VesinAlgorithm only has Auto, BruteForce, CellList). Cluster-pair is +/// selected automatically when N >= this threshold and algorithm is Auto. +#define CLUSTER_PAIR_THRESHOLD 256 + // used to store dynamically allocated error messages before giving a pointer // to them back to the user thread_local std::string LAST_ERROR; @@ -75,13 +82,26 @@ extern "C" int vesin_neighbors( {{box[2][0], box[2][1], box[2][2]}}, }}}; - vesin::cpu::neighbors( - reinterpret_cast(points), - n_points, - vesin::BoundingBox(matrix, periodic), - options, - *neighbors - ); + auto bounding_box = vesin::BoundingBox(matrix, periodic); + auto points_vec = reinterpret_cast(points); + + if (options.algorithm == VesinAutoAlgorithm && n_points >= CLUSTER_PAIR_THRESHOLD) { + vesin::cpu::cluster_pair_neighbors( + points_vec, + n_points, + bounding_box, + options, + *neighbors + ); + } else { + vesin::cpu::neighbors( + points_vec, + n_points, + 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..0b7e55a4 --- /dev/null +++ b/vesin/tests/cluster_pair.cpp @@ -0,0 +1,493 @@ +#include +#include +#include +#include +#include + +#include +#include + +using namespace Catch::Matchers; + +#include + +/// 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 (should use cluster-pair for N=64) + 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); +}