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

grid: Clean up GPU backend and fix race condition (#5338)

Browse source 
Co-authored-by: Mathieu Taillefumier <mathieu.taillefumier@free.fr>
This commit is contained in:
Taillefumier Mathieu 2026-06-22 09:58:47 +00:00 • committed by GitHub
parent e6ecf3375d
commit efe242d354
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 271 additions and 347 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -40,8 +40,7 @@ __device__ __inline__ void block_to_cab(const kernel_params &params,
// This is a T matrix product. Since the pab block can be quite large the
// two products are fused to conserve shared memory.
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int tid = thread_global_index();
for (int jco = task.first_cosetb + tid / 16; jco < task.ncosetb; jco += 4) {
for (int ico = task.first_coseta + (tid % 16); ico < task.ncoseta;
@ -82,15 +81,14 @@ __global__
__launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
// Copy task from global to shared memory and precompute some stuff.
extern __shared__ T shared_memory[];
const int number_of_tasks = dev_.num_tasks_per_block_dev[blockIdx.x];
const int number_of_tasks = dev_.num_tasks_per_block_dev[block_index()];
if (number_of_tasks == 0)
return;
T *smem_alpha = &shared_memory[0];
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int offset = dev_.sorted_blocks_offset_dev[blockIdx.x];
const int tid = thread_global_index();
const int offset = dev_.sorted_blocks_offset_dev[block_index()];
T *smem_cab = allocate_workspace<T>(dev_);
for (int tk = 0; tk < number_of_tasks; tk++) {
@ -107,8 +105,10 @@ __launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
z += blockDim.x * blockDim.y * blockDim.z)
smem_cab[z] = 0.0;
__syncthreads();
block_to_cab<T, IS_FUNC_AB>(dev_, task, smem_cab);
__syncthreads();
cab_to_cxyz(task, smem_alpha, smem_cab, coef_);
}
}
@ -141,25 +141,24 @@ __launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
calculate_coefficients. We only keep the non zero elements to same memory.
*/
template <typename T, typename T3, bool distributed__, bool orthorhombic_>
template <typename T, typename T3, bool distributed__, bool orthogonal_>
__global__
__launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
// Copy task from global to shared memory and precompute some stuff.
__shared__ smem_task_reduced<T, T3> task;
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
if (dev_.tasks[dev_.first_task + block_index()].skip_task)
return;
fill_smem_task_reduced(dev_, dev_.first_task + blockIdx.x, task);
fill_smem_task_reduced(dev_, dev_.first_task + block_index(), task);
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int tid = thread_global_index();
// Alloc shared memory.
extern __shared__ T coefs_[];
T *coef_ =
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset];
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + block_index()].coef_offset];
__shared__ T dh_[9], dh_inv_[9];
if (tid < 9) {
@ -176,26 +175,12 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
// the cube center is initialy expressed in lattice coordinates but we
// always do something like this. x = x + lower_corner + cube_center (+
// roffset) - grid_lower_corner so shift the cube center already
task.cube_center.z += task.lb_cube.z - dev_.grid_lower_corner_[0];
task.cube_center.y += task.lb_cube.y - dev_.grid_lower_corner_[1];
task.cube_center.x += task.lb_cube.x - dev_.grid_lower_corner_[2];
if (distributed__) {
if (task.apply_border_mask) {
compute_window_size(
dev_.grid_local_size_,
dev_.tasks[dev_.first_task + blockIdx.x].border_mask,
dev_.grid_border_width_, &task.window_size, &task.window_shift);
}
}
setup_task_cube_center<T, T3, distributed__>(dev_, task);
}
__syncthreads();
for (int z = threadIdx.z; z < task.cube_size.z; z += blockDim.z) {
int z2 = (z + task.cube_center.z) % dev_.grid_full_size_[0];
if (z2 < 0)
z2 += dev_.grid_full_size_[0];
int z2 = wrap_grid_index(z + task.cube_center.z, dev_.grid_full_size_.z);
if (distributed__) {
// check if the point is within the window
@ -211,23 +196,15 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
// compute the coordinates of the point in atomic coordinates
T kremain;
short int ymin = 0;
short int ymax = task.cube_size.y - 1;
int ymin = 0;
int ymax = task.cube_size.y - 1;
if (orthorhombic_ && !task.apply_border_mask) {
ymin = (2 * (z + task.lb_cube.z) - 1) / 2;
ymin *= ymin;
kremain = task.discrete_radius * task.discrete_radius -
((T)ymin) * dh_[8] * dh_[8];
ymin = ceil(-1.0e-8 - sqrt(fmax(0.0, kremain)) * dh_inv_[4]);
ymax = 1 - ymin - task.lb_cube.y;
ymin = ymin - task.lb_cube.y;
if (orthogonal_ && !task.apply_border_mask) {
kremain = calculate_ymix_ymax_boundaries(task, z, ymin, ymax);
}
for (int y = ymin + threadIdx.y; y <= ymax; y += blockDim.y) {
int y2 = (y + task.cube_center.y) % dev_.grid_full_size_[1];
if (y2 < 0)
y2 += dev_.grid_full_size_[1];
int y2 = wrap_grid_index(y + task.cube_center.y, dev_.grid_full_size_.y);
if (distributed__) {
if (task.apply_border_mask) {
@ -238,23 +215,15 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
}
}
short int xmin = 0;
short int xmax = task.cube_size.x - 1;
if (orthorhombic_ && !task.apply_border_mask) {
xmin = (2 * (y + task.lb_cube.y) - 1) / 2;
xmin *= xmin;
xmin =
ceil(-1.0e-8 - sqrt(fmax(0.0, kremain - xmin * dh_[4] * dh_[4])) *
dh_inv_[0]);
xmax = 1 - xmin - task.lb_cube.x;
xmin = xmin - task.lb_cube.x;
int xmin = 0;
int xmax = task.cube_size.x - 1;
if (orthogonal_ && !task.apply_border_mask) {
calculate_xmin_xmax_boundaries<T, T3>(task, y, kremain, xmin, xmax);
}
for (int x = xmin + threadIdx.x; x <= xmax; x += blockDim.x) {
int x2 = (x + task.cube_center.x) % dev_.grid_full_size_[2];
if (x2 < 0)
x2 += dev_.grid_full_size_[2];
int x2 =
wrap_grid_index(x + task.cube_center.x, dev_.grid_full_size_.x);
if (distributed__) {
if (task.apply_border_mask) {
@ -266,19 +235,13 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
}
}
// I make no distinction between orthorhombic and non orthorhombic
// I make no distinction between orthogonal and non orthogonal
// cases
T3 r3;
if (orthorhombic_) {
r3.x = (x + task.lb_cube.x + task.roffset.x) * dh_[0];
r3.y = (y + task.lb_cube.y + task.roffset.y) * dh_[4];
r3.z = (z + task.lb_cube.z + task.roffset.z) * dh_[8];
} else {
r3 = compute_coordinates(dh_, (x + task.lb_cube.x + task.roffset.x),
(y + task.lb_cube.y + task.roffset.y),
(z + task.lb_cube.z + task.roffset.z));
}
r3 = compute_coordinates(dh_, (x + task.lb_cube.x + task.roffset.x),
(y + task.lb_cube.y + task.roffset.y),
(z + task.lb_cube.z + task.roffset.z));
const T r3x2 = r3.x * r3.x;
const T r3y2 = r3.y * r3.y;
@ -290,11 +253,11 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
// the region of interest.
if (((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)) &&
(!orthorhombic_ || task.apply_border_mask))
(!orthogonal_ || task.apply_border_mask))
continue;
} else {
// we do not need to do this test for the orthorhombic case
if ((!orthorhombic_) &&
if ((!orthogonal_) &&
((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)))
continue;
}
@ -424,8 +387,8 @@ __launch_bounds__(64) void collocate_kernel(const kernel_params dev_) {
res *= exp(-(r3x2 + r3y2 + r3z2) * task.zetp);
atomicAdd(dev_.ptr_dev[1] +
(z2 * dev_.grid_local_size_[1] + y2) *
dev_.grid_local_size_[2] +
(z2 * dev_.grid_local_size_.y + y2) *
dev_.grid_local_size_.x +
x2,
res);
}
@ -484,7 +447,7 @@ void context_info::collocate_one_grid_level(const int level,
const dim3 threads_per_block(4, 4, 4);
if (grid_[level].is_distributed()) {
if (grid_[level].is_orthorhombic())
if (grid_[level].is_orthogonal())
collocate_kernel<double, double3, true, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
@ -493,7 +456,7 @@ void context_info::collocate_one_grid_level(const int level,
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
} else {
if (grid_[level].is_orthorhombic())
if (grid_[level].is_orthogonal())
collocate_kernel<double, double3, false, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);

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

@ -61,16 +61,16 @@ context_info::set_kernel_parameters(const int level,
if (level >= 0) {
params.ptr_dev[1] = grid_[level].data();
for (int i = 0; i < 3; i++) {
memcpy(params.dh_, grid_[level].dh(), 9 * sizeof(double));
memcpy(params.dh_inv_, grid_[level].dh_inv(), 9 * sizeof(double));
params.grid_full_size_[i] = grid_[level].full_size(i);
params.grid_local_size_[i] = grid_[level].local_size(i);
params.grid_lower_corner_[i] = grid_[level].lower_corner(i);
params.grid_border_width_[i] = grid_[level].border_width(i);
}
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];
params.grid_full_size_ = grid_[level].full_size();
params.grid_local_size_ = grid_[level].local_size();
params.grid_lower_corner_ = grid_[level].lower_corner();
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();
@ -132,7 +132,7 @@ extern "C" void grid_gpu_create_task_list(
shift_local + 3 * level, border_width + 3 * level);
ctx->grid_[level].is_distributed(false);
ctx->grid_[level].set_lattice_vectors(&dh[9 * level], &dh_inv[9 * level]);
ctx->grid_[level].check_orthorhombicity(ortho);
ctx->grid_[level].check_orthogonality(ortho);
for (int i = 0; i < 9; i++)
dh_max[level] = std::max(dh_max[level], std::abs(dh[9 * level + i]));
}
@ -283,7 +283,7 @@ extern "C" void grid_gpu_create_task_list(
tasks_host[i].apply_border_mask = (tasks_host[i].border_mask != 0);
if (grid.is_orthorhombic() && (tasks_host[i].border_mask == 0)) {
if (grid.is_orthogonal() && (tasks_host[i].border_mask == 0)) {
tasks_host[i].discrete_radius =
rocm_backend::compute_cube_properties<double, double3, true>(
tasks_host[i].radius, grid.dh(), grid.dh_inv(),
@ -527,7 +527,7 @@ extern "C" void grid_gpu_collocate_task_list(const grid_gpu_task_list *ptr,
for (int has_border_mask = 0; has_border_mask <= 1; has_border_mask++) {
for (int lp = 0; lp < 20; lp++) {
const int count = ctx->stats[has_border_mask][lp];
if (ctx->grid_[0].is_orthorhombic() && !has_border_mask) {
if (ctx->grid_[0].is_orthogonal() && !has_border_mask) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_COLLOCATE_ORTHO, count);
} else {
@ -613,7 +613,7 @@ extern "C" void grid_gpu_integrate_task_list(
for (int has_border_mask = 0; has_border_mask <= 1; has_border_mask++) {
for (int lp = 0; lp < 20; lp++) {
const int count = ctx->stats[has_border_mask][lp];
if (ctx->grid_[0].is_orthorhombic() && !has_border_mask) {
if (ctx->grid_[0].is_orthogonal() && !has_border_mask) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_INTEGRATE_ORTHO, count);
} else {

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

@ -215,14 +215,14 @@ public:
template <typename T> class grid_info {
private:
std::array<int, 3> full_size_;
std::array<int, 3> local_size_;
int3 full_size_;
int3 local_size_;
// origin of the local part of the grid in grid point
std::array<int, 3> lower_corner_;
std::array<int, 3> border_width_;
int3 lower_corner_;
int3 border_width_;
std::array<T, 9> dh_;
std::array<T, 9> dh_inv_;
bool orthorhombic_{false};
bool orthogonal_{false};
bool is_distributed_{false};
gpu_vector<T> grid_;
@ -279,9 +279,10 @@ public:
is_distributed_ = distributed__;
}
void check_orthorhombicity(const bool ortho) {
/// Check if the lattice vectors form a orthogonal basis
void check_orthogonality(const bool ortho) {
if (ortho) {
orthorhombic_ = true;
orthogonal_ = true;
return;
}
T norm1, norm2, norm3;
@ -307,7 +308,7 @@ public:
((fabs(dh_[0] * dh_[3] + dh_[1] * dh_[4] + dh_[2] * dh_[5]) * norm1 *
norm2) < 1e-12);
orthorhombic_ = orthogonal[0] && orthogonal[1] && orthogonal[2];
orthogonal_ = orthogonal[0] && orthogonal[1] && orthogonal[2];
}
inline void copy_to_host(T *data__, offloadStream_t &stream) {
@ -327,25 +328,13 @@ public:
}
inline bool is_distributed() { return is_distributed_; }
inline int full_size(const int i) {
assert(i < 3);
return full_size_[i];
}
inline int3 &full_size() { return full_size_; }
inline int local_size(const int i) {
assert(i < 3);
return local_size_[i];
}
inline int3 &local_size() { return local_size_; }
inline int lower_corner(const int i) {
assert(i < 3);
return lower_corner_[i];
}
inline int3 &lower_corner() { return lower_corner_; }
inline int border_width(const int i) {
assert(i < 3);
return border_width_[i];
}
inline int3 &border_width() { return border_width_; }
inline T *data() { return grid_.data(); }
inline const T *data() const { return grid_.data(); }
@ -359,35 +348,34 @@ public:
inline T *dh_inv() { return dh_inv_.data(); }
inline const T *dh_inv() const { return dh_inv_.data(); }
inline bool is_orthorhombic() const { return orthorhombic_; }
inline bool is_orthogonal() const { return orthogonal_; }
inline bool is_distributed() const { return is_distributed_; }
private:
void initialize(const int *const full_size__, const int *const local_size__,
const int *const roffset__, const int *const border_width__) {
// the calling code store things like this cube[z][y][x] (in fortran
// cube(x,y,z)) so all sizes are [x,y,z] while we are working in C/C++ so we
// have to permute the indices to get this right.
// cube(x,y,z)).
full_size_[2] = full_size__[0];
full_size_[1] = full_size__[1];
full_size_[0] = full_size__[2];
full_size_.x = full_size__[0];
full_size_.y = full_size__[1];
full_size_.z = full_size__[2];
local_size_[2] = local_size__[0];
local_size_[1] = local_size__[1];
local_size_[0] = local_size__[2];
local_size_.x = local_size__[0];
local_size_.y = local_size__[1];
local_size_.z = local_size__[2];
lower_corner_[0] = roffset__[2];
lower_corner_[1] = roffset__[1];
lower_corner_[2] = roffset__[0];
lower_corner_.x = roffset__[0];
lower_corner_.y = roffset__[1];
lower_corner_.z = roffset__[2];
is_distributed_ = (full_size_[2] != local_size_[2]) ||
(full_size_[1] != local_size_[1]) ||
(full_size_[0] != local_size_[0]);
is_distributed_ = (full_size_.x != local_size_.x) ||
(full_size_.y != local_size_.y) ||
(full_size_.z != local_size_.z);
border_width_[2] = border_width__[0];
border_width_[1] = border_width__[1];
border_width_[0] = border_width__[2];
border_width_.x = border_width__[0];
border_width_.y = border_width__[1];
border_width_.z = border_width__[2];
}
};
@ -434,10 +422,10 @@ struct kernel_params {
int first_task{0};
// max size of cab.
int cab_size_{0};
int grid_full_size_[3] = {0, 0, 0};
int grid_local_size_[3] = {0, 0, 0};
int grid_lower_corner_[3] = {0, 0, 0};
int grid_border_width_[3] = {0, 0, 0};
int3 grid_full_size_ = make_int3(0, 0, 0);
int3 grid_local_size_ = make_int3(0, 0, 0);
int3 grid_lower_corner_ = make_int3(0, 0, 0);
int3 grid_border_width_ = make_int3(0, 0, 0);
double dh_[9];
double dh_inv_[9];
task_info *tasks;

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

@ -31,71 +31,33 @@
namespace rocm_backend {
// do a warp reduction and return the final sum to thread_id = 0
// do a block reduction for a block of size 64 and return the final sum to
// thread_id = 0
template <typename T>
__device__ __inline__ T warp_reduce(T *table, const int tid) {
__device__ __inline__ T block_reduce_64(T *table, const T val_, const int tid) {
// AMD GPU have warp size of 64 while nvidia GPUs have warpSize of 32 so the
// first step is common to both platforms.
table[tid] = val_;
__syncthreads();
T val = 0.0;
if (tid < 32) {
table[tid] += table[tid + 32];
}
__syncthreads();
if (tid < 16) {
table[tid] += table[tid + 16];
}
__syncthreads();
if (tid < 8) {
table[tid] += table[tid + 8];
}
__syncthreads();
if (tid < 4) {
table[tid] += table[tid + 4];
}
__syncthreads();
if (tid < 2) {
table[tid] += table[tid + 2];
}
__syncthreads();
return (table[0] + table[1]);
}
// #endif
val = table[tid] + table[tid + 32];
template <typename T>
__device__ __inline__ T reduce_two_warps(const T in, const int tid) {
int lane = tid & 31; // lane within warp
int warp = tid >> 5; // warp index (0 or 1)
T x = in;
__shared__ T sdata[2];
// Step 1: per‑warp reduction via shuffles
for (int offset = 16; offset > 0; offset >>= 1) {
#if defined(__CUDACC__)
unsigned mask = 0xffffffff;
for (int offset = 16; offset > 0; offset >>= 1)
x += __shfl_down_sync(mask, x, offset);
val += __shfl_down_sync(0xffffffff, val, offset);
#else
for (int offset = warpSize / 2; offset > 0; offset >>= 1)
x += __shfl_down(x, offset);
val += __shfl_down(val, offset);
#endif
// Step 2: each warp writes its partial sum to shared memory
if (lane == 0)
sdata[warp] = x;
__syncthreads(); // ensures both partial sums are visible
// Step 3: first warp loads the partials and finishes reduction
if (warp == 0) {
T v = (lane < 2) ? sdata[lane] : 0; // 2 warps → 2 partials
#if defined(__CUDACC__)
// reduce across first warp only
for (int offset = 16; offset > 0; offset >>= 1)
v += __shfl_down_sync(0xFFFFFFFF, v, offset);
#else
for (int offset = warpSize / 2; offset > 0; offset >>= 1)
v += __shfl_down_sync(0xFFFFFFFFFFFFFFFF, v, offset);
#endif
return v;
}
}
// prevents threads >= 32 (which never touch table[]) from racing ahead and
// overwriting it for the next call before the reduction above has read it.
__syncthreads();
return val;
}
template <typename T, bool COMPUTE_TAU>
@ -103,15 +65,14 @@ __global__ __launch_bounds__(64) void compute_hab_v4(const kernel_params dev_) {
// Copy task from global to shared memory and precompute some stuff.
extern __shared__ T shared_memory[];
// T *smem_cab = &shared_memory[dev_.smem_cab_offset];
const int number_of_tasks = dev_.num_tasks_per_block_dev[blockIdx.x];
const int number_of_tasks = dev_.num_tasks_per_block_dev[block_index()];
if (number_of_tasks == 0)
return;
T *smem_alpha = &shared_memory[0];
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int offset = dev_.sorted_blocks_offset_dev[blockIdx.x];
const int tid = thread_global_index();
const int offset = dev_.sorted_blocks_offset_dev[block_index()];
T *smem_cab = allocate_workspace<T>(dev_);
for (int tk = 0; tk < number_of_tasks; tk++) {
@ -158,15 +119,14 @@ template <typename T, typename T3, bool COMPUTE_TAU, bool CALCULATE_FORCES>
__global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
// Copy task from global to shared memory and precompute some stuff.
extern __shared__ T shared_memory[];
const int number_of_tasks = dev_.num_tasks_per_block_dev[blockIdx.x];
const int number_of_tasks = dev_.num_tasks_per_block_dev[block_index()];
if (number_of_tasks == 0)
return;
T *smem_alpha = &shared_memory[0];
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int offset = dev_.sorted_blocks_offset_dev[blockIdx.x];
const int tid = thread_global_index();
const int offset = dev_.sorted_blocks_offset_dev[block_index()];
T fa[3], fb[3];
T virial[9];
@ -197,7 +157,6 @@ __global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
fill_smem_task_coef(dev_, task_id, task);
T *coef_ = &dev_.ptr_dev[2][dev_.tasks[task_id].coef_offset];
compute_alpha(task, smem_alpha);
__syncthreads();
cxyz_to_cab(task, smem_alpha, coef_, smem_cab);
@ -286,34 +245,23 @@ __global__ __launch_bounds__(64) void compute_hab_v2(const kernel_params dev_) {
if (dev_.ptr_dev[5] != nullptr) {
for (int i = 0; i < 9; i++) {
sum[tid] = virial[i];
__syncthreads();
virial[i] = warp_reduce<T>(sum, tid);
__syncthreads();
virial[i] = block_reduce_64<T>(sum, virial[i], tid);
if (tid == 0)
atomicAdd(dev_.ptr_dev[5] + i, virial[i]);
__syncthreads();
}
}
for (int i = 0; i < 3; i++) {
sum[tid] = fa[i];
__syncthreads();
fa[i] = warp_reduce<T>(sum, tid);
__syncthreads();
fa[i] = block_reduce_64<T>(sum, fa[i], tid);
if (tid == 0)
atomicAdd(forces_a + i, fa[i]);
__syncthreads();
sum[tid] = fb[i];
__syncthreads();
fb[i] = warp_reduce<T>(sum, tid);
__syncthreads();
fb[i] = block_reduce_64<T>(sum, fb[i], tid);
if (tid == 0)
atomicAdd(forces_b + i, fb[i]);
__syncthreads();
}
}
}
@ -344,15 +292,14 @@ So most of the code is the same except the core of the routine that is
specialized to the integration.
******************************************************************************/
template <typename T, typename T3, bool distributed__, bool orthorhombic_,
template <typename T, typename T3, bool distributed__, bool orthogonal_,
int lbatch = 10>
__global__
__launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
if (dev_.tasks[dev_.first_task + block_index()].skip_task)
return;
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
const int tid = thread_global_index();
__shared__ T dh_inv_[9];
__shared__ T dh_[9];
@ -364,21 +311,10 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
dh_[d] = dev_.dh_[d];
__shared__ smem_task_reduced<T, T3> task;
fill_smem_task_reduced(dev_, dev_.first_task + blockIdx.x, task);
fill_smem_task_reduced(dev_, dev_.first_task + block_index(), task);
if (tid == 0) {
task.cube_center.z += task.lb_cube.z - dev_.grid_lower_corner_[0];
task.cube_center.y += task.lb_cube.y - dev_.grid_lower_corner_[1];
task.cube_center.x += task.lb_cube.x - dev_.grid_lower_corner_[2];
if (distributed__) {
if (task.apply_border_mask) {
compute_window_size(
dev_.grid_local_size_,
dev_.tasks[dev_.first_task + blockIdx.x].border_mask,
dev_.grid_border_width_, &task.window_size, &task.window_shift);
}
}
setup_task_cube_center<T, T3, distributed__>(dev_, task);
}
__syncthreads();
@ -399,10 +335,7 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
__syncthreads();
for (int z = threadIdx.z; z < task.cube_size.z; z += blockDim.z) {
int z2 = (z + task.cube_center.z) % dev_.grid_full_size_[0];
if (z2 < 0)
z2 += dev_.grid_full_size_[0];
int z2 = wrap_grid_index(z + task.cube_center.z, dev_.grid_full_size_.z);
if (distributed__) {
// known at compile time. Will be stripped away
@ -419,21 +352,13 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
int ymax = task.cube_size.y - 1;
T kremain = 0.0;
if (orthorhombic_ && !task.apply_border_mask) {
ymin = (2 * (z + task.lb_cube.z) - 1) / 2;
ymin *= ymin;
kremain = task.discrete_radius * task.discrete_radius -
ymin * dh_[8] * dh_[8];
ymin = ceil(-1.0e-8 - sqrt(fmax(0.0, kremain)) * dh_inv_[4]);
ymax = 1 - ymin - task.lb_cube.y;
ymin = ymin - task.lb_cube.y;
if (orthogonal_ && !task.apply_border_mask) {
kremain = calculate_ymix_ymax_boundaries(task, z, ymin, ymax);
}
for (int y = ymin + threadIdx.y; y <= ymax; y += blockDim.y) {
int y2 = (y + task.cube_center.y) % dev_.grid_full_size_[1];
if (y2 < 0)
y2 += dev_.grid_full_size_[1];
int y2 =
wrap_grid_index(y + task.cube_center.y, dev_.grid_full_size_.y);
if (distributed__) {
/* check if the point is within the window */
@ -447,21 +372,13 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
int xmin = 0;
int xmax = task.cube_size.x - 1;
if (orthorhombic_ && !task.apply_border_mask) {
xmin = (2 * (y + task.lb_cube.y) - 1) / 2;
xmin *= xmin;
xmin =
ceil(-1.0e-8 - sqrt(fmax(0.0, kremain - xmin * dh_[4] * dh_[4])) *
dh_inv_[0]);
xmax = 1 - xmin - task.lb_cube.x;
xmin -= task.lb_cube.x;
if (orthogonal_ && !task.apply_border_mask) {
calculate_xmin_xmax_boundaries<T, T3>(task, y, kremain, xmin, xmax);
}
for (int x = xmin + threadIdx.x; x <= xmax; x += blockDim.x) {
int x2 = (x + task.cube_center.x) % dev_.grid_full_size_[2];
if (x2 < 0)
x2 += dev_.grid_full_size_[2];
int x2 =
wrap_grid_index(x + task.cube_center.x, dev_.grid_full_size_.x);
if (distributed__) {
/* check if the point is within the window */
@ -472,34 +389,33 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
}
}
// I make no distinction between orthorhombic and non orthorhombic
// I make no distinction between orthogonal and non orthogonal
// cases
T3 r3;
if (orthorhombic_) {
r3.x = (x + task.lb_cube.x + task.roffset.x) * dh_[0];
r3.y = (y + task.lb_cube.y + task.roffset.y) * dh_[4];
r3.z = (z + task.lb_cube.z + task.roffset.z) * dh_[8];
} else {
r3 = compute_coordinates(dh_, (x + task.lb_cube.x + task.roffset.x),
(y + task.lb_cube.y + task.roffset.y),
(z + task.lb_cube.z + task.roffset.z));
}
r3 = compute_coordinates(dh_, (x + task.lb_cube.x + task.roffset.x),
(y + task.lb_cube.y + task.roffset.y),
(z + task.lb_cube.z + task.roffset.z));
// check if the point is inside the sphere or not. Note that it does
// not apply for the orthorhombic case when the full sphere is inside
// not apply for the orthogonal case when the full sphere is inside
// the region of interest.
const T r3x2 = r3.x * r3.x;
const T r3y2 = r3.y * r3.y;
const T r3z2 = r3.z * r3.z;
if (distributed__) {
if (((task.radius * task.radius) <=
(r3.x * r3.x + r3.y * r3.y + r3.z * r3.z)) &&
(!orthorhombic_ || task.apply_border_mask))
// check if the point is inside the sphere or not. Note that it does
// not apply for the orthorhombic case when the full sphere is
// inside the region of interest.
if (((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)) &&
(!orthogonal_ || task.apply_border_mask))
continue;
} else {
if (!orthorhombic_) {
if ((task.radius * task.radius) <=
(r3.x * r3.x + r3.y * r3.y + r3.z * r3.z))
continue;
}
// we do not need to do this test for the orthorhombic case
if ((!orthogonal_) &&
((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)))
continue;
}
// read the next point assuming that reading is non blocking untill
@ -507,14 +423,10 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
// NVIDIA hardware
T grid_value =
__ldg(&dev_.ptr_dev[1][(z2 * dev_.grid_local_size_[1] + y2) *
dev_.grid_local_size_[2] +
__ldg(&dev_.ptr_dev[1][(z2 * dev_.grid_local_size_.y + y2) *
dev_.grid_local_size_.x +
x2]);
const T r3x2 = r3.x * r3.x;
const T r3y2 = r3.y * r3.y;
const T r3z2 = r3.z * r3.z;
const T r3xy = r3.x * r3.y;
const T r3xz = r3.x * r3.z;
const T r3yz = r3.y * r3.z;
@ -680,10 +592,20 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
// Load local value
T val = accumulator[i][tid];
// Warp-level reduction (32 lanes)
// After this loop, lane 0 of warp 0 holds the result
// Warp-level reduction (32 lanes) After this loop, lane 0 of warp 0
// holds the result All threads should execute this operation so
// __activemask() is incorrect here.
// The ifdef statement is needed because hip does not necessarily
// support the instruction either. Reverting back to __shfl_down is
// required.
//
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_down_sync(__activemask(), val, offset);
#if defined(__CUDACC__)
val += __shfl_down_sync(0xffffffff, val, offset);
#else
val += __shfl_down(val, offset);
#endif
}
if (tid == 0)
@ -692,43 +614,9 @@ __launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
}
__syncthreads();
// if (tid < 16) {
// for (int i = 0; i < min(length - ico, lbatch); i++) {
// accumulator[i][tid] += accumulator[i][tid + 16];
// }
// }
// __syncthreads();
// if (tid < 8) {
// for (int i = 0; i < min(length - ico, lbatch); i++) {
// accumulator[i][tid] += accumulator[i][tid + 8];
// }
// }
// __syncthreads();
// if (tid < 4) {
// for (int i = 0; i < min(length - ico, lbatch); i++) {
// accumulator[i][tid] += accumulator[i][tid + 4];
// }
// }
// __syncthreads();
// if (tid < 2) {
// for (int i = 0; i < min(length - ico, lbatch); i++) {
// accumulator[i][tid] += accumulator[i][tid + 2];
// }
// }
// __syncthreads();
// if (tid == 0) {
// for (int i = 0; i < min(length - ico, lbatch); i++) {
// sum[i] = accumulator[i][0] + accumulator[i][1];
// }
// }
// __syncthreads();
if (tid < min(length - ico, lbatch))
dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset +
dev_.ptr_dev[2][dev_.tasks[dev_.first_task + block_index()].coef_offset +
tid + ico] = sum[tid];
__syncthreads();
}
}
@ -759,7 +647,7 @@ void context_info::integrate_one_grid_level(const int level, int *lp_diff) {
const dim3 threads_per_block(4, 4, 4);
if (grid_[level].is_distributed()) {
if (grid_[level].is_orthorhombic()) {
if (grid_[level].is_orthogonal()) {
integrate_kernel<double, double3, true, true, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);
@ -769,7 +657,7 @@ void context_info::integrate_one_grid_level(const int level, int *lp_diff) {
level_streams[level]>>>(params);
}
} else {
if (grid_[level].is_orthorhombic()) {
if (grid_[level].is_orthogonal()) {
integrate_kernel<double, double3, false, true, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);

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

@ -123,7 +123,12 @@ template <typename T> struct smem_task {
* \brief data needed for collocate and integrate kernels
******************************************************************************/
template <typename T, typename T3> struct smem_task_reduced {
// radius: true cutoff used for the sphere-membership test.
// discrete_radius: box-aligned cutoff (always >= radius) used only for
// sizing the cube / trimming box boundaries.
T radius, discrete_radius;
T norm_lattice_vector_z_2, norm_lattice_vector_y_2;
T norm_inverse_lattice_vector_y, norm_inverse_lattice_vector_x;
int3 cube_center, lb_cube, cube_size, window_size, window_shift;
T3 roffset;
T zetp;
@ -283,7 +288,7 @@ __inline__ __device__ unsigned int block_index() {
return blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z);
}
// Calculating the global index of any given device thread
// Calculating the global index in the grid of any given device thread
__inline__ __device__ unsigned int thread_global_index() {
return threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
}
@ -436,8 +441,8 @@ __inline__ T compute_cube_properties(
const T disr_radius =
std::min(norm1, std::min(norm2, norm3)) *
std::max(1,
(int)ceil(radius / std::min(norm1, std::min(norm2, norm3))));
(std::max(1,
(int)ceil(radius / std::min(norm1, std::min(norm2, norm3)))));
rp2.x = cubecenter->x;
rp2.y = cubecenter->y;
@ -450,6 +455,10 @@ __inline__ T compute_cube_properties(
roffset->y -= rp->y;
roffset->z -= rp->z;
/* This potentially makes the cube bigger than necessary because disc_radius
is derived with the smallest lattice parameter. It is not an issue when
the lattice is cubic but could lead to more calculations than necessary
when the lattice is orthorhombic */
rp2.x = disr_radius;
rp2.y = disr_radius;
rp2.z = disr_radius;
@ -527,31 +536,67 @@ __inline__ T compute_cube_properties(
}
}
__inline__ __device__ void compute_window_size(const int *const grid_size,
__inline__ __device__ void compute_window_size(const int3 grid_size,
const int border_mask,
const int *border_width,
int3 *const window_size,
int3 *const window_shift) {
window_shift->x = 0;
window_shift->y = 0;
window_shift->z = 0;
const int3 &border_width,
int3 &window_size,
int3 &window_shift) {
window_shift.x = 0;
window_shift.y = 0;
window_shift.z = 0;
window_size->x = grid_size[2] - 1;
window_size->y = grid_size[1] - 1;
window_size->z = grid_size[0] - 1;
window_size.x = grid_size.x - 1;
window_size.y = grid_size.y - 1;
window_size.z = grid_size.z - 1;
if (border_mask & (1 << 0))
window_shift->x += border_width[2];
window_shift.x += border_width.x;
if (border_mask & (1 << 1))
window_size->x -= border_width[2];
window_size.x -= border_width.x;
if (border_mask & (1 << 2))
window_shift->y += border_width[1];
window_shift.y += border_width.y;
if (border_mask & (1 << 3))
window_size->y -= border_width[1];
window_size.y -= border_width.y;
if (border_mask & (1 << 4))
window_shift->z += border_width[0];
window_shift.z += border_width.z;
if (border_mask & (1 << 5))
window_size->z -= border_width[0];
window_size.z -= border_width.z;
}
/*******************************************************************************
* \brief Wraps a cube-local coordinate into the periodic full grid along one
* axis. Single point of entry shared by collocate_kernel and
* integrate_kernel.
******************************************************************************/
__inline__ __device__ int wrap_grid_index(const int idx, const int full_size) {
int idx2 = idx % full_size;
if (idx2 < 0)
idx2 += full_size;
return idx2;
}
/*******************************************************************************
* \brief Shifts task.cube_center into the local grid's coordinate system and,
* for distributed grids, derives the border window. Single point of
* entry shared by collocate_kernel and integrate_kernel so the two
* cannot drift apart on which axis maps to which.
******************************************************************************/
template <typename T, typename T3, bool distributed__>
__device__ __inline__ void
setup_task_cube_center(const kernel_params &dev_,
smem_task_reduced<T, T3> &task) {
task.cube_center.x += task.lb_cube.x - dev_.grid_lower_corner_.x;
task.cube_center.y += task.lb_cube.y - dev_.grid_lower_corner_.y;
task.cube_center.z += task.lb_cube.z - dev_.grid_lower_corner_.z;
if (distributed__) {
if (task.apply_border_mask) {
compute_window_size(
dev_.grid_local_size_,
dev_.tasks[dev_.first_task + block_index()].border_mask,
dev_.grid_border_width_, task.window_size, task.window_shift);
}
}
}
/*******************************************************************************
@ -592,15 +637,14 @@ cab_to_cxyz(const smem_task<T> &task, const T *__restrict__ alpha,
const auto &b = coset_inv[jco];
for (int ico = 0; ico < ncoset(task.la_max); ico++) {
const auto &a = coset_inv[ico];
const T p = task.prefactor *
alpha[0 * s1 + b.l[0] * s2 + a.l[0] * s3 + co.l[0]] *
const T p = alpha[0 * s1 + b.l[0] * s2 + a.l[0] * s3 + co.l[0]] *
alpha[1 * s1 + b.l[1] * s2 + a.l[1] * s3 + co.l[1]] *
alpha[2 * s1 + b.l[2] * s2 + a.l[2] * s3 + co.l[2]];
reg += p * cab[jco * task.n1 + ico]; // collocate
}
}
cxyz[i] = reg;
cxyz[i] = task.prefactor * reg;
}
__syncthreads(); // because of concurrent writes to cxyz / cab
}
@ -640,14 +684,14 @@ cxyz_to_cab(const smem_task<T> &task, const T *__restrict__ alpha,
T reg = 0.0; // accumulate into a register
for (int ic = 0; ic < ncoset(task.lp); ic++) {
const auto &co = coset_inv[ic];
const T p = task.prefactor *
alpha[b.l[0] * s2 + a.l[0] * s3 + co.l[0]] *
const T p = alpha[b.l[0] * s2 + a.l[0] * s3 + co.l[0]] *
alpha[s1 + b.l[1] * s2 + a.l[1] * s3 + co.l[1]] *
alpha[2 * s1 + b.l[2] * s2 + a.l[2] * s3 + co.l[2]];
reg += p * cxyz[ic]; // integrate
}
cab[jco * task.n1 + ico] = reg; // partial loop coverage -> zero it
cab[jco * task.n1 + ico] =
task.prefactor * reg; // partial loop coverage -> zero it
}
}
}
@ -690,6 +734,20 @@ fill_smem_task_reduced(const kernel_params &dev, const int task_id,
task.lb_cube.y = glb_task.lb_cube.y;
task.lb_cube.z = glb_task.lb_cube.z;
task.norm_lattice_vector_z_2 = dev.dh_[6] * dev.dh_[6] +
dev.dh_[7] * dev.dh_[7] +
dev.dh_[8] * dev.dh_[8];
task.norm_lattice_vector_y_2 = dev.dh_[3] * dev.dh_[3] +
dev.dh_[4] * dev.dh_[4] +
dev.dh_[5] * dev.dh_[5];
task.norm_inverse_lattice_vector_y =
sqrt(dev.dh_inv_[3] * dev.dh_inv_[3] + dev.dh_inv_[4] * dev.dh_inv_[4] +
dev.dh_inv_[5] * dev.dh_inv_[5]);
task.norm_inverse_lattice_vector_x =
sqrt(dev.dh_inv_[0] * dev.dh_inv_[0] + dev.dh_inv_[1] * dev.dh_inv_[1] +
dev.dh_inv_[2] * dev.dh_inv_[2]);
task.apply_border_mask = glb_task.apply_border_mask;
}
__syncthreads();
@ -721,8 +779,6 @@ __device__ __inline__ void fill_smem_task_coef(const kernel_params &dev,
}
task.prefactor = glb_task.prefactor;
// task.radius = glb_task.radius;
task.off_diag_twice = glb_task.off_diag_twice;
// angular momentum range of basis set
@ -844,5 +900,34 @@ __inline__ __device__ T *allocate_workspace(const kernel_params &dev_) {
return (T *)(dev_.ptr_dev[6] + offset);
}
template <typename T, typename T3>
__device__ __inline__ T
calculate_ymix_ymax_boundaries(smem_task_reduced<T, T3> &task, const int z,
int &ymin, int &ymax) {
T kremain = 0.0;
ymin = (2 * (z + task.lb_cube.z) - 1) / 2;
ymin *= ymin;
kremain = task.discrete_radius * task.discrete_radius -
((T)ymin) * task.norm_lattice_vector_z_2;
ymin = ceil(-1.0e-8 -
sqrt(fmax(0.0, kremain)) * task.norm_inverse_lattice_vector_y);
ymax = 1 - ymin - task.lb_cube.y;
ymin = ymin - task.lb_cube.y;
return kremain;
}
template <typename T, typename T3>
__device__ __inline__ void
calculate_xmin_xmax_boundaries(smem_task_reduced<T, T3> &task, const int y,
const T kremain, int &xmin, int &xmax) {
xmin = (2 * (y + task.lb_cube.y) - 1) / 2;
xmin *= xmin;
xmin = ceil(-1.0e-8 -
sqrt(fmax(0.0, kremain - xmin * task.norm_lattice_vector_y_2)) *
task.norm_inverse_lattice_vector_x);
xmax = 1 - xmin - task.lb_cube.x;
xmin -= task.lb_cube.x;
}
} // namespace rocm_backend
#endif