Finalize cell list code and add tests
Some checks failed
Build and Test / build-and-test (push) Failing after 5m3s

This commit is contained in:
Alex Selimov 2025-09-26 19:30:12 -04:00
parent 0eeb9d305e
commit 6927026f87
Signed by: aselimov
GPG key ID: 3DDB9C3E023F1F31
3 changed files with 97 additions and 22 deletions

View file

@ -53,16 +53,6 @@ struct CellList {
cudaFree(particle_cells);
}
// Get cell index from 3D coordinates
// TODO; Maybe update this to use Morton Encodings in the future to improve
// 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_from_cell_coords(int3 cell_coords) const {
return cell_coords.z * grid_size.x * grid_size.y +
cell_coords.y * grid_size.x + cell_coords.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),
@ -78,17 +68,31 @@ struct CellList {
return std::make_pair(grid_size, cell_size);
}
__host__ __device__ int3 get_cell_coords_from_position(float3 pos) const {
// Get cell index from 3D coordinates
// TODO; Maybe update this to use Morton Encodings in the future to improve
// locality of particle indices. Unclear how much of a benefit this will add,
// but would be cool to do
__host__ __device__ static int
get_cell_index_from_cell_coords(int3 cell_coords, int3 grid_size) {
return cell_coords.z * grid_size.x * grid_size.y +
cell_coords.y * grid_size.x + cell_coords.x;
}
__host__ __device__ static int3
get_cell_coords_from_position(float3 pos, float3 box_min, float3 cell_size) {
return make_int3((int)((pos.x - box_min.x) / cell_size.x),
(int)((pos.y - box_min.y) / cell_size.y),
(int)((pos.z - box_min.z) / cell_size.z));
}
__host__ __device__ int get_cell_index_from_position(float3 pos) const {
return get_cell_index_from_cell_coords(get_cell_coords_from_position(pos));
__host__ __device__ static int
get_cell_index_from_position(float3 pos, int3 grid_size, float3 box_min,
float3 cell_size) {
return get_cell_index_from_cell_coords(
get_cell_coords_from_position(pos, box_min, cell_size), grid_size);
}
void assign_particles_to_cells(float3 *positions) {
void assign_particles_to_cells(float3 *d_positions) {
thrust::device_ptr<int> particle_cells_ptr(particle_cells);
thrust::device_ptr<int> sorted_particles_ptr(sorted_particles);
thrust::device_ptr<int> cell_starts_ptr(cell_starts);
@ -96,9 +100,16 @@ struct CellList {
thrust::sequence(sorted_particles_ptr, sorted_particles_ptr + n_particles);
for (size_t i = 0; i < n_particles; i++) {
particle_cells[i] = get_cell_index_from_position(positions[i]);
}
int3 grid_size = this->grid_size;
float3 box_min = this->box_min;
float3 cell_size = this->cell_size;
thrust::transform(d_positions, d_positions + n_particles,
particle_cells_ptr,
[grid_size, box_min, cell_size] __device__(auto pos) {
return CellList::get_cell_index_from_position(
pos, grid_size, box_min, cell_size);
});
thrust::sort_by_key(particle_cells_ptr, particle_cells_ptr + n_particles,
sorted_particles_ptr);