diff --git a/src/grid/gpu/grid_gpu_collocate.cu b/src/grid/gpu/grid_gpu_collocate.cu index 9218e1ea6c..63949052ca 100644 --- a/src/grid/gpu/grid_gpu_collocate.cu +++ b/src/grid/gpu/grid_gpu_collocate.cu @@ -99,7 +99,8 @@ __launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) { continue; fill_smem_task_coef(dev_, task_id, task); - T *__restrict__ coef_ = &dev_.ptr_dev[2][dev_.tasks[task_id].coef_offset]; + T *__restrict__ coef_ = + &dev_.buffers_dev.coef[dev_.tasks[task_id].coef_offset]; compute_alpha(task, smem_alpha); @@ -160,8 +161,9 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) { // Alloc shared memory. extern __shared__ T coefs_[]; - T *coef_ = - &dev_.ptr_dev[2][dev_.tasks[dev_.first_task + block_index()].coef_offset]; + const size_t coef_offset = + dev_.tasks[dev_.first_task + block_index()].coef_offset; + T *coef_ = &dev_.buffers_dev.coef[coef_offset]; __shared__ T dh_[9]; if (tid < 9) { @@ -387,7 +389,7 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) { } res *= exp(-(r3x2 + r3y2 + r3z2) * task.zetp); - atomicAdd(dev_.ptr_dev[1] + + atomicAdd(dev_.buffers_dev.grid + (z2 * dev_.grid_local_size_.y + y2) * dev_.grid_local_size_.x + x2, diff --git a/src/grid/gpu/grid_gpu_context.cu b/src/grid/gpu/grid_gpu_context.cu index 70733164f9..85d31f4769 100644 --- a/src/grid/gpu/grid_gpu_context.cu +++ b/src/grid/gpu/grid_gpu_context.cu @@ -62,10 +62,10 @@ context_info::set_kernel_parameters(const int level, params.la_max_diff = smem_params.ldiffs().la_max_diff; params.lb_max_diff = smem_params.ldiffs().lb_max_diff; - params.ptr_dev[0] = pab_block_.data(); + params.buffers_dev.pab_block = pab_block_.data(); if (level >= 0) { - params.ptr_dev[1] = grid_[level].data(); + params.buffers_dev.grid = grid_[level].data(); memcpy(params.dh_, grid_[level].dh(), 9 * sizeof(double)); memcpy(params.dh_inv_, grid_[level].dh_inv(), 9 * sizeof(double)); params.first_task = first_task_per_level_[level]; @@ -76,11 +76,11 @@ context_info::set_kernel_parameters(const int level, params.grid_border_width_ = grid_[level].border_width(); } - params.ptr_dev[2] = this->coef_dev_.data(); - params.ptr_dev[3] = hab_block_.data(); - params.ptr_dev[4] = forces_.data(); - params.ptr_dev[5] = virial_.data(); - params.ptr_dev[6] = this->cab_dev_.data(); + params.buffers_dev.coef = this->coef_dev_.data(); + params.buffers_dev.hab_block = hab_block_.data(); + params.buffers_dev.forces = forces_.data(); + params.buffers_dev.virial = virial_.data(); + params.buffers_dev.cab = this->cab_dev_.data(); params.cab_block_offset_dev = this->cab_block_offset_dev.data(); params.sphi_dev = this->sphi_dev.data(); return params; diff --git a/src/grid/gpu/grid_gpu_context.h b/src/grid/gpu/grid_gpu_context.h index d7b51eb07c..0d77428c5d 100644 --- a/src/grid/gpu/grid_gpu_context.h +++ b/src/grid/gpu/grid_gpu_context.h @@ -414,6 +414,20 @@ struct task_info { bool skip_task; }; +/******************************************************************************* + * \brief Device buffers used by the collocate and integrate kernels. + * Which fields are non-null depends on which kernel is launched. + ******************************************************************************/ +struct kernel_buffers { + double *pab_block{nullptr}; + double *grid{nullptr}; + double *coef{nullptr}; + double *hab_block{nullptr}; + double *forces{nullptr}; + double *virial{nullptr}; + double *cab{nullptr}; +}; + /******************************************************************************* * \brief Parameters of the collocate kernel. ******************************************************************************/ @@ -435,8 +449,7 @@ struct kernel_params { char la_max_diff{0}; char lb_max_diff{0}; enum grid_func func; - double *ptr_dev[7] = {nullptr, nullptr, nullptr, nullptr, - nullptr, nullptr, nullptr}; + kernel_buffers buffers_dev; double **sphi_dev{nullptr}; int ntasks{0}; int *task_sorted_by_blocks_dev{nullptr}; diff --git a/src/grid/gpu/grid_gpu_integrate.cu b/src/grid/gpu/grid_gpu_integrate.cu index 55daaf2b74..e8dd6c7486 100644 --- a/src/grid/gpu/grid_gpu_integrate.cu +++ b/src/grid/gpu/grid_gpu_integrate.cu @@ -90,7 +90,7 @@ __global__ __launch_bounds__(64) void compute_hab(const kernel_params dev_) { fill_smem_task_coef(dev_, task_id, task); T *__restrict__ coef_ = reinterpret_cast(__builtin_assume_aligned( - &dev_.ptr_dev[2][dev_.tasks[task_id].coef_offset], 32)); + &dev_.buffers_dev.coef[dev_.tasks[task_id].coef_offset], 32)); __syncthreads(); compute_alpha(task, smem_alpha); @@ -166,7 +166,8 @@ __launch_bounds__(64) void compute_hab_forces(const kernel_params dev_) { continue; fill_smem_task_coef(dev_, task_id, task); - T *__restrict__ coef_ = &dev_.ptr_dev[2][dev_.tasks[task_id].coef_offset]; + T *__restrict__ coef_ = + &dev_.buffers_dev.coef[dev_.tasks[task_id].coef_offset]; __syncthreads(); compute_alpha(task, smem_alpha); __syncthreads(); @@ -215,7 +216,7 @@ __launch_bounds__(64) void compute_hab_forces(const kernel_params dev_) { get_force_b(a, b, 2, task.zeta, task.zetb, task.rab, task.n1, smem_cab); - if (dev_.ptr_dev[5] != nullptr) { + if (dev_.buffers_dev.virial != nullptr) { virial[0] += sphia_times_sphib * (get_virial_a(a, b, 0, 0, task.zeta, @@ -300,17 +301,17 @@ __launch_bounds__(64) void compute_hab_forces(const kernel_params dev_) { const auto &glb_task = dev_.tasks[task_id]; const int iatom = glb_task.iatom; const int jatom = glb_task.jatom; - T *forces_a = &dev_.ptr_dev[4][3 * iatom]; - T *forces_b = &dev_.ptr_dev[4][3 * jatom]; + T *forces_a = &dev_.buffers_dev.forces[3 * iatom]; + T *forces_b = &dev_.buffers_dev.forces[3 * jatom]; T *sum = (T *)shared_memory; - if (dev_.ptr_dev[5] != nullptr) { + if (dev_.buffers_dev.virial != nullptr) { for (int i = 0; i < 9; i++) { virial[i] = block_reduce_64(sum, virial[i], tid); if (tid == 0) - atomicAdd(dev_.ptr_dev[5] + i, virial[i]); + atomicAdd(dev_.buffers_dev.virial + i, virial[i]); } } @@ -354,7 +355,7 @@ specialized to the integration. ******************************************************************************/ template + int lbatch = 20> __global__ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) { if (dev_.tasks[dev_.first_task + block_index()].skip_task) @@ -482,10 +483,10 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) { // the register is actually needed for computation. This is true on // NVIDIA hardware - T grid_value = - __ldg(&dev_.ptr_dev[1][(z2 * dev_.grid_local_size_.y + y2) * - dev_.grid_local_size_.x + - x2]); + const int grid_index = + (z2 * dev_.grid_local_size_.y + y2) * dev_.grid_local_size_.x + + x2; + T grid_value = __ldg(&dev_.buffers_dev.grid[grid_index]); const T r3xy = r3.x * r3.y; const T r3xz = r3.x * r3.z; @@ -673,12 +674,20 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) { if (tid == 0) accumulator[i][0] = val; } +#if defined(__CUDACC__) __syncwarp(); +#endif } - if (tid < min(length - ico, lbatch)) - dev_.ptr_dev[2][dev_.tasks[dev_.first_task + block_index()].coef_offset + - tid + ico] = accumulator[tid][0]; +#if !defined(__CUDACC__) + __syncthreads(); +#endif + + if (tid < min(length - ico, lbatch)) { + const size_t coef_offset = + dev_.tasks[dev_.first_task + block_index()].coef_offset; + dev_.buffers_dev.coef[coef_offset + tid + ico] = accumulator[tid][0]; + } __syncthreads(); } } diff --git a/src/grid/gpu/grid_gpu_internal_header.h b/src/grid/gpu/grid_gpu_internal_header.h index 43e54e35c9..9664efca10 100644 --- a/src/grid/gpu/grid_gpu_internal_header.h +++ b/src/grid/gpu/grid_gpu_internal_header.h @@ -826,13 +826,15 @@ __device__ __inline__ void fill_smem_task_coef(const kernel_params &dev, // Locate current matrix block within the buffer. const int block_offset = dev.block_offsets[glb_task.block_num]; task.block_transposed = glb_task.block_transposed; - task.pab_block = dev.ptr_dev[0] + block_offset + glb_task.subblock_offset; + task.pab_block = + dev.buffers_dev.pab_block + block_offset + glb_task.subblock_offset; - if (dev.ptr_dev[3] != nullptr) { - task.hab_block = dev.ptr_dev[3] + block_offset + glb_task.subblock_offset; - if (dev.ptr_dev[4] != nullptr) { - task.forces_a = &dev.ptr_dev[4][3 * iatom]; - task.forces_b = &dev.ptr_dev[4][3 * jatom]; + if (dev.buffers_dev.hab_block != nullptr) { + task.hab_block = + dev.buffers_dev.hab_block + block_offset + glb_task.subblock_offset; + if (dev.buffers_dev.forces != nullptr) { + task.forces_a = &dev.buffers_dev.forces[3 * iatom]; + task.forces_b = &dev.buffers_dev.forces[3 * jatom]; } } } @@ -896,7 +898,7 @@ public: template __inline__ __device__ T *allocate_workspace(const kernel_params &dev_) { unsigned int offset = dev_.cab_block_offset_dev[block_index()]; - return (T *)(dev_.ptr_dev[6] + offset); + return (T *)(dev_.buffers_dev.cab + offset); } template