comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-21 06:25:15 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make the grid gpu code less opaque (#5540)

Browse source 
Co-authored-by: Mathieu Taillefumier <mathieu.taillefumier@free.fr>
This commit is contained in:
Taillefumier Mathieu 2026-07-08 13:29:50 +00:00 • committed by GitHub
parent 39e6d08f13
commit 7909f53bfa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 61 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/grid/gpu/grid_gpu_collocate.cu
View file

@ -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,

14
src/grid/gpu/grid_gpu_context.cu
View file

@ -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;

17
src/grid/gpu/grid_gpu_context.h
View file

@ -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};

39
src/grid/gpu/grid_gpu_integrate.cu
View file

@ -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<T *>(__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<COMPUTE_TAU, T>(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<COMPUTE_TAU, T>(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<T>(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 <typename T, typename T3, bool distributed__, bool orthogonal_,
int lbatch = 10>
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();
}
}

16
src/grid/gpu/grid_gpu_internal_header.h
View file

@ -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 <typename T>
__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 <typename T, typename T3>