From cd4a00ebed75f1c837c1fef1ad5d20bd85e21090 Mon Sep 17 00:00:00 2001 From: Alex Selimov Date: Thu, 25 Sep 2025 20:06:29 -0400 Subject: [PATCH] Refactor type arguments to compute cell_index --- kernels/neighbor_list.cuh | 6 ++++-- tests/cuda_unit_tests/test_neighbor_list.cu | 3 ++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/kernels/neighbor_list.cuh b/kernels/neighbor_list.cuh index 73039be..14d1201 100644 --- a/kernels/neighbor_list.cuh +++ b/kernels/neighbor_list.cuh @@ -55,8 +55,10 @@ struct CellList { // 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(int x, int y, int z) const { - return z * grid_size.x * grid_size.y + y * grid_size.x + x; + __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 calc_grid_and_cell_size(Box &box, diff --git a/tests/cuda_unit_tests/test_neighbor_list.cu b/tests/cuda_unit_tests/test_neighbor_list.cu index 9818075..ad44eab 100644 --- a/tests/cuda_unit_tests/test_neighbor_list.cu +++ b/tests/cuda_unit_tests/test_neighbor_list.cu @@ -27,5 +27,6 @@ TEST(CellListTest, GetCellIndex) { 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); + EXPECT_EQ(cell_list.get_cell_index_from_cell_coords({x, y, z}), + expected_index); }