kernels/neighbor_list.cu

@@ -1 +1,194 @@
 #include "neighbor_list.cuh"
+
+/**
+ * Step 1: Assign particles to cells
+ */
+__global__ void assign_particles_to_cells_kernel(const float4 *positions,
+                                                 int *particle_cells,
+                                                 const CellList cell_list,
+                                                 size_t n_particles) {
+  size_t i = get_thread_id();
+  if (i >= n_particles)
+    return;
+
+  float4 pos = positions[i];
+  float3 pos3 = make_float3(pos.x, pos.y, pos.z);
+
+  int3 cell_coords = cell_list.get_cell_coords(pos3);
+
+  // Clamp to valid range (handle edge cases)
+  cell_coords.x = max(0, min(cell_coords.x, cell_list.grid_size.x - 1));
+  cell_coords.y = max(0, min(cell_coords.y, cell_list.grid_size.y - 1));
+  cell_coords.z = max(0, min(cell_coords.z, cell_list.grid_size.z - 1));
+
+  particle_cells[i] =
+      cell_list.get_cell_index(cell_coords.x, cell_coords.y, cell_coords.z);
+}
+
+/**
+ * Step 2: Find cell boundaries after sorting
+ */
+__global__ void find_cell_boundaries_kernel(const int *sorted_particle_cells,
+                                            int *cell_starts, int *cell_ends,
+                                            size_t n_particles,
+                                            size_t total_cells) {
+  size_t i = get_thread_id();
+  if (i >= n_particles)
+    return;
+
+  int cell = sorted_particle_cells[i];
+
+  // Check if this is the start of a new cell
+  if (i == 0 || sorted_particle_cells[i - 1] != cell) {
+    cell_starts[cell] = i;
+  }
+
+  // Check if this is the end of a cell
+  if (i == n_particles - 1 || sorted_particle_cells[i + 1] != cell) {
+    cell_ends[cell] = i + 1;
+  }
+}
+
+/**
+ * Step 3: Build actual neighbor lists using cell lists
+ */
+__global__ void
+build_neighbor_lists_kernel(const float4 *positions, const CellList cell_list,
+                            NeighborList neighbor_list, float cutoff_squared,
+                            const float3 *box_lengths, size_t n_particles) {
+  size_t i = get_thread_id();
+  if (i >= n_particles)
+    return;
+
+  float4 pos_i = positions[i];
+  float3 my_pos = make_float3(pos_i.x, pos_i.y, pos_i.z);
+
+  int3 my_cell = cell_list.get_cell_coords(my_pos);
+  int neighbor_count = 0;
+  int max_neighbors = neighbor_list.max_neighbors;
+  int *my_neighbors = neighbor_list.get_neighbors(i);
+
+  // Search neighboring cells (3x3x3 = 27 cells including self)
+  for (int dz = -1; dz <= 1; dz++) {
+    for (int dy = -1; dy <= 1; dy++) {
+      for (int dx = -1; dx <= 1; dx++) {
+
+        int3 neighbor_cell =
+            make_int3(my_cell.x + dx, my_cell.y + dy, my_cell.z + dz);
+
+        // Check bounds
+        if (neighbor_cell.x < 0 || neighbor_cell.x >= cell_list.grid_size.x ||
+            neighbor_cell.y < 0 || neighbor_cell.y >= cell_list.grid_size.y ||
+            neighbor_cell.z < 0 || neighbor_cell.z >= cell_list.grid_size.z) {
+          continue;
+        }
+
+        int cell_idx = cell_list.get_cell_index(
+            neighbor_cell.x, neighbor_cell.y, neighbor_cell.z);
+        int start = cell_list.cell_starts[cell_idx];
+        int end = cell_list.cell_ends[cell_idx];
+
+        // Check all particles in this cell
+        for (int idx = start; idx < end; idx++) {
+          int j = cell_list.sorted_particles[idx];
+
+          if (i >= j)
+            continue; // Avoid double counting and self-interaction
+
+          float4 pos_j = positions[j];
+          float3 other_pos = make_float3(pos_j.x, pos_j.y, pos_j.z);
+
+          // Calculate distance with periodic boundary conditions
+          float dx = my_pos.x - other_pos.x;
+          float dy = my_pos.y - other_pos.y;
+          float dz = my_pos.z - other_pos.z;
+
+          // Apply PBC
+          dx -= box_lengths->x * roundf(dx / box_lengths->x);
+          dy -= box_lengths->y * roundf(dy / box_lengths->y);
+          dz -= box_lengths->z * roundf(dz / box_lengths->z);
+
+          float r_squared = dx * dx + dy * dy + dz * dz;
+
+          if (r_squared <= cutoff_squared && neighbor_count < max_neighbors) {
+            my_neighbors[neighbor_count] = j;
+            neighbor_count++;
+          }
+        }
+      }
+    }
+  }
+
+  neighbor_list.neighbor_counts[i] = neighbor_count;
+}
+
+// =============================================================================
+// HOST FUNCTIONS
+// =============================================================================
+
+void build_cell_list(const float4 *positions, CellList &cell_list,
+                     const float3 *box_lengths) {
+
+  // Step 1: Assign particles to cells
+  auto config = get_launch_config(cell_list.n_particles);
+
+  assign_particles_to_cells_kernel<<<config.blocks, config.threads>>>(
+      positions, cell_list.particle_cells, cell_list, cell_list.n_particles);
+
+  // Step 2: Sort particles by cell index using Thrust
+  thrust::device_ptr<int> particle_cells_ptr(cell_list.particle_cells);
+  thrust::device_ptr<int> sorted_particles_ptr(cell_list.sorted_particles);
+
+  // Initialize particle indices 0, 1, 2, ...
+  thrust::sequence(sorted_particles_ptr,
+                   sorted_particles_ptr + cell_list.n_particles);
+
+  // Sort particle indices by their cell assignments
+  thrust::sort_by_key(particle_cells_ptr,
+                      particle_cells_ptr + cell_list.n_particles,
+                      sorted_particles_ptr);
+
+  // Step 3: Initialize cell boundaries to -1
+  cudaMemset(cell_list.cell_starts, -1, cell_list.total_cells * sizeof(int));
+  cudaMemset(cell_list.cell_ends, -1, cell_list.total_cells * sizeof(int));
+
+  // Step 4: Find cell boundaries
+  config = get_launch_config(cell_list.n_particles);
+  find_cell_boundaries_kernel<<<config.blocks, config.threads>>>(
+      cell_list.particle_cells, cell_list.cell_starts, cell_list.cell_ends,
+      cell_list.n_particles, cell_list.total_cells);
+}
+
+void build_neighbor_list(const float4 *positions, const CellList &cell_list,
+                         NeighborList &neighbor_list, float cutoff_squared,
+                         const float3 *box_lengths) {
+
+  // Initialize neighbor counts to 0
+  cudaMemset(neighbor_list.neighbor_counts, 0,
+             neighbor_list.n_particles * sizeof(int));
+
+  auto config = get_launch_config(neighbor_list.n_particles);
+
+  build_neighbor_lists_kernel<<<config.blocks, config.threads>>>(
+      positions, cell_list, neighbor_list, cutoff_squared, box_lengths,
+      neighbor_list.n_particles);
+}
+
+void build_neighbor_list_optimized(const float4 *positions, size_t n_particles,
+                                   const float3 *box_lengths, float cutoff,
+                                   NeighborList &neighbor_list) {
+
+  // Get box size for cell list construction
+  float3 box_size;
+  cudaMemcpy(&box_size, box_lengths, sizeof(float3), cudaMemcpyDeviceToHost);
+
+  // Create cell list
+  CellList cell_list(n_particles, box_size, cutoff);
+
+  // Build cell list
+  build_cell_list(positions, cell_list, box_lengths);
+
+  // Build neighbor list using cell list
+  build_neighbor_list(positions, cell_list, neighbor_list, cutoff * cutoff,
+                      box_lengths);
+}

kernels/neighbor_list.cuh

@@ -1,21 +1,68 @@
 #ifndef NEIGHBOR_LIST_CUH
 #define NEIGHBOR_LIST_CUH
 
-#include "box.hpp"
+#include "kernel_config.cuh" // From previous artifact
-#include "kernel_config.cuh"
-#include "utils.cuh"
 #include <cuda_runtime.h>
 #include <thrust/device_vector.h>
-#include <thrust/pair.h>
 #include <thrust/sequence.h>
 #include <thrust/sort.h>
 
+/**
+ * Simple kernel to initialize neighbor offsets
+ */
+__global__ void init_neighbor_offsets(int *offsets, size_t n_particles,
+                                      int max_neighbors) {
+  size_t i = get_thread_id();
+  if (i <= n_particles) { // Note: <= because we need n_particles + 1 elements
+    offsets[i] = i * max_neighbors;
+  }
+}
+
+/**
+ * Neighbor list data structure
+ * Uses a compact format similar to CSR sparse matrices
+ */
+struct NeighborList {
+  int *neighbor_offsets; // Size: n_particles + 1, offset into neighbor_indices
+  int *neighbor_indices; // Size: total_neighbors, actual neighbor particle IDs
+  int *neighbor_counts;  // Size: n_particles, number of neighbors per particle
+  size_t max_neighbors;  // Maximum neighbors allocated per particle
+  size_t n_particles;
+
+  // Constructor
+  NeighborList(size_t n_particles, size_t max_neighbors_per_particle)
+      : max_neighbors(max_neighbors_per_particle), n_particles(n_particles) {
+
+    cudaMalloc(&neighbor_offsets, (n_particles + 1) * sizeof(int));
+    cudaMalloc(&neighbor_indices, n_particles * max_neighbors * sizeof(int));
+    cudaMalloc(&neighbor_counts, n_particles * sizeof(int));
+
+    // Initialize offsets
+    auto kernel_config = get_launch_config(n_particles);
+    init_neighbor_offsets<<<kernel_config.blocks, kernel_config.threads>>>(
+        neighbor_offsets, n_particles, max_neighbors);
+  }
+
+  ~NeighborList() {
+    cudaFree(neighbor_offsets);
+    cudaFree(neighbor_indices);
+    cudaFree(neighbor_counts);
+  }
+
+  // Get neighbors for particle i (device function)
+  __device__ int *get_neighbors(int i) const {
+    return &neighbor_indices[neighbor_offsets[i]];
+  }
+
+  __device__ int get_neighbor_count(int i) const { return neighbor_counts[i]; }
+};
+
 /**
  * Cell list structure for spatial hashing
  */
 struct CellList {
   int *cell_starts;      // Size: total_cells, start index in sorted_particles
-  int *cell_count;       // Size: total_cells, end index in sorted_particles
+  int *cell_ends;        // Size: total_cells, end index in sorted_particles
   int *sorted_particles; // Size: n_particles, particle IDs sorted by cell
   int *particle_cells; // Size: n_particles, which cell each particle belongs to
 
@@ -25,55 +72,40 @@ struct CellList {
   size_t total_cells;
   size_t n_particles;
 
-  CellList(size_t n_particles, Box &box, float r_cutoff)
+  CellList(size_t n_particles, float3 box_size, float cutoff)
       : n_particles(n_particles) {
 
-    box_min.x = box.xlo;
+    // Calculate grid dimensions (cell size slightly larger than cutoff)
-    box_min.y = box.ylo;
+    float cell_margin = 1.001f; // Small safety margin
-    box_min.z = box.zlo;
+    cell_size = make_float3(cutoff * cell_margin, cutoff * cell_margin,
+                            cutoff * cell_margin);
 
-    auto [grid_size, cell_size] = calc_grid_and_cell_size(box, r_cutoff);
+    grid_size.x = (int)ceilf(box_size.x / cell_size.x);
-    this->grid_size = grid_size;
+    grid_size.y = (int)ceilf(box_size.y / cell_size.y);
-    this->cell_size = cell_size;
+    grid_size.z = (int)ceilf(box_size.z / cell_size.z);
 
-    total_cells = grid_size.x * grid_size.y * grid_size.z;
+    total_cells = (size_t)grid_size.x * grid_size.y * grid_size.z;
+    box_min = make_float3(0.0f, 0.0f, 0.0f); // Assume box starts at origin
 
+    // Allocate memory
     cudaMalloc(&cell_starts, total_cells * sizeof(int));
-    cudaMalloc(&cell_count, total_cells * sizeof(int));
+    cudaMalloc(&cell_ends, total_cells * sizeof(int));
     cudaMalloc(&sorted_particles, n_particles * sizeof(int));
     cudaMalloc(&particle_cells, n_particles * sizeof(int));
   }
 
   ~CellList() {
     cudaFree(cell_starts);
-    cudaFree(cell_count);
+    cudaFree(cell_ends);
     cudaFree(sorted_particles);
     cudaFree(particle_cells);
   }
 
   // Get cell index from 3D coordinates
-  // TODO; Maybe update this to use Morton Encodings in the future to improve
+  __device__ int get_cell_index(int x, int y, int z) const {
-  // locality of particle indices. Unclear how much of a benefit this will add,
-  // but would be cool to do
-  __host__ __device__ int get_cell_index(int x, int y, int z) const {
     return z * grid_size.x * grid_size.y + y * grid_size.x + x;
   }
 
-  std::pair<int3, float3> calc_grid_and_cell_size(Box &box,
-                                                  float r_cutoff) const {
-    int3 grid_size = {utils::max((int)(box.xhi - box.xlo) / r_cutoff, 1),
-                      utils::max((int)(box.yhi - box.ylo) / r_cutoff, 1),
-                      utils::max((int)(box.zhi - box.zlo) / r_cutoff, 1)};
-
-    float3 cell_size = {
-        (box.xhi - box.xlo) / grid_size.x,
-        (box.yhi - box.ylo) / grid_size.y,
-        (box.zhi - box.zlo) / grid_size.z,
-    };
-
-    return std::make_pair(grid_size, cell_size);
-  }
-
   // Get cell coordinates from position
   __device__ int3 get_cell_coords(float3 pos) const {
     return make_int3((int)((pos.x - box_min.x) / cell_size.x),
@@ -82,4 +114,19 @@ struct CellList {
   }
 };
 
+// Forward declarations
+void build_cell_list(const float4 *positions, CellList &cell_list,
+                     const float3 *box_lengths);
+
+void build_neighbor_list(const float4 *positions, const CellList &cell_list,
+                         NeighborList &neighbor_list, float cutoff_squared,
+                         const float3 *box_lengths);
+
+/**
+ * High-level interface - builds complete neighbor list
+ */
+void build_neighbor_list_optimized(const float4 *positions, size_t n_particles,
+                                   const float3 *box_lengths, float cutoff,
+                                   NeighborList &neighbor_list);
+
 #endif

kernels/utils.cuh

@@ -1,8 +0,0 @@
-#ifndef UTILS_CUH
-#define UTILS_CUH
-
-namespace utils {
-__device__ __host__ inline int max(int a, int b) { return (a > b) ? a : b; }
-} // namespace utils
-
-#endif

src/box.hpp

@@ -18,16 +18,6 @@ struct Box {
   bool x_is_periodic;
   bool y_is_periodic;
   bool z_is_periodic;
-
-  Box(real xlo, real xhi, real ylo, real yhi, real zlo, real zhi,
-      bool x_is_periodic, bool y_is_periodic, bool z_is_periodic)
-      : xlo(xlo), xhi(xhi), ylo(ylo), yhi(yhi), zlo(zlo), zhi(zhi),
-        x_is_periodic(x_is_periodic), y_is_periodic(y_is_periodic),
-        z_is_periodic(z_is_periodic) {}
-
-  Box(real xlo, real xhi, real ylo, real yhi, real zlo, real zhi)
-      : xlo(xlo), xhi(xhi), ylo(ylo), yhi(yhi), zlo(zlo), zhi(zhi),
-        x_is_periodic(true), y_is_periodic(true), z_is_periodic(true) {}
 };
 
 #endif

tests/cuda_unit_tests/CMakeLists.txt

@@ -4,7 +4,6 @@ add_executable(${NAME}_cuda_tests
     test_potential.cu
     test_forces.cu
     test_kernel_config.cu
-    test_neighbor_list.cu
 )
 
 target_link_libraries(${NAME}_cuda_tests gtest gtest_main)

tests/cuda_unit_tests/test_neighbor_list.cu

@@ -1,31 +0,0 @@
-
-#include "box.hpp"
-#include "neighbor_list.cuh"
-#include "gtest/gtest.h"
-
-TEST(CellListTest, Constructor) {
-  // Test case 1: Simple case
-  Box box(0, 10, 0, 10, 0, 10);
-  float cutoff = 2.5;
-  CellList cell_list(100, box, cutoff);
-
-  EXPECT_EQ(cell_list.grid_size.x, 4);
-  EXPECT_EQ(cell_list.grid_size.y, 4);
-  EXPECT_EQ(cell_list.grid_size.z, 4);
-  EXPECT_FLOAT_EQ(cell_list.cell_size.x, 2.5);
-  EXPECT_FLOAT_EQ(cell_list.cell_size.y, 2.5);
-  EXPECT_FLOAT_EQ(cell_list.cell_size.z, 2.5);
-  EXPECT_EQ(cell_list.total_cells, 64);
-}
-
-TEST(CellListTest, GetCellIndex) {
-  Box box(0, 10, 0, 10, 0, 10);
-  float cutoff = 2.5;
-  CellList cell_list(100, box, cutoff);
-
-  int x = 1, y = 2, z = 3;
-  int expected_index = z * cell_list.grid_size.x * cell_list.grid_size.y +
-                       y * cell_list.grid_size.x + x;
-
-  EXPECT_EQ(cell_list.get_cell_index(x, y, z), expected_index);
-}