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

Grid gpu engines merge (#5206)

Browse source 
Co-authored-by: Mathieu Taillefumier <mathieu.taillefumier@free.fr>
This commit is contained in:
Taillefumier Mathieu 2026-05-16 17:18:57 +00:00 • committed by GitHub
parent 024a11eaca
commit 695da86d0c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 1173 additions and 3557 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/CMakeLists.txt
View file

@ -1500,11 +1500,7 @@ set(CP2K_SRCS_GPU "")
set(CP2K_GRID_SRCS_GPU
grid/gpu/grid_gpu_collocate.cu grid/gpu/grid_gpu_integrate.cu
grid/gpu/grid_gpu_task_list.cu)
set(CP2K_GRID_SRCS_HIP
grid/hip/grid_hip_collocate.cu grid/hip/grid_hip_integrate.cu
grid/hip/grid_hip_context.cu)
grid/gpu/grid_gpu_context.cu)
set(CP2K_DBM_SRCS_GPU dbm/dbm_multiply_gpu_kernel.cu)
set(CP2K_DBM_SRCS_GPU_C dbm/dbm_multiply_opencl.c dbm/dbm_multiply_gpu.c)
@ -1557,9 +1553,6 @@ if(CP2K_USE_CUDA
endif()
if(CP2K_ENABLE_GRID_GPU)
if(CP2K_USE_HIP)
list(APPEND CP2K_SRCS_GPU ${CP2K_GRID_SRCS_HIP})
endif()
list(APPEND CP2K_SRCS_GPU ${CP2K_GRID_SRCS_GPU})
endif()
if(CP2K_USE_HIP)
@ -1974,7 +1967,7 @@ foreach(__app grid_miniapp grid_unittest dbm_miniapp)
set_target_properties(
${__app}
PROPERTIES POSITION_INDEPENDENT_CODE ON
LINKER_LANGUAGE C
LINKER_LANGUAGE CXX
POSITION_INDEPENDENT_CODE ON
OUTPUT_NAME "${__app}.${__cp2k_ext}")

4
src/environment.F
View file

@ -65,7 +65,6 @@ MODULE environment
GRID_BACKEND_CPU,&
GRID_BACKEND_DGEMM,&
GRID_BACKEND_GPU,&
GRID_BACKEND_HIP,&
GRID_BACKEND_REF
USE header, ONLY: cp2k_footer,&
cp2k_header
@ -920,9 +919,6 @@ CONTAINS
CASE (GRID_BACKEND_GPU)
WRITE (UNIT=output_unit, FMT="(T2,A,T75,A6)") &
start_section_label//"| Grid backend", "GPU"
CASE (GRID_BACKEND_HIP)
WRITE (UNIT=output_unit, FMT="(T2,A,T75,A6)") &
start_section_label//"| Grid backend", "HIP"
CASE (GRID_BACKEND_REF)
WRITE (UNIT=output_unit, FMT="(T2,A,T75,A6)") &
start_section_label//"| Grid backend", "REF"

1
src/grid/common/grid_constants.h
View file

@ -51,7 +51,6 @@ enum grid_backend {
GRID_BACKEND_CPU = 12,
GRID_BACKEND_DGEMM = 13,
GRID_BACKEND_GPU = 14,
GRID_BACKEND_HIP = 15,
};
#endif

2
src/grid/common/grid_library.c
View file

@ -230,7 +230,7 @@ void grid_library_print_stats(const int fortran_comm,
const char *kernel_names[] = {"collocate ortho", "integrate ortho",
"collocate general", "integrate general"};
const char *backend_names[] = {"REF", "CPU", "DGEMM", "GPU", "HIP"};
const char *backend_names[] = {"REF", "CPU", "DGEMM", "GPU"};
for (int i = 0; i < ncounters; i++) {
if (counters[i][0] == 0)

559
src/grid/gpu/grid_gpu_collint.h
View file

@ -1,559 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#include <algorithm>
#include <assert.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "../common/grid_basis_set.h"
#include "../common/grid_common.h"
#include "grid_gpu_task_list.h"
#if (GRID_DO_COLLOCATE)
#define GRID_CONST_WHEN_COLLOCATE const
#define GRID_CONST_WHEN_INTEGRATE
#else
#define GRID_CONST_WHEN_COLLOCATE
#define GRID_CONST_WHEN_INTEGRATE const
#endif
/*******************************************************************************
* \brief Forward declarations for inner-most loop bodies.
* Implementations are in grid_gpu_collocate.cu / grid_gpu_integrate.cu.
* \author Ole Schuett
******************************************************************************/
#if (GRID_DO_COLLOCATE)
// collocate
typedef double cxyz_store;
__device__ static void cxyz_to_gridpoint(const double dx, const double dy,
const double dz, const double zetp,
const int lp, const cxyz_store *cxyz,
double *gridpoint);
#else
// integrate
typedef struct {
double *regs;
int offset;
} cxyz_store;
__device__ static void gridpoint_to_cxyz(const double dx, const double dy,
const double dz, const double zetp,
const int lp, const double *gridpoint,
cxyz_store *store);
#endif
/*******************************************************************************
* \brief Atomic add for doubles that also works prior to compute capability 6.
* \author Ole Schuett
******************************************************************************/
__device__ static void atomicAddDouble(double *address, double val) {
if (val == 0.0)
return;
#if __CUDA_ARCH__ >= 600
atomicAdd(address, val); // part of cuda library
#else
// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
unsigned long long int *address_as_ull = (unsigned long long int *)address;
unsigned long long int old = *address_as_ull, assumed;
do {
assumed = old;
old = atomicCAS(address_as_ull, assumed,
__double_as_longlong(val + __longlong_as_double(assumed)));
// Uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
} while (assumed != old);
#endif
}
/*******************************************************************************
* \brief Cab matrix container to be passed through prepare_pab to cab_add.
* \author Ole Schuett
******************************************************************************/
typedef struct {
double *data;
const int n1;
} cab_store;
/*******************************************************************************
* \brief Returns matrix element cab[idx(b)][idx(a)].
* \author Ole Schuett
******************************************************************************/
__device__ static inline double cab_get(const cab_store *cab, const orbital a,
const orbital b) {
const int i = idx(b) * cab->n1 + idx(a);
return cab->data[i];
}
/*******************************************************************************
* \brief Adds given value to matrix element cab[idx(b)][idx(a)].
* \author Ole Schuett
******************************************************************************/
__device__ static inline void cab_add(cab_store *cab, const orbital a,
const orbital b, const double value) {
const int i = idx(b) * cab->n1 + idx(a);
atomicAddDouble(&cab->data[i], value);
}
/*******************************************************************************
* \brief Parameters of the collocate kernel.
* \author Ole Schuett
******************************************************************************/
typedef struct {
int smem_cab_length;
int smem_cab_offset;
int smem_alpha_offset;
int smem_cxyz_offset;
int first_task;
int npts_global[3];
int npts_local[3];
int shift_local[3];
double dh[3][3];
double dh_inv[3][3];
const grid_gpu_task *tasks;
const double *pab_blocks;
GRID_CONST_WHEN_INTEGRATE double *grid;
int la_min_diff;
int lb_min_diff;
int la_max_diff;
int lb_max_diff;
#if (GRID_DO_COLLOCATE)
// collocate
enum grid_func func;
#else
// integrate
double *hab_blocks;
double *forces;
double *virial;
#endif
} kernel_params;
/*******************************************************************************
* \brief Shared memory representation of a task.
* \author Ole Schuett
******************************************************************************/
typedef struct smem_task_struct : grid_gpu_task_struct {
// angular momentum range of actual collocate / integrate operation
int la_max;
int lb_max;
int la_min;
int lb_min;
int lp;
// size of the cab matrix
int n1;
int n2;
const double *pab_block;
#if (!GRID_DO_COLLOCATE)
// integrate
double *hab_block;
double *forces_a;
double *forces_b;
#endif
} smem_task;
/*******************************************************************************
* \brief Tabulated functions kept in constant memory to reduce register count.
* \author Ole Schuett
******************************************************************************/
__constant__ orbital coset_inv[1330];
__constant__ double binomial_coef[19][19];
/*******************************************************************************
* \brief Initializes the device's constant memory.
* \author Ole Schuett
******************************************************************************/
static void init_constant_memory(void) {
static bool initialized = false;
if (initialized) {
return; // constant memory has to be initialized only once
}
// Inverse coset mapping
orbital coset_inv_host[1330];
for (int lx = 0; lx <= 18; lx++) {
for (int ly = 0; ly <= 18 - lx; ly++) {
for (int lz = 0; lz <= 18 - lx - ly; lz++) {
const int i = coset(lx, ly, lz);
coset_inv_host[i] = {{lx, ly, lz}};
}
}
}
offloadMemcpyToSymbol(coset_inv, &coset_inv_host, sizeof(coset_inv_host));
// Binomial coefficient
double binomial_coef_host[19][19];
memset(binomial_coef_host, 0, sizeof(binomial_coef_host));
for (int n = 0; n <= 18; n++) {
for (int k = 0; k <= n; k++) {
binomial_coef_host[n][k] = fac(n) / fac(k) / fac(n - k);
}
}
offloadMemcpyToSymbol(binomial_coef, &binomial_coef_host,
sizeof(binomial_coef_host));
initialized = true;
}
/*******************************************************************************
* \brief Collocates coefficients C_xyz onto the grid for orthorhombic case.
* \author Ole Schuett
******************************************************************************/
__device__ static void
ortho_cxyz_to_grid(const kernel_params *params, const smem_task *task,
GRID_CONST_WHEN_COLLOCATE cxyz_store *cxyz,
GRID_CONST_WHEN_INTEGRATE double *grid) {
// The cube contains an even number of grid points in each direction and
// collocation is always performed on a pair of two opposing grid points.
// Hence, the points with index 0 and 1 are both assigned distance zero via
// the formular distance=(2*index-1)/2.
const int kmin = ceil(-1e-8 - task->disr_radius * params->dh_inv[2][2]);
for (int k = threadIdx.z + kmin; k <= 1 - kmin; k += blockDim.z) {
const int ka = task->cube_center_shifted[2] + k;
const int kg =
modulo(ka, params->npts_global[2]); // target location on the grid
const int kd = (2 * k - 1) / 2; // distance from center in grid points
const double kr = kd * params->dh[2][2]; // distance from center in a.u.
const double kremain = task->disr_radius * task->disr_radius - kr * kr;
const int jmin =
ceil(-1e-8 - sqrt(fmax(0.0, kremain)) * params->dh_inv[1][1]);
for (int j = threadIdx.y + jmin; j <= 1 - jmin; j += blockDim.y) {
const int ja = task->cube_center_shifted[1] + j;
const int jg =
modulo(ja, params->npts_global[1]); // target location on the grid
const int jd = (2 * j - 1) / 2; // distance from center in grid points
const double jr = jd * params->dh[1][1]; // distance from center in a.u.
const double jremain = kremain - jr * jr;
const int imin =
ceil(-1e-8 - sqrt(fmax(0.0, jremain)) * params->dh_inv[0][0]);
for (int i = threadIdx.x + imin; i <= 1 - imin; i += blockDim.x) {
const int ia = task->cube_center_shifted[0] + i;
const int ig =
modulo(ia, params->npts_global[0]); // target location on the grid
// The distances above (kd, jd) do not take roffset into account,
// ie. they always snap to the next grid point. This allowed the legacy
// implementation to cache the loop bounds.
// For the calculation of the grid value we'll use the true distance.
const double dx = i * params->dh[0][0] + task->cube_offset[0];
const double dy = j * params->dh[1][1] + task->cube_offset[1];
const double dz = k * params->dh[2][2] + task->cube_offset[2];
const int grid_index =
kg * params->npts_local[1] * params->npts_local[0] +
jg * params->npts_local[0] + ig;
#if (GRID_DO_COLLOCATE)
// collocate
cxyz_to_gridpoint(dx, dy, dz, task->zetp, task->lp, cxyz,
&grid[grid_index]);
#else
// integrate
gridpoint_to_cxyz(dx, dy, dz, task->zetp, task->lp, &grid[grid_index],
cxyz);
#endif
}
}
}
__syncthreads(); // because of concurrent writes to grid / cxyz
}
/*******************************************************************************
* \brief Collocates coefficients C_xyz onto the grid for general case.
* \author Ole Schuett
******************************************************************************/
__device__ static void
general_cxyz_to_grid(const kernel_params *params, const smem_task *task,
GRID_CONST_WHEN_COLLOCATE cxyz_store *cxyz,
GRID_CONST_WHEN_INTEGRATE double *grid) {
// Go over the grid
const int k_start = threadIdx.z + task->index_min[2];
for (int k = k_start; k <= task->index_max[2]; k += blockDim.z) {
const int kg = modulo(k - params->shift_local[2], params->npts_global[2]);
if (kg < task->bounds_k[0] || task->bounds_k[1] < kg) {
continue;
}
const int j_start = threadIdx.y + task->index_min[1];
for (int j = j_start; j <= task->index_max[1]; j += blockDim.y) {
const int jg = modulo(j - params->shift_local[1], params->npts_global[1]);
if (jg < task->bounds_j[0] || task->bounds_j[1] < jg) {
continue;
}
const int i_start = threadIdx.x + task->index_min[0];
for (int i = i_start; i <= task->index_max[0]; i += blockDim.x) {
const int ig =
modulo(i - params->shift_local[0], params->npts_global[0]);
if (ig < task->bounds_i[0] || task->bounds_i[1] < ig) {
continue;
}
// Compute distance of current grid point from center of Gaussian.
const double di = i - task->gp[0];
double dx = di * params->dh[0][0];
double dy = di * params->dh[0][1];
double dz = di * params->dh[0][2];
const double dj = j - task->gp[1];
dx += dj * params->dh[1][0];
dy += dj * params->dh[1][1];
dz += dj * params->dh[1][2];
const double dk = k - task->gp[2];
dx += dk * params->dh[2][0];
dy += dk * params->dh[2][1];
dz += dk * params->dh[2][2];
// Cycle if the point is not within the radius.
const double r2 = dx * dx + dy * dy + dz * dz;
if (r2 >= task->radius2) {
continue;
}
const int stride = params->npts_local[1] * params->npts_local[0];
const int grid_index = kg * stride + jg * params->npts_local[0] + ig;
#if (GRID_DO_COLLOCATE)
// collocate
cxyz_to_gridpoint(dx, dy, dz, task->zetp, task->lp, cxyz,
&grid[grid_index]);
#else
// integrate
gridpoint_to_cxyz(dx, dy, dz, task->zetp, task->lp, &grid[grid_index],
cxyz);
#endif
}
}
}
__syncthreads(); // because of concurrent writes to grid / cxyz
}
/*******************************************************************************
* \brief Computes the polynomial expansion coefficients:
* (x-a)**lxa (x-b)**lxb -> sum_{ls} alpha(ls,lxa,lxb,1)*(x-p)**ls
* \author Ole Schuett
******************************************************************************/
__device__ static void compute_alpha(const smem_task *task, double *alpha) {
// strides for accessing alpha
const int s3 = (task->lp + 1);
const int s2 = (task->la_max + 1) * s3;
const int s1 = (task->lb_max + 1) * s2;
for (int idir = threadIdx.z; idir < 3; idir += blockDim.z) {
const double drpa = task->rp[idir] - task->ra[idir];
const double drpb = task->rp[idir] - task->rb[idir];
for (int la = threadIdx.y; la <= task->la_max; la += blockDim.y) {
for (int lb = threadIdx.x; lb <= task->lb_max; lb += blockDim.x) {
for (int i = 0; i <= task->lp; i++) {
const int base = idir * s1 + lb * s2 + la * s3;
alpha[base + i] = 0.0;
}
double a = 1.0;
for (int k = 0; k <= la; k++) {
double b = 1.0;
for (int l = 0; l <= lb; l++) {
const int base = idir * s1 + lb * s2 + la * s3;
alpha[base + la - l + lb - k] +=
a * b * binomial_coef[la][k] * binomial_coef[lb][l];
b *= drpb;
}
a *= drpa;
}
}
}
}
__syncthreads(); // because of concurrent writes to alpha
}
/*******************************************************************************
* \brief Transforms coefficients C_ab into C_xyz.
* \author Ole Schuett
******************************************************************************/
__device__ static void cab_to_cxyz(const smem_task *task, const double *alpha,
GRID_CONST_WHEN_COLLOCATE cab_store *cab,
GRID_CONST_WHEN_INTEGRATE double *cxyz) {
// *** initialise the coefficient matrix, we transform the sum
//
// sum_{lxa,lya,lza,lxb,lyb,lzb} P_{lxa,lya,lza,lxb,lyb,lzb} *
// (x-a_x)**lxa (y-a_y)**lya (z-a_z)**lza (x-b_x)**lxb (y-a_y)**lya
// (z-a_z)**lza
//
// into
//
// sum_{lxp,lyp,lzp} P_{lxp,lyp,lzp} (x-p_x)**lxp (y-p_y)**lyp (z-p_z)**lzp
//
// where p is center of the product gaussian, and lp = la_max + lb_max
// (current implementation is l**7)
// strides for accessing alpha
const int s3 = (task->lp + 1);
const int s2 = (task->la_max + 1) * s3;
const int s1 = (task->lb_max + 1) * s2;
// TODO: Maybe we can transpose alpha to index it directly with ico and jco.
#if (GRID_DO_COLLOCATE)
// collocate
for (int lzp = threadIdx.z; lzp <= task->lp; lzp += blockDim.z) {
for (int lyp = threadIdx.y; lyp <= task->lp - lzp; lyp += blockDim.y) {
for (int lxp = threadIdx.x; lxp <= task->lp - lzp - lyp;
lxp += blockDim.x) {
double reg = 0.0; // accumulate into a register
const int jco_start = ncoset(task->lb_min - 1);
const int jco_end = ncoset(task->lb_max);
for (int jco = jco_start; jco < jco_end; jco++) {
const orbital b = coset_inv[jco];
const int ico_start = ncoset(task->la_min - 1);
const int ico_end = ncoset(task->la_max);
for (int ico = ico_start; ico < ico_end; ico++) {
const orbital a = coset_inv[ico];
#else
// integrate
if (threadIdx.z == 0) { // TODO: How bad is this?
const int jco_start = ncoset(task->lb_min - 1) + threadIdx.y;
const int jco_end = ncoset(task->lb_max);
for (int jco = jco_start; jco < jco_end; jco += blockDim.y) {
const orbital b = coset_inv[jco];
const int ico_start = ncoset(task->la_min - 1) + threadIdx.x;
const int ico_end = ncoset(task->la_max);
for (int ico = ico_start; ico < ico_end; ico += blockDim.x) {
const orbital a = coset_inv[ico];
double reg = 0.0; // accumulate into a register
for (int lzp = 0; lzp <= task->lp; lzp++) {
for (int lyp = 0; lyp <= task->lp - lzp; lyp++) {
for (int lxp = 0; lxp <= task->lp - lzp - lyp; lxp++) {
#endif
const double p = task->prefactor *
alpha[0 * s1 + b.l[0] * s2 + a.l[0] * s3 + lxp] *
alpha[1 * s1 + b.l[1] * s2 + a.l[1] * s3 + lyp] *
alpha[2 * s1 + b.l[2] * s2 + a.l[2] * s3 + lzp];
#if (GRID_DO_COLLOCATE)
reg += p * cab_get(cab, a, b); // collocate
#else
reg += p * cxyz[coset(lxp, lyp, lzp)]; // integrate
#endif
}
}
#if (GRID_DO_COLLOCATE)
// collocate
cxyz[coset(lxp, lyp, lzp)] = reg; // overwrite - no zeroing needed.
}
#else
// integrate
}
cab_add(cab, a, b, reg); // partial loop coverage -> zero it
}
#endif
}
}
__syncthreads(); // because of concurrent writes to cxyz / cab
}
/*******************************************************************************
* \brief Allocates Cab block from global memory.
* \author Ole Schuett
******************************************************************************/
__device__ static double *malloc_cab(const smem_task *task) {
__shared__ double *gmem_cab;
if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
gmem_cab = (double *)malloc(task->n1 * task->n2 * sizeof(double));
assert(gmem_cab != NULL &&
"MallocHeapSize too low, please increase it in grid_library_init()");
}
__syncthreads(); // wait for write to shared gmem_cab variable
return gmem_cab;
}
/*******************************************************************************
* \brief Frees Cab block. Only needed if it was allocated from global memory.
* \author Ole Schuett
******************************************************************************/
__device__ static void free_cab(double *gmem_cab) {
__syncthreads(); // Ensure that all threads have finished using cab.
if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
free(gmem_cab);
}
}
/*******************************************************************************
* \brief Initializes the cab matrix with zeros.
* \author Ole Schuett
******************************************************************************/
__device__ static void zero_cab(cab_store *cab, const int cab_len) {
if (threadIdx.z == 0 && threadIdx.y == 0) {
for (int i = threadIdx.x; i < cab_len; i += blockDim.x) {
cab->data[i] = 0.0;
}
}
__syncthreads(); // because of concurrent writes to cab
}
/*******************************************************************************
* \brief Copies a task from global to shared memory and does precomputations.
* \author Ole Schuett
******************************************************************************/
__device__ static void load_task(const kernel_params *params, smem_task *task) {
// Parallel copy of base task from global to shared memory.
if (threadIdx.y == 0 && threadIdx.z == 0) {
const int itask = params->first_task + blockIdx.x;
const grid_gpu_task *src_task = &params->tasks[itask];
grid_gpu_task *dest_task = task; // Upcast to base struct.
for (int i = threadIdx.x; i < (int)sizeof(grid_gpu_task); i += blockDim.x) {
((char *)dest_task)[i] = ((const char *)src_task)[i];
}
}
__syncthreads(); // because of concurrent writes to task
if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
// angular momentum range for the actual collocate/integrate opteration.
task->la_max = task->la_max_basis + params->la_max_diff;
task->lb_max = task->lb_max_basis + params->lb_max_diff;
task->la_min = imax(task->la_min_basis + params->la_min_diff, 0);
task->lb_min = imax(task->lb_min_basis + params->lb_min_diff, 0);
task->lp = task->la_max + task->lb_max;
// size of the cab matrix
task->n1 = ncoset(task->la_max);
task->n2 = ncoset(task->lb_max);
task->pab_block = &params->pab_blocks[task->ab_block_offset];
// integrate
#if (!GRID_DO_COLLOCATE)
task->hab_block = &params->hab_blocks[task->ab_block_offset];
if (params->forces != NULL) {
task->forces_a = &params->forces[3 * task->iatom];
task->forces_b = &params->forces[3 * task->jatom];
}
#endif
}
__syncthreads(); // because of concurrent writes to task
}
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
// EOF

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

@ -5,377 +5,479 @@
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
/*
* Authors :
- Mathieu Taillefumier (ETH Zurich / CSCS)
- Advanced Micro Devices, Inc.
- Ole Schuett
*/
#include <algorithm>
#include <assert.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#define GRID_DO_COLLOCATE 1
#include "../common/grid_common.h"
#include "grid_gpu_collint.h"
#include "grid_gpu_collocate.h"
// This has to be included after grid_gpu_collint.h
#include "../common/grid_prepare_pab.h"
#include "grid_gpu_internal_header.h"
#include "grid_gpu_prepare_pab.h"
#if defined(_OMP_H)
#error "OpenMP should not be used in .cu files to accommodate HIP."
#endif
/*******************************************************************************
* \brief Collocate a single grid point with distance d{xyz} from center.
* \author Ole Schuett
******************************************************************************/
__device__ static void cxyz_to_gridpoint(const double dx, const double dy,
const double dz, const double zetp,
const int lp, const cxyz_store *cxyz,
double *gridpoint) {
// Squared distance of point from center.
const double r2 = dx * dx + dy * dy + dz * dz;
const double gaussian = exp(-zetp * r2);
// accumulate into register
double gridpoint_reg = 0.0;
// Manually unrolled loops based on terms in coset_inv.
// For lp > 6 the register usage increases and the kernel becomes slower.
gridpoint_reg += cxyz[0];
if (lp >= 1) {
gridpoint_reg += cxyz[1] * dx;
gridpoint_reg += cxyz[2] * dy;
gridpoint_reg += cxyz[3] * dz;
if (lp >= 2) {
const double dx2 = dx * dx;
const double dy2 = dy * dy;
const double dz2 = dz * dz;
gridpoint_reg += cxyz[4] * dx2;
gridpoint_reg += cxyz[5] * dx * dy;
gridpoint_reg += cxyz[6] * dx * dz;
gridpoint_reg += cxyz[7] * dy2;
gridpoint_reg += cxyz[8] * dy * dz;
gridpoint_reg += cxyz[9] * dz2;
if (lp >= 3) {
const double dx3 = dx2 * dx;
const double dy3 = dy2 * dy;
const double dz3 = dz2 * dz;
gridpoint_reg += cxyz[10] * dx3;
gridpoint_reg += cxyz[11] * dx2 * dy;
gridpoint_reg += cxyz[12] * dx2 * dz;
gridpoint_reg += cxyz[13] * dx * dy2;
gridpoint_reg += cxyz[14] * dx * dy * dz;
gridpoint_reg += cxyz[15] * dx * dz2;
gridpoint_reg += cxyz[16] * dy3;
gridpoint_reg += cxyz[17] * dy2 * dz;
gridpoint_reg += cxyz[18] * dy * dz2;
gridpoint_reg += cxyz[19] * dz3;
if (lp >= 4) {
const double dx4 = dx3 * dx;
const double dy4 = dy3 * dy;
const double dz4 = dz3 * dz;
gridpoint_reg += cxyz[20] * dx4;
gridpoint_reg += cxyz[21] * dx3 * dy;
gridpoint_reg += cxyz[22] * dx3 * dz;
gridpoint_reg += cxyz[23] * dx2 * dy2;
gridpoint_reg += cxyz[24] * dx2 * dy * dz;
gridpoint_reg += cxyz[25] * dx2 * dz2;
gridpoint_reg += cxyz[26] * dx * dy3;
gridpoint_reg += cxyz[27] * dx * dy2 * dz;
gridpoint_reg += cxyz[28] * dx * dy * dz2;
gridpoint_reg += cxyz[29] * dx * dz3;
gridpoint_reg += cxyz[30] * dy4;
gridpoint_reg += cxyz[31] * dy3 * dz;
gridpoint_reg += cxyz[32] * dy2 * dz2;
gridpoint_reg += cxyz[33] * dy * dz3;
gridpoint_reg += cxyz[34] * dz4;
if (lp >= 5) {
const double dx5 = dx4 * dx;
const double dy5 = dy4 * dy;
const double dz5 = dz4 * dz;
gridpoint_reg += cxyz[35] * dx5;
gridpoint_reg += cxyz[36] * dx4 * dy;
gridpoint_reg += cxyz[37] * dx4 * dz;
gridpoint_reg += cxyz[38] * dx3 * dy2;
gridpoint_reg += cxyz[39] * dx3 * dy * dz;
gridpoint_reg += cxyz[40] * dx3 * dz2;
gridpoint_reg += cxyz[41] * dx2 * dy3;
gridpoint_reg += cxyz[42] * dx2 * dy2 * dz;
gridpoint_reg += cxyz[43] * dx2 * dy * dz2;
gridpoint_reg += cxyz[44] * dx2 * dz3;
gridpoint_reg += cxyz[45] * dx * dy4;
gridpoint_reg += cxyz[46] * dx * dy3 * dz;
gridpoint_reg += cxyz[47] * dx * dy2 * dz2;
gridpoint_reg += cxyz[48] * dx * dy * dz3;
gridpoint_reg += cxyz[49] * dx * dz4;
gridpoint_reg += cxyz[50] * dy5;
gridpoint_reg += cxyz[51] * dy4 * dz;
gridpoint_reg += cxyz[52] * dy3 * dz2;
gridpoint_reg += cxyz[53] * dy2 * dz3;
gridpoint_reg += cxyz[54] * dy * dz4;
gridpoint_reg += cxyz[55] * dz5;
if (lp >= 6) {
const double dx6 = dx5 * dx;
const double dy6 = dy5 * dy;
const double dz6 = dz5 * dz;
gridpoint_reg += cxyz[56] * dx6;
gridpoint_reg += cxyz[57] * dx5 * dy;
gridpoint_reg += cxyz[58] * dx5 * dz;
gridpoint_reg += cxyz[59] * dx4 * dy2;
gridpoint_reg += cxyz[60] * dx4 * dy * dz;
gridpoint_reg += cxyz[61] * dx4 * dz2;
gridpoint_reg += cxyz[62] * dx3 * dy3;
gridpoint_reg += cxyz[63] * dx3 * dy2 * dz;
gridpoint_reg += cxyz[64] * dx3 * dy * dz2;
gridpoint_reg += cxyz[65] * dx3 * dz3;
gridpoint_reg += cxyz[66] * dx2 * dy4;
gridpoint_reg += cxyz[67] * dx2 * dy3 * dz;
gridpoint_reg += cxyz[68] * dx2 * dy2 * dz2;
gridpoint_reg += cxyz[69] * dx2 * dy * dz3;
gridpoint_reg += cxyz[70] * dx2 * dz4;
gridpoint_reg += cxyz[71] * dx * dy5;
gridpoint_reg += cxyz[72] * dx * dy4 * dz;
gridpoint_reg += cxyz[73] * dx * dy3 * dz2;
gridpoint_reg += cxyz[74] * dx * dy2 * dz3;
gridpoint_reg += cxyz[75] * dx * dy * dz4;
gridpoint_reg += cxyz[76] * dx * dz5;
gridpoint_reg += cxyz[77] * dy6;
gridpoint_reg += cxyz[78] * dy5 * dz;
gridpoint_reg += cxyz[79] * dy4 * dz2;
gridpoint_reg += cxyz[80] * dy3 * dz3;
gridpoint_reg += cxyz[81] * dy2 * dz4;
gridpoint_reg += cxyz[82] * dy * dz5;
gridpoint_reg += cxyz[83] * dz6;
}
}
}
}
}
}
// Handle higher values of lp.
if (lp >= 7) {
for (int i = 84; i < ncoset(lp); i++) {
double val = cxyz[i];
const orbital a = coset_inv[i];
for (int j = 0; j < a.l[0]; j++) {
val *= dx;
}
for (int j = 0; j < a.l[1]; j++) {
val *= dy;
}
for (int j = 0; j < a.l[2]; j++) {
val *= dz;
}
gridpoint_reg += val;
}
}
atomicAddDouble(gridpoint, gridpoint_reg * gaussian);
}
/*******************************************************************************
* \brief Collocates coefficients C_xyz onto the grid.
* \author Ole Schuett
******************************************************************************/
__device__ static void cxyz_to_grid(const kernel_params *params,
const smem_task *task, const double *cxyz,
double *grid) {
if (task->use_orthorhombic_kernel) {
ortho_cxyz_to_grid(params, task, cxyz, grid);
} else {
general_cxyz_to_grid(params, task, cxyz, grid);
}
}
namespace rocm_backend {
/*******************************************************************************
* \brief Decontracts the subblock, going from spherical to cartesian harmonics.
* \author Ole Schuett
******************************************************************************/
template <bool IS_FUNC_AB>
__device__ static void block_to_cab(const kernel_params *params,
const smem_task *task, cab_store *cab) {
template <typename T, bool IS_FUNC_AB>
__device__ __inline__ void block_to_cab(const kernel_params &params,
const smem_task<T> &task, T *cab) {
// The carthesian index runs over exponents and then over angular momentum.
// The spherical index runs over angular momentum and then over contractions.
// The cartesian index runs over exponents and then over angular momentum.
// Decontract block, apply prepare_pab, and store in cab.
// This is a double matrix product. Since the pab block can be quite large the
// 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);
// TODO find if we can make better bounds for ico and jco because
// we only need a submatrix of cab.
if (threadIdx.z == 0) { // TODO: How bad is this?
const int jco_start = ncoset(task->lb_min_basis - 1) + threadIdx.y;
const int jco_end = ncoset(task->lb_max_basis);
for (int jco = jco_start; jco < jco_end; jco += blockDim.y) {
const orbital b = coset_inv[jco];
const int ico_start = ncoset(task->la_min_basis - 1) + threadIdx.x;
const int ico_end = ncoset(task->la_max_basis);
for (int ico = ico_start; ico < ico_end; ico += blockDim.x) {
const orbital a = coset_inv[ico];
double pab_val = 0.0;
for (int i = 0; i < task->nsgf_setb; i++) {
const double sphib = task->sphib[i * task->maxcob + idx(b)];
for (int j = 0; j < task->nsgf_seta; j++) {
double block_val;
if (task->block_transposed) {
block_val =
task->pab_block[j * task->nsgfb + i] * task->off_diag_twice;
} else {
block_val =
task->pab_block[i * task->nsgfa + j] * task->off_diag_twice;
}
const double sphia = task->sphia[j * task->maxcoa + idx(a)];
pab_val += block_val * sphia * sphib;
for (int jco = task.first_cosetb + tid / 16; jco < task.ncosetb; jco += 4) {
for (int ico = task.first_coseta + (tid % 16); ico < task.ncoseta;
ico += 16) {
T pab_val = 0.0;
for (int i = 0; i < task.nsgf_setb; i++) {
const T sphib = task.sphib[i * task.maxcob + jco];
T tmp_val = 0.0;
for (int j = 0; j < task.nsgf_seta; j++) {
T block_val;
if (task.block_transposed) {
block_val =
task.pab_block[j * task.nsgfb + i] * task.off_diag_twice;
} else {
block_val =
task.pab_block[i * task.nsgfa + j] * task.off_diag_twice;
}
// const T sphia = task.sphia[j * task.maxcoa + ico];
tmp_val += block_val * task.sphia[j * task.maxcoa + ico];
}
if (IS_FUNC_AB) {
// fast path for common case
prepare_pab(GRID_FUNC_AB, a, b, task->zeta, task->zetb, pab_val, cab);
} else {
// Since prepare_pab is a register hog we use it only when needed.
prepare_pab(params->func, a, b, task->zeta, task->zetb, pab_val, cab);
}
pab_val += tmp_val * sphib;
}
if (IS_FUNC_AB) {
cab[jco * task.ncoseta + ico] += pab_val;
} else {
const auto a = coset_inv[ico];
const auto b = coset_inv[jco];
// remains the atomicAdd used in prepare_pab. Are they needed ?
prepare_pab(params.func, a, b, task.zeta, task.zetb, pab_val, task.n1,
cab);
}
}
}
__syncthreads(); // because of concurrent writes to cab
}
/*******************************************************************************
* \brief Cuda kernel for collocating all tasks of one grid level.
* \author Ole Schuett
******************************************************************************/
template <bool IS_FUNC_AB>
__device__ static void collocate_kernel(const kernel_params *params) {
template <typename T, bool IS_FUNC_AB>
__global__
__launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
__shared__ smem_task<T> task;
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
return;
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
fill_smem_task_coef(dev_, dev_.first_task + blockIdx.x, task);
extern __shared__ T shared_memory[];
// T *smem_cab = &shared_memory[dev_.smem_cab_offset];
T *smem_alpha = &shared_memory[dev_.smem_alpha_offset];
T *coef_ =
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset];
T *smem_cab =
&dev_.ptr_dev[6][dev_.tasks[dev_.first_task + blockIdx.x].cab_offset];
compute_alpha(task, smem_alpha);
for (int z = tid; z < task.n1 * task.n2;
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_);
}
/*
\brief compute the real space representation of an operator expressed in the
gaussian basis
this kernel does the following operation
n_{ijk} = \f[
\sum_{p,\alpha,\beta,\gamma} C^p_{\alpha\beta\gamma} X_{\alpha,i} Y_{\beta, j}
Z_{\gamma, k} \exp\left(- \eta (r_{ijk} - r_c) ^ 2\right)
]
where $X_{\alpha,i}, Y_{\beta, j}, Z_{\gamma, k}$ are polynomials of degree
$\alpha,\beta,\gamma$ and $r_{ijk}% a point (in cartesian coordinates) on a 3D
grid. C^p_{\alpha\beta\gamma} are also constrained such that 0 <= \alpha +
\beta + \gamma <= lmax. It means in practice that we need store (lmax + 1) *
(lamax + 2) * (lmax + 3) / 6 coefficients all the other coefficients are zero
to reduce computation, a spherical cutoff is applied such that all points
$|r_{ijk} - r_c| > radius$ are not computed. The sum over p extends over all
relevant pairs of gaussians (which are called task in the code).
the kernel computes the polynomials and the gaussian then sums the result
back to the grid.
the coefficients $C^p_{\alpha\beta\gamma}$ are computed by
calculate_coefficients. We only keep the non zero elements to same memory.
*/
template <typename T, typename T3, bool distributed__, bool orthorhombic_>
__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 task;
load_task(params, &task);
__shared__ smem_task_reduced<T, T3> task;
// Check if radius is below the resolution of the grid.
if (2.0 * task.radius < task.dh_max) {
return; // nothing to do
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
return;
fill_smem_task_reduced(dev_, dev_.first_task + blockIdx.x, task);
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
// Alloc shared memory.
extern __shared__ T coefs_[];
T *coef_ =
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset];
__shared__ T dh_[9], dh_inv_[9];
if (tid < 9) {
// matrix from lattice coordinates to cartesian coordinates
dh_[tid] = dev_.dh_[tid];
// matrix from cartesian coordinates to lattice coordinates.
dh_inv_[tid] = dev_.dh_inv_[tid];
}
// Allot dynamic shared memory.
extern __shared__ double shared_memory[];
double *smem_cab = &shared_memory[params->smem_cab_offset];
double *smem_alpha = &shared_memory[params->smem_alpha_offset];
double *smem_cxyz = &shared_memory[params->smem_cxyz_offset];
for (int i = tid; i < ncoset(6); i += blockDim.x * blockDim.y * blockDim.z)
coefs_[i] = coef_[i];
// Allocate Cab from global memory if it does not fit into shared memory.
cab_store cab = {.data = NULL, .n1 = task.n1};
if (params->smem_cab_length < task.n1 * task.n2) {
cab.data = malloc_cab(&task);
} else {
cab.data = smem_cab;
if (tid == 0) {
// 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);
}
}
}
__syncthreads();
zero_cab(&cab, task.n1 * task.n2);
compute_alpha(&task, smem_alpha);
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];
block_to_cab<IS_FUNC_AB>(params, &task, &cab);
cab_to_cxyz(&task, smem_alpha, &cab, smem_cxyz);
cxyz_to_grid(params, &task, smem_cxyz, params->grid);
if (z2 < 0)
z2 += dev_.grid_full_size_[0];
if (params->smem_cab_length < task.n1 * task.n2) {
free_cab(cab.data);
if (distributed__) {
// check if the point is within the window
if (task.apply_border_mask) {
// this test is only relevant when the grid is split over several mpi
// ranks. in that case we take only the points contributing to local
// part of the grid.
if ((z2 < task.window_shift.z) || (z2 > task.window_size.z)) {
continue;
}
}
}
// compute the coordinates of the point in atomic coordinates
T kremain;
short int ymin = 0;
short 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;
}
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];
if (distributed__) {
if (task.apply_border_mask) {
// check if the point is within the window
if ((y2 < task.window_shift.y) || (y2 > task.window_size.y)) {
continue;
}
}
}
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;
}
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];
if (distributed__) {
if (task.apply_border_mask) {
// check if the point is within the window (only true or false
// when using mpi) otherwise MPI=1 always true
if ((x2 < task.window_shift.x) || (x2 > task.window_size.x)) {
continue;
}
}
}
// I make no distinction between orthorhombic and non orthorhombic
// 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));
}
const T r3x2 = r3.x * r3.x;
const T r3y2 = r3.y * r3.y;
const T r3z2 = r3.z * r3.z;
if (distributed__) {
// 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)) &&
(!orthorhombic_ || task.apply_border_mask))
continue;
} else {
// we do not need to do this test for the orthorhombic case
if ((!orthorhombic_) &&
((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)))
continue;
}
// allow computation of the address in parallel to starting the
// computations
T res = coefs_[0];
if (task.lp >= 1) {
res += coefs_[1] * r3.x;
res += coefs_[2] * r3.y;
res += coefs_[3] * r3.z;
}
const T r3xy = r3.x * r3.y;
const T r3xz = r3.x * r3.z;
const T r3yz = r3.y * r3.z;
if (task.lp >= 2) {
res += coefs_[4] * r3x2;
res += coefs_[5] * r3xy;
res += coefs_[6] * r3xz;
res += coefs_[7] * r3y2;
res += coefs_[8] * r3yz;
res += coefs_[9] * r3z2;
}
if (task.lp >= 3) {
res += r3x2 *
(coefs_[10] * r3.x + coefs_[11] * r3.y + coefs_[12] * r3.z);
res += r3.x * (coefs_[13] * r3y2 + coefs_[15] * r3z2);
res += coefs_[14] * r3xy * r3.z;
res += r3y2 * (coefs_[16] * r3.y + coefs_[17] * r3.z);
res += r3z2 * (coefs_[18] * r3.y + coefs_[19] * r3.z);
}
if (task.lp >= 4) {
res += r3x2 *
(coefs_[20] * r3x2 + coefs_[21] * r3xy + coefs_[22] * r3xz +
coefs_[23] * r3y2 + coefs_[24] * r3yz + coefs_[25] * r3z2);
res += r3y2 *
(coefs_[26] * r3xy + coefs_[27] * r3xz + coefs_[30] * r3y2 +
coefs_[31] * r3yz + coefs_[32] * r3z2);
res += r3z2 * (coefs_[28] * r3xy + coefs_[29] * r3xz +
coefs_[33] * r3yz + coefs_[34] * r3z2);
}
if (task.lp >= 5) {
const T r3x4 = r3x2 * r3x2;
const T r3y4 = r3y2 * r3y2;
const T r3z4 = r3z2 * r3z2;
res += r3x4 * (coefs_[35] * r3.x + // x^5
coefs_[36] * r3.y + // x^4 y
coefs_[37] * r3.z); // x^4 z
res += r3x2 * (r3.x * (coefs_[38] * r3y2 + // x^3 y^2
coefs_[39] * r3yz + // x^3 y z
coefs_[40] * r3z2) + // x^3 z^2
r3y2 * (coefs_[41] * r3.y + // x^2 y^3
coefs_[42] * r3.z) + // x^2 y^2 z
r3z2 * (coefs_[43] * r3.y + // x^2 y z^2
coefs_[44] * r3.z)); // x^2 z^2 z
res += r3.x * (coefs_[45] * r3y4 + // x y^4
r3y2 * (coefs_[46] * r3yz + // x y^3 z
coefs_[47] * r3z2) + // x y^2 z^2
r3z2 * (coefs_[48] * r3yz + // x y z^3
coefs_[49] * r3z2)); // x z^4
res += r3y2 * (r3y2 * (coefs_[50] * r3.y + // y^5
coefs_[51] * r3.z) + // y^4 z
r3z2 * (coefs_[52] * r3.y + // y^3 z^2
coefs_[53] * r3.z)); // y^2 z^3
res += r3z4 * (coefs_[54] * r3.y + // y z^4
coefs_[55] * r3.z); // z^5
if (task.lp >= 6) {
res += r3x4 * (coefs_[56] * r3x2 + // x^6
coefs_[57] * r3xy + // x^ 5 y
coefs_[58] * r3xz + // x^ 5 z
coefs_[59] * r3y2 + // x^4 y^2
coefs_[60] * r3yz + // x^4 yz
coefs_[61] * r3z2); // x^4 z^2
res += r3x2 * (coefs_[62] * r3y2 * r3xy + // x^3 y^3
coefs_[63] * r3y2 * r3xz + // x^3 y^2 z
coefs_[64] * r3xy * r3z2 + // x^3 y z^2
coefs_[65] * r3z2 * r3xz + // x^3 z^3
coefs_[66] * r3y4 + // x^2 y^4
coefs_[67] * r3y2 * r3yz + // x^2 y^3 z
coefs_[68] * r3y2 * r3z2 + // x^2 y^2 z^2
coefs_[69] * r3yz * r3z2 + // x^2 y z^3
coefs_[70] * r3z4); // x^2 z^4
res += r3y2 * r3z2 *
(coefs_[73] * r3xy + // x y^3 z^2
coefs_[74] * r3xz + // x y^2 z^3
coefs_[80] * r3yz + // y^3 z^3
coefs_[81] * r3z2); // y^2 z^4
res += r3y4 * (coefs_[71] * r3xy + // x y^5
coefs_[72] * r3xz + // x y^4 z
coefs_[77] * r3y2 + // y^6
coefs_[78] * r3yz + // y^5 z
coefs_[79] * r3z2); // y^4 z^2
res += r3z4 * (coefs_[75] * r3xy + // x y z^4
coefs_[76] * r3xz + // x z^5
coefs_[82] * r3yz + // y z^5
coefs_[83] * r3z2); // z^6
}
if (task.lp >= 7) {
for (int ic = ncoset(6); ic < ncoset(task.lp); ic++) {
T tmp1 = coef_[ic];
auto &co = coset_inv[ic];
T tmp = 1.0;
for (int po = 0; po < (co.l[2] >> 1); po++)
tmp *= r3z2;
if (co.l[2] & 0x1)
tmp *= r3.z;
for (int po = 0; po < (co.l[1] >> 1); po++)
tmp *= r3y2;
if (co.l[1] & 0x1)
tmp *= r3.y;
for (int po = 0; po < (co.l[0] >> 1); po++)
tmp *= r3x2;
if (co.l[0] & 0x1)
tmp *= r3.x;
res += tmp * tmp1;
}
}
}
res *= exp(-(r3x2 + r3y2 + r3z2) * task.zetp);
atomicAdd(dev_.ptr_dev[1] +
(z2 * dev_.grid_local_size_[1] + y2) *
dev_.grid_local_size_[2] +
x2,
res);
}
}
}
}
/*******************************************************************************
* \brief Specialized Cuda kernel that can only collocate GRID_FUNC_AB.
* \author Ole Schuett
******************************************************************************/
__global__ static void collocate_kernel_density(const kernel_params params) {
collocate_kernel<true>(&params);
}
/*******************************************************************************
* \brief Cuda kernel that can collocate any function, ie. GRID_FUNC_*.
* \author Ole Schuett
******************************************************************************/
__global__ static void collocate_kernel_anyfunc(const kernel_params params) {
collocate_kernel<false>(&params);
}
/*******************************************************************************
* \brief Launches the Cuda kernel that collocates all tasks of one grid level.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_collocate_one_grid_level(
const grid_gpu_task_list *task_list, const int first_task,
const int last_task, const enum grid_func func,
const grid_gpu_layout *layout, const offloadStream_t stream,
const double *pab_blocks_dev, double *grid_dev, int *lp_diff) {
void context_info::collocate_one_grid_level(const int level,
const enum grid_func func,
int *lp_diff) {
if (number_of_tasks_per_level_[level] == 0)
return;
// Compute max angular momentum.
const prepare_ldiffs ldiffs = prepare_get_ldiffs(func);
*lp_diff = ldiffs.la_max_diff + ldiffs.lb_max_diff; // for reporting stats
const int la_max = task_list->lmax + ldiffs.la_max_diff;
const int lb_max = task_list->lmax + ldiffs.lb_max_diff;
const int lp_max = la_max + lb_max;
const int ntasks = last_task - first_task + 1;
if (ntasks == 0) {
return; // Nothing to do and lp_diff already set.
}
const ldiffs_value ldiffs = prepare_get_ldiffs(func);
smem_parameters smem_params(ldiffs, lmax());
*lp_diff = smem_params.lp_diff();
init_constant_memory();
// Small Cab blocks are stored in shared mem, larger ones in global memory.
const int CAB_SMEM_LIMIT = ncoset(5) * ncoset(5); // = 56 * 56 = 3136
// Compute required shared memory.
const int alpha_len = 3 * (lb_max + 1) * (la_max + 1) * (lp_max + 1);
const int cxyz_len = ncoset(lp_max);
const int cab_len = imin(CAB_SMEM_LIMIT, ncoset(lb_max) * ncoset(la_max));
const size_t smem_per_block =
(alpha_len + cxyz_len + cab_len) * sizeof(double);
// kernel parameters
kernel_params params;
params.smem_cab_length = cab_len;
params.smem_cab_offset = 0;
params.smem_alpha_offset = params.smem_cab_offset + cab_len;
params.smem_cxyz_offset = params.smem_alpha_offset + alpha_len;
params.first_task = first_task;
kernel_params params = set_kernel_parameters(level, smem_params);
params.func = func;
params.grid = grid_dev;
params.la_min_diff = ldiffs.la_min_diff;
params.lb_min_diff = ldiffs.lb_min_diff;
params.la_max_diff = ldiffs.la_max_diff;
params.lb_max_diff = ldiffs.lb_max_diff;
params.tasks = task_list->tasks_dev;
params.pab_blocks = pab_blocks_dev;
memcpy(params.dh, layout->dh, 9 * sizeof(double));
memcpy(params.dh_inv, layout->dh_inv, 9 * sizeof(double));
memcpy(params.npts_global, layout->npts_global, 3 * sizeof(int));
memcpy(params.npts_local, layout->npts_local, 3 * sizeof(int));
memcpy(params.shift_local, layout->shift_local, 3 * sizeof(int));
// Launch !
const int nblocks = ntasks;
const dim3 threads_per_block(4, 4, 4);
if (func == GRID_FUNC_AB) {
collocate_kernel_density<<<nblocks, threads_per_block, smem_per_block,
stream>>>(params);
calculate_coefficients<double, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
smem_params.smem_per_block(), level_streams[level]>>>(params);
} else {
collocate_kernel_anyfunc<<<nblocks, threads_per_block, smem_per_block,
stream>>>(params);
calculate_coefficients<double, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
smem_params.smem_per_block(), level_streams[level]>>>(params);
}
OFFLOAD_CHECK(offloadGetLastError());
}
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
// EOF
if (grid_[level].is_distributed()) {
if (grid_[level].is_orthorhombic())
collocate_kernel<double, double3, true, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
else
collocate_kernel<double, double3, true, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
} else {
if (grid_[level].is_orthorhombic())
collocate_kernel<double, double3, false, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
else
collocate_kernel<double, double3, false, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
}
}
} // namespace rocm_backend

35
src/grid/gpu/grid_gpu_collocate.h
View file

@ -1,35 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#ifndef GRID_GPU_COLLOCATE_H
#define GRID_GPU_COLLOCATE_H
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#include "grid_gpu_task_list.h"
#ifdef __cplusplus
extern "C" {
#endif
/*******************************************************************************
* \brief Launches the Cuda kernel that collocates all tasks of one grid level.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_collocate_one_grid_level(
const grid_gpu_task_list *task_list, const int first_task,
const int last_task, const enum grid_func func,
const grid_gpu_layout *layout, const offloadStream_t stream,
const double *pab_blocks_dev, double *grid_dev, int *lp_diff);
#ifdef __cplusplus
}
#endif
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#endif
// EOF

36
src/grid/hip/grid_hip_context.cu → src/grid/gpu/grid_gpu_context.cu
View file

@ -7,16 +7,15 @@
/*
* Authors :
- Dr Mathieu Taillefumier (ETH Zurich / CSCS)
- Mathieu Taillefumier (ETH Zurich / CSCS)
- Advanced Micro Devices, Inc.
- Ole Schuett
*/
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <hip/hip_runtime_api.h>
#include <iostream>
#include "../../offload/offload_library.h"
@ -26,10 +25,10 @@ extern "C" {
#include "../common/grid_library.h"
}
#include "grid_hip_context.h"
#include "grid_hip_internal_header.h"
#include "grid_gpu_context.h"
#include "grid_gpu_internal_header.h"
#include "grid_hip_task_list.h"
#include "grid_gpu_task_list.h"
#if defined(_OMP_H)
#error "OpenMP should not be used in .cu files to accommodate HIP."
@ -39,7 +38,7 @@ extern "C" {
* \brief Allocates a task list for the GPU backend.
* See grid_ctx.h for details.
******************************************************************************/
extern "C" void grid_hip_create_task_list(
extern "C" void grid_gpu_create_task_list(
const bool ortho, const int ntasks, const int nlevels, const int natoms,
const int nkinds, const int nblocks, const int *block_offsets,
const double *atom_positions, const int *atom_kinds,
@ -49,7 +48,7 @@ extern "C" void grid_hip_create_task_list(
const int *border_mask_list, const int *block_num_list,
const double *radius_list, const double *rab_list, const int *npts_global,
const int *npts_local, const int *shift_local, const int *border_width,
const double *dh, const double *dh_inv, void *ptr) {
const double *dh, const double *dh_inv, grid_gpu_task_list **ptr) {
rocm_backend::context_info **ctx_out = (rocm_backend::context_info **)ptr;
// Select GPU device.
@ -367,7 +366,7 @@ extern "C" void grid_hip_create_task_list(
/*******************************************************************************
* \brief destroy a context
******************************************************************************/
extern "C" void grid_hip_free_task_list(void *ptr) {
extern "C" void grid_gpu_free_task_list(grid_gpu_task_list *ptr) {
rocm_backend::context_info *ctx = (rocm_backend::context_info *)ptr;
// Select GPU device.
@ -381,7 +380,7 @@ extern "C" void grid_hip_free_task_list(void *ptr) {
/*******************************************************************************
* \brief Collocate all tasks of in given list onto given grids.
******************************************************************************/
extern "C" void grid_hip_collocate_task_list(const void *ptr,
extern "C" void grid_gpu_collocate_task_list(const grid_gpu_task_list *ptr,
const enum grid_func func,
const int nlevels,
const offload_buffer *pab_blocks,
@ -438,10 +437,10 @@ extern "C" void grid_hip_collocate_task_list(const void *ptr,
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) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_HIP,
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_COLLOCATE_ORTHO, count);
} else {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_HIP,
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_COLLOCATE_GENERAL, count);
}
}
@ -458,8 +457,8 @@ extern "C" void grid_hip_collocate_task_list(const void *ptr,
* \brief Integrate all tasks of in given list onto given grids.
* See grid_ctx.h for details.
******************************************************************************/
extern "C" void grid_hip_integrate_task_list(
const void *ptr, const bool compute_tau, const int nlevels,
extern "C" void grid_gpu_integrate_task_list(
const grid_gpu_task_list *ptr, const bool compute_tau, const int nlevels,
const offload_buffer *pab_blocks, const offload_buffer **grids,
offload_buffer *hab_blocks, double *forces, double *virial) {
@ -523,10 +522,10 @@ extern "C" void grid_hip_integrate_task_list(
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) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_HIP,
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_INTEGRATE_ORTHO, count);
} else {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_HIP,
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_INTEGRATE_GENERAL, count);
}
}
@ -535,8 +534,7 @@ extern "C" void grid_hip_integrate_task_list(
// need to wait for all streams to finish
for (int level = 0; level < ctx->nlevels; level++) {
if (ctx->number_of_tasks_per_level_[level])
ctx->synchronize(ctx->level_streams[level]);
ctx->synchronize(ctx->level_streams[level]);
}
// computing the hab coefficients does not depend on the number of grids so we
// can run these calculations on the main stream
@ -552,5 +550,3 @@ extern "C" void grid_hip_integrate_task_list(
ctx->synchronize(ctx->main_stream);
}
#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)

16
src/grid/hip/grid_hip_context.h → src/grid/gpu/grid_gpu_context.h
View file

@ -11,10 +11,14 @@
- Advanced Micro Devices, Inc.
*/
#ifndef GRID_HIP_CONTEXT_H
#define GRID_HIP_CONTEXT_H
#ifndef GRID_GPU_CONTEXT_H
#define GRID_GPU_CONTEXT_H
#ifdef __OFFLOAD_HIP
#include <hip/hip_runtime_api.h>
#else
#include <cuda_runtime.h>
#endif
#include <vector>
extern "C" {
@ -442,7 +446,8 @@ public:
~context_info() { clear(); }
void clear() {
hipSetDevice(device_id_);
offload_set_chosen_device(device_id_);
offload_activate_chosen_device();
tasks_dev.reset();
block_offsets_dev.reset();
coef_dev_.reset();
@ -533,7 +538,10 @@ public:
offloadDeviceSynchronize();
}
void set_device() { hipSetDevice(device_id_); }
void set_device() {
offload_set_chosen_device(device_id_);
offload_activate_chosen_device();
}
void collocate_one_grid_level(const int level, const enum grid_func func,
int *lp_diff);

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

File diff suppressed because it is too large Load diff

36
src/grid/gpu/grid_gpu_integrate.h
View file

@ -1,36 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#ifndef GRID_GPU_INTEGRATE_H
#define GRID_GPU_INTEGRATE_H
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#include "grid_gpu_task_list.h"
#ifdef __cplusplus
extern "C" {
#endif
/*******************************************************************************
* \brief Launches the Cuda kernel that integrates all tasks of one grid level.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_integrate_one_grid_level(
const grid_gpu_task_list *task_list, const int first_task,
const int last_task, const bool compute_tau, const grid_gpu_layout *layout,
const offloadStream_t stream, const double *pab_blocks_dev,
const double *grid_dev, double *hab_blocks_dev, double *forces_dev,
double *virial_dev, int *lp_diff);
#ifdef __cplusplus
}
#endif
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#endif
// EOF

20
src/grid/hip/grid_hip_internal_header.h → src/grid/gpu/grid_gpu_internal_header.h
View file

@ -11,12 +11,11 @@
- Advanced Micro Devices, Inc.
*/
#ifndef GRID_HIP_INTERNAL_HEADER_H
#define GRID_HIP_INTERNAL_HEADER_H
#ifndef GRID_GPU_INTERNAL_HEADER_H
#define GRID_GPU_INTERNAL_HEADER_H
#include <algorithm>
#include <assert.h>
#include <hip/hip_runtime.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
@ -29,7 +28,7 @@ extern "C" {
#include "../common/grid_constants.h"
}
#include "grid_hip_context.h"
#include "grid_gpu_context.h"
namespace rocm_backend {
@ -243,10 +242,8 @@ inline static void init_constant_memory() {
}
}
}
hipError_t error =
hipMemcpyToSymbol(coset_inv, &coset_inv_host, sizeof(coset_inv_host), 0,
hipMemcpyHostToDevice);
assert(error == hipSuccess);
offloadMemcpyToSymbol(coset_inv, &coset_inv_host, sizeof(coset_inv_host));
// Binomial coefficient
int binomial_coef_host[19][19] = {
@ -276,11 +273,8 @@ inline static void init_constant_memory() {
6188, 2380, 680, 136, 17, 1, 0},
{1, 18, 153, 816, 3060, 8568, 18564, 31824, 43758, 48620, 43758, 31824,
18564, 8568, 3060, 816, 153, 18, 1}};
error =
hipMemcpyToSymbol(binomial_coef, &binomial_coef_host[0][0],
sizeof(binomial_coef_host), 0, hipMemcpyHostToDevice);
assert(error == hipSuccess);
offloadMemcpyToSymbol(binomial_coef, &binomial_coef_host[0][0],
sizeof(binomial_coef_host));
initialized = true;
}

4
src/grid/hip/grid_hip_prepare_pab.h → src/grid/gpu/grid_gpu_prepare_pab.h
View file

@ -12,6 +12,9 @@
- Advanced Micro Devices, Inc.
*/
#ifndef GRID_GPU_PROCESS_PAB_H
#define GRID_GPU_PROCESS_PAB_H
#include "../common/grid_constants.h"
#include <stdbool.h>
#include <stdio.h>
@ -404,3 +407,4 @@ inline ldiffs_value prepare_get_ldiffs(const enum grid_func func) {
return ldiffs;
}
} // namespace rocm_backend
#endif

6
src/grid/hip/grid_hip_process_vab.h → src/grid/gpu/grid_gpu_process_vab.h
View file

@ -12,10 +12,10 @@
- Advanced Micro Devices, Inc.
*/
#ifndef GRID_HIP_PROCESS_VAB_H
#define GRID_HIP_PROCESS_VAB_H
#ifndef GRID_GPU_PROCESS_VAB_H
#define GRID_GPU_PROCESS_VAB_H
#include "grid_hip_internal_header.h"
#include "grid_gpu_internal_header.h"
/* taken from ../common/grid_process_vab.h but added template parameter */
namespace rocm_backend {

556
src/grid/gpu/grid_gpu_task_list.cu
View file

@ -1,556 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#include "grid_gpu_collocate.h"
#include "grid_gpu_integrate.h"
#include "grid_gpu_task_list.h"
#include "../../offload/offload_library.h"
#include "../../offload/offload_mempool.h"
#include "../common/grid_common.h"
#include "../common/grid_constants.h"
#include "../common/grid_library.h"
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#if defined(_OMP_H)
#error "OpenMP should not be used in .cu files to accommodate HIP."
#endif
/*******************************************************************************
* \brief Create a single task and precompute as much as possible.
* \author Ole Schuett
******************************************************************************/
static void
create_tasks(const bool orthorhombic, const int ntasks,
const int block_offsets[], const double atom_positions[][3],
const int atom_kinds[], const grid_basis_set *basis_sets[],
const int level_list[], const int iatom_list[],
const int jatom_list[], const int iset_list[],
const int jset_list[], const int ipgf_list[],
const int jpgf_list[], const int border_mask_list[],
const int block_num_list[], const double radius_list[],
const double rab_list[][3], const int npts_local[][3],
const int shift_local[][3], const int border_width[][3],
const double dh[][3][3], const double dh_inv[][3][3],
const double *sphis_dev[], grid_gpu_task tasks[]) {
for (int itask = 0; itask < ntasks; itask++) {
grid_gpu_task *task = &tasks[itask];
task->iatom = iatom_list[itask] - 1;
task->jatom = jatom_list[itask] - 1;
const int iset = iset_list[itask] - 1;
const int jset = jset_list[itask] - 1;
const int ipgf = ipgf_list[itask] - 1;
const int jpgf = jpgf_list[itask] - 1;
const int ikind = atom_kinds[task->iatom] - 1;
const int jkind = atom_kinds[task->jatom] - 1;
const grid_basis_set *ibasis = basis_sets[ikind];
const grid_basis_set *jbasis = basis_sets[jkind];
const int level = level_list[itask] - 1;
const int border_mask = border_mask_list[itask];
task->use_orthorhombic_kernel = (orthorhombic && border_mask == 0);
task->zeta = ibasis->zet[iset * ibasis->maxpgf + ipgf];
task->zetb = jbasis->zet[jset * jbasis->maxpgf + jpgf];
task->zetp = task->zeta + task->zetb;
const double f = task->zetb / task->zetp;
task->rab2 = 0.0;
for (int i = 0; i < 3; i++) {
task->rab[i] = rab_list[itask][i];
task->rab2 += task->rab[i] * task->rab[i];
task->ra[i] = atom_positions[task->iatom][i];
task->rb[i] = task->ra[i] + task->rab[i];
task->rp[i] = task->ra[i] + task->rab[i] * f;
}
// center in grid coords, gp = MATMUL(dh_inv, rp)
for (int i = 0; i < 3; i++) {
task->gp[i] = dh_inv[level][0][i] * task->rp[0] +
dh_inv[level][1][i] * task->rp[1] +
dh_inv[level][2][i] * task->rp[2];
}
// resolution of the grid
task->dh_max = 0.0;
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++) {
task->dh_max = fmax(task->dh_max, fabs(dh[level][i][j]));
}
}
task->radius = radius_list[itask];
task->radius2 = task->radius * task->radius;
task->prefactor = exp(-task->zeta * f * task->rab2);
task->off_diag_twice = (task->iatom == task->jatom) ? 1.0 : 2.0;
// angular momentum range of basis set
task->la_max_basis = ibasis->lmax[iset];
task->lb_max_basis = jbasis->lmax[jset];
task->la_min_basis = ibasis->lmin[iset];
task->lb_min_basis = jbasis->lmin[jset];
// size of entire spherical basis
task->nsgfa = ibasis->nsgf;
task->nsgfb = jbasis->nsgf;
// size of spherical set
task->nsgf_seta = ibasis->nsgf_set[iset];
task->nsgf_setb = jbasis->nsgf_set[jset];
// strides of the sphi transformation matrices
task->maxcoa = ibasis->maxco;
task->maxcob = jbasis->maxco;
// start of spherical set within the basis
const int sgfa = ibasis->first_sgf[iset] - 1;
const int sgfb = jbasis->first_sgf[jset] - 1;
// start of exponent within the cartesian set
const int o1 = ipgf * ncoset(task->la_max_basis);
const int o2 = jpgf * ncoset(task->lb_max_basis);
// transformations from contracted spherical to primitiv carthesian basis
task->sphia = &sphis_dev[ikind][sgfa * task->maxcoa + o1];
task->sphib = &sphis_dev[jkind][sgfb * task->maxcob + o2];
// Locate current matrix block within the buffer.
const int block_num = block_num_list[itask] - 1;
task->block_transposed = (task->iatom > task->jatom);
const int block_offset = block_offsets[block_num];
const int subblock_offset = (task->block_transposed)
? sgfa * task->nsgfb + sgfb
: sgfb * task->nsgfa + sgfa;
task->ab_block_offset = block_offset + subblock_offset;
// Stuff for the ortho kernel ----------------------------------------------
if (orthorhombic) {
// Discretize the radius.
const double drmin =
fmin(dh[level][0][0], fmin(dh[level][1][1], dh[level][2][2]));
const int imr = imax(1, (int)ceil(task->radius / drmin));
task->disr_radius = drmin * imr;
// Position of the gaussian product:
// this is the actual definition of the position on the grid
// i.e. a point rp(:) gets here grid coordinates
// MODULO(rp(:)/dr(:),npts_global(:))+1
// hence (0.0,0.0,0.0) in real space is rsgrid%lb on the rsgrid in Fortran
// and (1,1,1) on grid here in C.
//
// cubecenter(:) = FLOOR(MATMUL(dh_inv, rp))
for (int i = 0; i < 3; i++) {
const int cubecenter = floor(task->gp[i]);
task->cube_center_shifted[i] = cubecenter - shift_local[level][i];
task->cube_offset[i] = cubecenter * dh[level][i][i] - task->rp[i];
}
}
// Stuff for the general kernel --------------------------------------------
//
// get the min max indices that contain at least the cube that contains a
// sphere around rp of radius radius if the cell is very non-orthogonal this
// implies that many useless points are included this estimate can be
// improved (i.e. not box but sphere should be used)
for (int i = 0; i < 3; i++) {
task->index_min[i] = INT_MAX;
task->index_max[i] = INT_MIN;
}
for (int i = -1; i <= 1; i++) {
for (int j = -1; j <= 1; j++) {
for (int k = -1; k <= 1; k++) {
const double x = task->rp[0] + i * task->radius;
const double y = task->rp[1] + j * task->radius;
const double z = task->rp[2] + k * task->radius;
for (int idir = 0; idir < 3; idir++) {
const double resc = dh_inv[level][0][idir] * x +
dh_inv[level][1][idir] * y +
dh_inv[level][2][idir] * z;
task->index_min[idir] =
imin(task->index_min[idir], (int)floor(resc));
task->index_max[idir] =
imax(task->index_max[idir], (int)ceil(resc));
}
}
}
}
// Defaults for border_mask == 0.
task->bounds_i[0] = 0;
task->bounds_i[1] = npts_local[level][0] - 1;
task->bounds_j[0] = 0;
task->bounds_j[1] = npts_local[level][1] - 1;
task->bounds_k[0] = 0;
task->bounds_k[1] = npts_local[level][2] - 1;
// See also rs_find_node() in task_list_methods.F.
// If the bit is set then we need to exclude the border in that direction.
if (border_mask & (1 << 0)) {
task->bounds_i[0] += border_width[level][0];
}
if (border_mask & (1 << 1)) {
task->bounds_i[1] -= border_width[level][0];
}
if (border_mask & (1 << 2)) {
task->bounds_j[0] += border_width[level][1];
}
if (border_mask & (1 << 3)) {
task->bounds_j[1] -= border_width[level][1];
}
if (border_mask & (1 << 4)) {
task->bounds_k[0] += border_width[level][2];
}
if (border_mask & (1 << 5)) {
task->bounds_k[1] -= border_width[level][2];
}
}
}
/*******************************************************************************
* \brief Allocates a task list for the GPU backend.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_create_task_list(
const bool orthorhombic, const int ntasks, const int nlevels,
const int natoms, const int nkinds, const int nblocks,
const int block_offsets[], const double atom_positions[][3],
const int atom_kinds[], const grid_basis_set *basis_sets[],
const int level_list[], const int iatom_list[], const int jatom_list[],
const int iset_list[], const int jset_list[], const int ipgf_list[],
const int jpgf_list[], const int border_mask_list[],
const int block_num_list[], const double radius_list[],
const double rab_list[][3], const int npts_global[][3],
const int npts_local[][3], const int shift_local[][3],
const int border_width[][3], const double dh[][3][3],
const double dh_inv[][3][3], grid_gpu_task_list **task_list_out) {
// Select GPU device.
offload_activate_chosen_device();
if (*task_list_out != NULL) {
// This is actually an opportunity to reuse some buffers.
grid_gpu_free_task_list(*task_list_out);
}
grid_gpu_task_list *task_list =
(grid_gpu_task_list *)malloc(sizeof(grid_gpu_task_list));
task_list->orthorhombic = orthorhombic;
task_list->ntasks = ntasks;
task_list->nlevels = nlevels;
task_list->natoms = natoms;
task_list->nkinds = nkinds;
task_list->nblocks = nblocks;
// Store grid layouts.
size_t size = nlevels * sizeof(grid_gpu_layout);
task_list->layouts = (grid_gpu_layout *)malloc(size);
for (int level = 0; level < nlevels; level++) {
for (int i = 0; i < 3; i++) {
task_list->layouts[level].npts_global[i] = npts_global[level][i];
task_list->layouts[level].npts_local[i] = npts_local[level][i];
task_list->layouts[level].shift_local[i] = shift_local[level][i];
task_list->layouts[level].border_width[i] = border_width[level][i];
for (int j = 0; j < 3; j++) {
task_list->layouts[level].dh[i][j] = dh[level][i][j];
task_list->layouts[level].dh_inv[i][j] = dh_inv[level][i][j];
}
}
}
// Upload basis set sphi matrices to device.
task_list->sphis_dev = (double **)malloc(nkinds * sizeof(double *));
for (int i = 0; i < nkinds; i++) {
size = basis_sets[i]->nsgf * basis_sets[i]->maxco * sizeof(double);
task_list->sphis_dev[i] = (double *)offload_mempool_device_malloc(size);
offloadMemcpyHtoD(task_list->sphis_dev[i], basis_sets[i]->sphi, size);
}
size = ntasks * sizeof(grid_gpu_task);
grid_gpu_task *tasks_host = (grid_gpu_task *)malloc(size);
create_tasks(orthorhombic, ntasks, block_offsets, atom_positions, atom_kinds,
basis_sets, level_list, iatom_list, jatom_list, iset_list,
jset_list, ipgf_list, jpgf_list, border_mask_list,
block_num_list, radius_list, rab_list, npts_local, shift_local,
border_width, dh, dh_inv, (const double **)task_list->sphis_dev,
tasks_host);
task_list->tasks_dev = (grid_gpu_task *)offload_mempool_device_malloc(size);
offloadMemcpyHtoD(task_list->tasks_dev, tasks_host, size);
free(tasks_host);
// Count tasks per level.
size = nlevels * sizeof(int);
task_list->tasks_per_level = (int *)malloc(size);
memset(task_list->tasks_per_level, 0, size);
for (int i = 0; i < ntasks; i++) {
task_list->tasks_per_level[level_list[i] - 1]++;
assert(i == 0 || level_list[i] >= level_list[i - 1]); // expect ordered list
}
// Find largest angular momentum.
task_list->lmax = 0;
for (int ikind = 0; ikind < nkinds; ikind++) {
for (int iset = 0; iset < basis_sets[ikind]->nset; iset++) {
task_list->lmax = imax(task_list->lmax, basis_sets[ikind]->lmax[iset]);
}
}
// collect stats
memset(task_list->stats, 0, 2 * 20 * sizeof(int));
for (int itask = 0; itask < ntasks; itask++) {
const int iatom = iatom_list[itask] - 1;
const int jatom = jatom_list[itask] - 1;
const int ikind = atom_kinds[iatom] - 1;
const int jkind = atom_kinds[jatom] - 1;
const int iset = iset_list[itask] - 1;
const int jset = jset_list[itask] - 1;
const int la_max = basis_sets[ikind]->lmax[iset];
const int lb_max = basis_sets[jkind]->lmax[jset];
const int lp = imin(la_max + lb_max, 19);
const bool has_border_mask = (border_mask_list[itask] != 0);
task_list->stats[has_border_mask][lp]++;
}
// allocate main cuda stream
offloadStreamCreate(&task_list->main_stream);
// allocate one cuda stream per grid level
size = nlevels * sizeof(offloadStream_t);
task_list->level_streams = (offloadStream_t *)malloc(size);
for (int i = 0; i < nlevels; i++) {
offloadStreamCreate(&task_list->level_streams[i]);
}
// return newly created task list
*task_list_out = task_list;
}
/*******************************************************************************
* \brief Deallocates given task list, basis_sets have to be freed separately.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_free_task_list(grid_gpu_task_list *task_list) {
offload_mempool_device_free(task_list->tasks_dev);
offloadStreamDestroy(task_list->main_stream);
for (int i = 0; i < task_list->nlevels; i++) {
offloadStreamDestroy(task_list->level_streams[i]);
}
free(task_list->level_streams);
for (int i = 0; i < task_list->nkinds; i++) {
offload_mempool_device_free(task_list->sphis_dev[i]);
}
free(task_list->sphis_dev);
free(task_list->tasks_per_level);
free(task_list->layouts);
free(task_list);
}
/*******************************************************************************
* \brief Collocate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_collocate_task_list(const grid_gpu_task_list *task_list,
const enum grid_func func, const int nlevels,
const offload_buffer *pab_blocks,
offload_buffer *grids[]) {
// Select GPU device.
offload_activate_chosen_device();
// Upload blocks buffer using the main stream
offloadMemcpyAsyncHtoD(pab_blocks->device_buffer, pab_blocks->host_buffer,
pab_blocks->size, task_list->main_stream);
// record an event so the level streams can wait for the blocks to be uploaded
offloadEvent_t input_ready_event;
offloadEventCreate(&input_ready_event);
offloadEventRecord(input_ready_event, task_list->main_stream);
int lp_diff;
int first_task = 0;
assert(task_list->nlevels == nlevels);
for (int level = 0; level < task_list->nlevels; level++) {
const int last_task = first_task + task_list->tasks_per_level[level] - 1;
const offloadStream_t level_stream = task_list->level_streams[level];
const grid_gpu_layout *layout = &task_list->layouts[level];
offload_buffer *grid = grids[level];
// zero grid device buffer
offloadMemsetAsync(grid->device_buffer, 0, grid->size, level_stream);
// launch kernel, but only after blocks have arrived
offloadStreamWaitEvent(level_stream, input_ready_event);
grid_gpu_collocate_one_grid_level(
task_list, first_task, last_task, func, layout, level_stream,
pab_blocks->device_buffer, grid->device_buffer, &lp_diff);
first_task = last_task + 1;
}
// update counters while we wait for kernels to finish
for (int has_border_mask = 0; has_border_mask <= 1; has_border_mask++) {
for (int lp = 0; lp < 20; lp++) {
const int count = task_list->stats[has_border_mask][lp];
if (task_list->orthorhombic && !has_border_mask) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_COLLOCATE_ORTHO, count);
} else {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_COLLOCATE_GENERAL, count);
}
}
}
// download result from device to host.
for (int level = 0; level < task_list->nlevels; level++) {
offload_buffer *grid = grids[level];
offloadMemcpyAsyncDtoH(grid->host_buffer, grid->device_buffer, grid->size,
task_list->level_streams[level]);
}
// clean up
offloadEventDestroy(input_ready_event);
// wait for all the streams to finish
offloadDeviceSynchronize();
}
/*******************************************************************************
* \brief Integrate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_integrate_task_list(const grid_gpu_task_list *task_list,
const bool compute_tau, const int natoms,
const int nlevels,
const offload_buffer *pab_blocks,
const offload_buffer *grids[],
offload_buffer *hab_blocks,
double forces[][3], double virial[3][3]) {
// Select GPU device.
offload_activate_chosen_device();
// Prepare shared buffers using the main stream
double *forces_dev = NULL;
double *virial_dev = NULL;
double *pab_blocks_dev = NULL;
const size_t forces_size = 3 * natoms * sizeof(double);
const size_t virial_size = 9 * sizeof(double);
if (forces != NULL || virial != NULL) {
offloadMemcpyAsyncHtoD(pab_blocks->device_buffer, pab_blocks->host_buffer,
pab_blocks->size, task_list->main_stream);
pab_blocks_dev = pab_blocks->device_buffer;
}
if (forces != NULL) {
forces_dev = (double *)offload_mempool_device_malloc(forces_size);
offloadMemsetAsync(forces_dev, 0, forces_size, task_list->main_stream);
}
if (virial != NULL) {
virial_dev = (double *)offload_mempool_device_malloc(virial_size);
offloadMemsetAsync(virial_dev, 0, virial_size, task_list->main_stream);
}
// zero device hab blocks buffers
offloadMemsetAsync(hab_blocks->device_buffer, 0, hab_blocks->size,
task_list->main_stream);
// record event so other streams can wait for hab, pab, virial etc to be ready
offloadEvent_t input_ready_event;
offloadEventCreate(&input_ready_event);
offloadEventRecord(input_ready_event, task_list->main_stream);
int lp_diff;
int first_task = 0;
assert(task_list->nlevels == nlevels);
for (int level = 0; level < task_list->nlevels; level++) {
const int last_task = first_task + task_list->tasks_per_level[level] - 1;
const offloadStream_t level_stream = task_list->level_streams[level];
const grid_gpu_layout *layout = &task_list->layouts[level];
const offload_buffer *grid = grids[level];
// upload grid
offloadMemcpyAsyncHtoD(grid->device_buffer, grid->host_buffer, grid->size,
level_stream);
// launch kernel, but only after hab, pab, virial, etc are ready
offloadStreamWaitEvent(level_stream, input_ready_event);
grid_gpu_integrate_one_grid_level(
task_list, first_task, last_task, compute_tau, layout, level_stream,
pab_blocks_dev, grid->device_buffer, hab_blocks->device_buffer,
forces_dev, virial_dev, &lp_diff);
// Have main stream wait for level to complete before downloading results.
offloadEvent_t level_done_event;
offloadEventCreate(&level_done_event);
offloadEventRecord(level_done_event, level_stream);
offloadStreamWaitEvent(task_list->main_stream, level_done_event);
offloadEventDestroy(level_done_event);
first_task = last_task + 1;
}
// update counters while we wait for kernels to finish
for (int has_border_mask = 0; has_border_mask <= 1; has_border_mask++) {
for (int lp = 0; lp < 20; lp++) {
const int count = task_list->stats[has_border_mask][lp];
if (task_list->orthorhombic && !has_border_mask) {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_INTEGRATE_ORTHO, count);
} else {
grid_library_counter_add(lp + lp_diff, GRID_BACKEND_GPU,
GRID_INTEGRATE_GENERAL, count);
}
}
}
// download result from device to host using main stream.
offloadMemcpyAsyncDtoH(hab_blocks->host_buffer, hab_blocks->device_buffer,
hab_blocks->size, task_list->main_stream);
if (forces != NULL) {
offloadMemcpyAsyncDtoH(forces, forces_dev, forces_size,
task_list->main_stream);
}
if (virial != NULL) {
offloadMemcpyAsyncDtoH(virial, virial_dev, virial_size,
task_list->main_stream);
}
// wait for all the streams to finish
offloadDeviceSynchronize();
// clean up
offloadEventDestroy(input_ready_event);
if (forces != NULL) {
offload_mempool_device_free(forces_dev);
}
if (virial != NULL) {
offload_mempool_device_free(virial_dev);
}
}
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
// EOF

152
src/grid/gpu/grid_gpu_task_list.h
View file

@ -7,8 +7,7 @@
#ifndef GRID_GPU_TASK_LIST_H
#define GRID_GPU_TASK_LIST_H
#include "../../offload/offload_runtime.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#include <stdbool.h>
#ifdef __cplusplus
extern "C" {
@ -17,160 +16,49 @@ extern "C" {
#include "../../offload/offload_buffer.h"
#include "../common/grid_basis_set.h"
#include "../common/grid_constants.h"
#include <stdbool.h>
/*******************************************************************************
* \brief Internal representation of a task.
* \author Ole Schuett
******************************************************************************/
typedef struct grid_gpu_task_struct {
bool use_orthorhombic_kernel;
bool block_transposed;
double radius;
double radius2;
double ra[3];
double rb[3];
double rp[3];
double rab[3];
double gp[3];
double rab2;
double zeta;
double zetb;
double zetp;
double prefactor;
double off_diag_twice;
double dh_max;
// angular momentum range of basis set
int la_max_basis;
int lb_max_basis;
int la_min_basis;
int lb_min_basis;
// size of entire spherical basis
int nsgfa;
int nsgfb;
// size of spherical set
int nsgf_seta;
int nsgf_setb;
// strides of the sphi transformation matrices
int maxcoa;
int maxcob;
// offset of the pab and hab block relative to buffer pointer.
int ab_block_offset;
// atoms to which the forces and virial should be added
int iatom;
int jatom;
// pointers basis set matrices
const double *sphia;
const double *sphib;
// Stuff for the ortho kernel.
double disr_radius;
int cube_center_shifted[3];
double cube_offset[3];
// Stuff for the general kernel.
int index_min[3];
int index_max[3];
int bounds_i[2];
int bounds_j[2];
int bounds_k[2];
} grid_gpu_task;
/*******************************************************************************
* \brief Internal representation of a grid layout.
* \author Ole Schuett
******************************************************************************/
typedef struct {
int npts_global[3];
int npts_local[3];
int shift_local[3];
int border_width[3];
double dh[3][3];
double dh_inv[3][3];
} grid_gpu_layout;
/*******************************************************************************
* \brief Internal representation of a task list.
* \author Ole Schuett
******************************************************************************/
typedef struct {
bool orthorhombic;
int ntasks;
int nlevels;
int natoms;
int nkinds;
int nblocks;
grid_gpu_layout *layouts;
int *tasks_per_level;
offloadStream_t *level_streams;
offloadStream_t main_stream;
int lmax;
int stats[2][20]; // [has_border_mask][lp]
// device pointers
double **sphis_dev;
grid_gpu_task *tasks_dev;
} grid_gpu_task_list;
typedef void grid_gpu_task_list;
/*******************************************************************************
* \brief Allocates a task list for the GPU backend.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_create_task_list(
const bool orthorhombic, const int ntasks, const int nlevels,
const int natoms, const int nkinds, const int nblocks,
const int block_offsets[], const double atom_positions[][3],
const int atom_kinds[], const grid_basis_set *basis_sets[],
const int level_list[], const int iatom_list[], const int jatom_list[],
const int iset_list[], const int jset_list[], const int ipgf_list[],
const int jpgf_list[], const int border_mask_list[],
const int block_num_list[], const double radius_list[],
const double rab_list[][3], const int npts_global[][3],
const int npts_local[][3], const int shift_local[][3],
const int border_width[][3], const double dh[][3][3],
const double dh_inv[][3][3], grid_gpu_task_list **task_list);
const bool ortho, const int ntasks, const int nlevels, const int natoms,
const int nkinds, const int nblocks, const int *block_offsets,
const double *atom_positions, const int *atom_kinds,
const grid_basis_set **basis_sets, const int *level_list,
const int *iatom_list, const int *jatom_list, const int *iset_list,
const int *jset_list, const int *ipgf_list, const int *jpgf_list,
const int *border_mask_list, const int *block_num_list,
const double *radius_list, const double *rab_list, const int *npts_global,
const int *npts_local, const int *shift_local, const int *border_width,
const double *dh, const double *dh_inv, grid_gpu_task_list **ptr);
/*******************************************************************************
* \brief Deallocates given task list, basis_sets have to be freed separately.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_free_task_list(grid_gpu_task_list *task_list);
void grid_gpu_free_task_list(grid_gpu_task_list *ptr);
/*******************************************************************************
* \brief Collocate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_collocate_task_list(const grid_gpu_task_list *task_list,
void grid_gpu_collocate_task_list(const grid_gpu_task_list *ptr,
const enum grid_func func, const int nlevels,
const offload_buffer *pab_blocks,
offload_buffer *grids[]);
offload_buffer **grids);
/*******************************************************************************
* \brief Integrate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
* \author Ole Schuett
******************************************************************************/
void grid_gpu_integrate_task_list(const grid_gpu_task_list *task_list,
const bool compute_tau, const int natoms,
const int nlevels,
void grid_gpu_integrate_task_list(const grid_gpu_task_list *ptr,
const bool compute_tau, const int nlevels,
const offload_buffer *pab_blocks,
const offload_buffer *grids[],
offload_buffer *hab_blocks,
double forces[][3], double virial[3][3]);
const offload_buffer **grids,
offload_buffer *hab_blocks, double *forces,
double *virial);
#ifdef __cplusplus
}
#endif
#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#endif
// EOF

1
src/grid/grid_api.F
View file

@ -66,7 +66,6 @@ MODULE grid_api
INTEGER, PARAMETER, PUBLIC :: GRID_BACKEND_CPU = 12
INTEGER, PARAMETER, PUBLIC :: GRID_BACKEND_DGEMM = 13
INTEGER, PARAMETER, PUBLIC :: GRID_BACKEND_GPU = 14
INTEGER, PARAMETER, PUBLIC :: GRID_BACKEND_HIP = 15
PUBLIC :: grid_library_init, grid_library_finalize
PUBLIC :: grid_library_set_config, grid_library_print_stats

47
src/grid/grid_task_list.c
View file

@ -48,10 +48,7 @@ void grid_create_task_list(
// Resolve AUTO to a concrete backend.
if (config.backend == GRID_BACKEND_AUTO) {
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
task_list->backend = GRID_BACKEND_HIP;
#elif defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
task_list->backend = GRID_BACKEND_GPU;
#else
task_list->backend = GRID_BACKEND_CPU;
@ -104,31 +101,15 @@ void grid_create_task_list(
case GRID_BACKEND_GPU:
#if (defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID))
grid_gpu_create_task_list(
orthorhombic, ntasks, nlevels, natoms, nkinds, nblocks, block_offsets,
atom_positions, atom_kinds, basis_sets, level_list, iatom_list,
jatom_list, iset_list, jset_list, ipgf_list, jpgf_list,
border_mask_list, block_num_list, radius_list, rab_list, npts_global,
npts_local, shift_local, border_width, dh, dh_inv, &task_list->gpu);
#else
fprintf(stderr,
"Error: The GPU grid backend is not available. "
"Please re-compile with -D__OFFLOAD_CUDA or -D__OFFLOAD_HIP");
abort();
#endif
break;
case GRID_BACKEND_HIP:
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
grid_hip_create_task_list(
orthorhombic, ntasks, nlevels, natoms, nkinds, nblocks, block_offsets,
&atom_positions[0][0], atom_kinds, basis_sets, level_list, iatom_list,
jatom_list, iset_list, jset_list, ipgf_list, jpgf_list,
border_mask_list, block_num_list, radius_list, &rab_list[0][0],
&npts_global[0][0], &npts_local[0][0], &shift_local[0][0],
&border_width[0][0], &dh[0][0][0], &dh_inv[0][0][0], &task_list->hip);
&border_width[0][0], &dh[0][0][0], &dh_inv[0][0][0], &task_list->gpu);
#else
fprintf(stderr, "Error: The HIP grid backend is not available. "
"Please re-compile with -D__OFFLOAD_HIP");
fprintf(stderr, "Error: The GPU grid backend is not available. "
"Please re-compile with -D__OFFLOAD");
abort();
#endif
break;
@ -166,12 +147,6 @@ void grid_free_task_list(grid_task_list *task_list) {
task_list->gpu = NULL;
}
#endif
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
if (task_list->hip != NULL) {
grid_hip_free_task_list(task_list->hip);
task_list->hip = NULL;
}
#endif
free(task_list->npts_local);
free(task_list);
@ -214,12 +189,6 @@ void grid_collocate_task_list(const grid_task_list *task_list,
grid_gpu_collocate_task_list(task_list->gpu, func, nlevels, pab_blocks,
grids);
break;
#endif
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
case GRID_BACKEND_HIP:
grid_hip_collocate_task_list(task_list->hip, func, nlevels, pab_blocks,
grids);
break;
#endif
default:
printf("Error: Unknown grid backend: %i.\n", task_list->backend);
@ -299,13 +268,7 @@ void grid_integrate_task_list(
switch (task_list->backend) {
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
case GRID_BACKEND_GPU:
grid_gpu_integrate_task_list(task_list->gpu, compute_tau, natoms, nlevels,
pab_blocks, grids, hab_blocks, forces, virial);
break;
#endif
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
case GRID_BACKEND_HIP:
grid_hip_integrate_task_list(task_list->hip, compute_tau, nlevels,
grid_gpu_integrate_task_list(task_list->gpu, compute_tau, nlevels,
pab_blocks, grids, hab_blocks, &forces[0][0],
&virial[0][0]);
break;

4
src/grid/grid_task_list.h
View file

@ -16,7 +16,6 @@
#include "cpu/grid_cpu_task_list.h"
#include "dgemm/grid_dgemm_task_list.h"
#include "gpu/grid_gpu_task_list.h"
#include "hip/grid_hip_task_list.h"
#include "ref/grid_ref_task_list.h"
/*******************************************************************************
@ -32,9 +31,6 @@ typedef struct {
grid_dgemm_task_list *dgemm;
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID)
grid_gpu_task_list *gpu;
#endif
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
grid_hip_task_list *hip;
#endif
// more backends to be added here
} grid_task_list;

486
src/grid/hip/grid_hip_collocate.cu
View file

@ -1,486 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
/*
* Authors :
- Dr Mathieu Taillefumier (ETH Zurich / CSCS)
- Advanced Micro Devices, Inc.
*/
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
#include <algorithm>
#include <assert.h>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <hip/hip_runtime.h>
#include "grid_hip_internal_header.h"
#include "grid_hip_prepare_pab.h"
#if defined(_OMP_H)
#error "OpenMP should not be used in .cu files to accommodate HIP."
#endif
namespace rocm_backend {
/*******************************************************************************
* \brief Decontracts the subblock, going from spherical to cartesian harmonics.
******************************************************************************/
template <typename T, bool IS_FUNC_AB>
__device__ __inline__ void block_to_cab(const kernel_params &params,
const smem_task<T> &task, T *cab) {
// The spherical index runs over angular momentum and then over contractions.
// The cartesian index runs over exponents and then over angular momentum.
// 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);
for (int jco = task.first_cosetb + tid / 16; jco < task.ncosetb; jco += 4) {
for (int ico = task.first_coseta + (tid % 16); ico < task.ncoseta;
ico += 16) {
T pab_val = 0.0;
for (int i = 0; i < task.nsgf_setb; i++) {
const T sphib = task.sphib[i * task.maxcob + jco];
T tmp_val = 0.0;
for (int j = 0; j < task.nsgf_seta; j++) {
T block_val;
if (task.block_transposed) {
block_val =
task.pab_block[j * task.nsgfb + i] * task.off_diag_twice;
} else {
block_val =
task.pab_block[i * task.nsgfa + j] * task.off_diag_twice;
}
// const T sphia = task.sphia[j * task.maxcoa + ico];
tmp_val += block_val * task.sphia[j * task.maxcoa + ico];
}
pab_val += tmp_val * sphib;
}
if (IS_FUNC_AB) {
cab[jco * task.ncoseta + ico] += pab_val;
} else {
const auto a = coset_inv[ico];
const auto b = coset_inv[jco];
// remains the atomicAdd used in prepare_pab. Are they needed ?
prepare_pab(params.func, a, b, task.zeta, task.zetb, pab_val, task.n1,
cab);
}
}
}
}
template <typename T, bool IS_FUNC_AB>
__global__
__launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) {
__shared__ smem_task<T> task;
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
return;
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
fill_smem_task_coef(dev_, dev_.first_task + blockIdx.x, task);
extern __shared__ T shared_memory[];
// T *smem_cab = &shared_memory[dev_.smem_cab_offset];
T *smem_alpha = &shared_memory[dev_.smem_alpha_offset];
T *coef_ =
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset];
T *smem_cab =
&dev_.ptr_dev[6][dev_.tasks[dev_.first_task + blockIdx.x].cab_offset];
compute_alpha(task, smem_alpha);
for (int z = tid; z < task.n1 * task.n2;
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_);
}
/*
\brief compute the real space representation of an operator expressed in the
gaussian basis
this kernel does the following operation
n_{ijk} = \f[
\sum_{p,\alpha,\beta,\gamma} C^p_{\alpha\beta\gamma} X_{\alpha,i} Y_{\beta, j}
Z_{\gamma, k} \exp\left(- \eta (r_{ijk} - r_c) ^ 2\right)
]
where $X_{\alpha,i}, Y_{\beta, j}, Z_{\gamma, k}$ are polynomials of degree
$\alpha,\beta,\gamma$ and $r_{ijk}% a point (in cartesian coordinates) on a 3D
grid. C^p_{\alpha\beta\gamma} are also constrained such that 0 <= \alpha +
\beta + \gamma <= lmax. It means in practice that we need store (lmax + 1) *
(lamax + 2) * (lmax + 3) / 6 coefficients all the other coefficients are zero
to reduce computation, a spherical cutoff is applied such that all points
$|r_{ijk} - r_c| > radius$ are not computed. The sum over p extends over all
relevant pairs of gaussians (which are called task in the code).
the kernel computes the polynomials and the gaussian then sums the result
back to the grid.
the coefficients $C^p_{\alpha\beta\gamma}$ are computed by
calculate_coefficients. We only keep the non zero elements to same memory.
*/
template <typename T, typename T3, bool distributed__, bool orthorhombic_>
__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)
return;
fill_smem_task_reduced(dev_, dev_.first_task + blockIdx.x, task);
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
// Alloc shared memory.
extern __shared__ T coefs_[];
T *coef_ =
&dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset];
__shared__ T dh_[9], dh_inv_[9];
if (tid < 9) {
// matrix from lattice coordinates to cartesian coordinates
dh_[tid] = dev_.dh_[tid];
// matrix from cartesian coordinates to lattice coordinates.
dh_inv_[tid] = dev_.dh_inv_[tid];
}
for (int i = tid; i < ncoset(6); i += blockDim.x * blockDim.y * blockDim.z)
coefs_[i] = coef_[i];
if (tid == 0) {
// 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);
}
}
}
__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];
if (distributed__) {
// check if the point is within the window
if (task.apply_border_mask) {
// this test is only relevant when the grid is split over several mpi
// ranks. in that case we take only the points contributing to local
// part of the grid.
if ((z2 < task.window_shift.z) || (z2 > task.window_size.z)) {
continue;
}
}
}
// compute the coordinates of the point in atomic coordinates
T kremain;
short int ymin = 0;
short 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;
}
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];
if (distributed__) {
if (task.apply_border_mask) {
// check if the point is within the window
if ((y2 < task.window_shift.y) || (y2 > task.window_size.y)) {
continue;
}
}
}
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;
}
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];
if (distributed__) {
if (task.apply_border_mask) {
// check if the point is within the window (only true or false
// when using mpi) otherwise MPI=1 always true
if ((x2 < task.window_shift.x) || (x2 > task.window_size.x)) {
continue;
}
}
}
// I make no distinction between orthorhombic and non orthorhombic
// 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));
}
const T r3x2 = r3.x * r3.x;
const T r3y2 = r3.y * r3.y;
const T r3z2 = r3.z * r3.z;
if (distributed__) {
// 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)) &&
(!orthorhombic_ || task.apply_border_mask))
continue;
} else {
// we do not need to do this test for the orthorhombic case
if ((!orthorhombic_) &&
((task.radius * task.radius) <= (r3x2 + r3y2 + r3z2)))
continue;
}
// allow computation of the address in parallel to starting the
// computations
T res = coefs_[0];
if (task.lp >= 1) {
res += coefs_[1] * r3.x;
res += coefs_[2] * r3.y;
res += coefs_[3] * r3.z;
}
const T r3xy = r3.x * r3.y;
const T r3xz = r3.x * r3.z;
const T r3yz = r3.y * r3.z;
if (task.lp >= 2) {
res += coefs_[4] * r3x2;
res += coefs_[5] * r3xy;
res += coefs_[6] * r3xz;
res += coefs_[7] * r3y2;
res += coefs_[8] * r3yz;
res += coefs_[9] * r3z2;
}
if (task.lp >= 3) {
res += r3x2 *
(coefs_[10] * r3.x + coefs_[11] * r3.y + coefs_[12] * r3.z);
res += r3.x * (coefs_[13] * r3y2 + coefs_[15] * r3z2);
res += coefs_[14] * r3xy * r3.z;
res += r3y2 * (coefs_[16] * r3.y + coefs_[17] * r3.z);
res += r3z2 * (coefs_[18] * r3.y + coefs_[19] * r3.z);
}
if (task.lp >= 4) {
res += r3x2 *
(coefs_[20] * r3x2 + coefs_[21] * r3xy + coefs_[22] * r3xz +
coefs_[23] * r3y2 + coefs_[24] * r3yz + coefs_[25] * r3z2);
res += r3y2 *
(coefs_[26] * r3xy + coefs_[27] * r3xz + coefs_[30] * r3y2 +
coefs_[31] * r3yz + coefs_[32] * r3z2);
res += r3z2 * (coefs_[28] * r3xy + coefs_[29] * r3xz +
coefs_[33] * r3yz + coefs_[34] * r3z2);
}
if (task.lp >= 5) {
const T r3x4 = r3x2 * r3x2;
const T r3y4 = r3y2 * r3y2;
const T r3z4 = r3z2 * r3z2;
res += r3x4 * (coefs_[35] * r3.x + // x^5
coefs_[36] * r3.y + // x^4 y
coefs_[37] * r3.z); // x^4 z
res += r3x2 * (r3.x * (coefs_[38] * r3y2 + // x^3 y^2
coefs_[39] * r3yz + // x^3 y z
coefs_[40] * r3z2) + // x^3 z^2
r3y2 * (coefs_[41] * r3.y + // x^2 y^3
coefs_[42] * r3.z) + // x^2 y^2 z
r3z2 * (coefs_[43] * r3.y + // x^2 y z^2
coefs_[44] * r3.z)); // x^2 z^2 z
res += r3.x * (coefs_[45] * r3y4 + // x y^4
r3y2 * (coefs_[46] * r3yz + // x y^3 z
coefs_[47] * r3z2) + // x y^2 z^2
r3z2 * (coefs_[48] * r3yz + // x y z^3
coefs_[49] * r3z2)); // x z^4
res += r3y2 * (r3y2 * (coefs_[50] * r3.y + // y^5
coefs_[51] * r3.z) + // y^4 z
r3z2 * (coefs_[52] * r3.y + // y^3 z^2
coefs_[53] * r3.z)); // y^2 z^3
res += r3z4 * (coefs_[54] * r3.y + // y z^4
coefs_[55] * r3.z); // z^5
if (task.lp >= 6) {
res += r3x4 * (coefs_[56] * r3x2 + // x^6
coefs_[57] * r3xy + // x^ 5 y
coefs_[58] * r3xz + // x^ 5 z
coefs_[59] * r3y2 + // x^4 y^2
coefs_[60] * r3yz + // x^4 yz
coefs_[61] * r3z2); // x^4 z^2
res += r3x2 * (coefs_[62] * r3y2 * r3xy + // x^3 y^3
coefs_[63] * r3y2 * r3xz + // x^3 y^2 z
coefs_[64] * r3xy * r3z2 + // x^3 y z^2
coefs_[65] * r3z2 * r3xz + // x^3 z^3
coefs_[66] * r3y4 + // x^2 y^4
coefs_[67] * r3y2 * r3yz + // x^2 y^3 z
coefs_[68] * r3y2 * r3z2 + // x^2 y^2 z^2
coefs_[69] * r3yz * r3z2 + // x^2 y z^3
coefs_[70] * r3z4); // x^2 z^4
res += r3y2 * r3z2 *
(coefs_[73] * r3xy + // x y^3 z^2
coefs_[74] * r3xz + // x y^2 z^3
coefs_[80] * r3yz + // y^3 z^3
coefs_[81] * r3z2); // y^2 z^4
res += r3y4 * (coefs_[71] * r3xy + // x y^5
coefs_[72] * r3xz + // x y^4 z
coefs_[77] * r3y2 + // y^6
coefs_[78] * r3yz + // y^5 z
coefs_[79] * r3z2); // y^4 z^2
res += r3z4 * (coefs_[75] * r3xy + // x y z^4
coefs_[76] * r3xz + // x z^5
coefs_[82] * r3yz + // y z^5
coefs_[83] * r3z2); // z^6
}
if (task.lp >= 7) {
for (int ic = ncoset(6); ic < ncoset(task.lp); ic++) {
T tmp1 = coef_[ic];
auto &co = coset_inv[ic];
T tmp = 1.0;
for (int po = 0; po < (co.l[2] >> 1); po++)
tmp *= r3z2;
if (co.l[2] & 0x1)
tmp *= r3.z;
for (int po = 0; po < (co.l[1] >> 1); po++)
tmp *= r3y2;
if (co.l[1] & 0x1)
tmp *= r3.y;
for (int po = 0; po < (co.l[0] >> 1); po++)
tmp *= r3x2;
if (co.l[0] & 0x1)
tmp *= r3.x;
res += tmp * tmp1;
}
}
}
res *= exp(-(r3x2 + r3y2 + r3z2) * task.zetp);
atomicAdd(dev_.ptr_dev[1] +
(z2 * dev_.grid_local_size_[1] + y2) *
dev_.grid_local_size_[2] +
x2,
res);
}
}
}
}
/*******************************************************************************
* \brief Launches the Cuda kernel that collocates all tasks of one grid level.
******************************************************************************/
void context_info::collocate_one_grid_level(const int level,
const enum grid_func func,
int *lp_diff) {
if (number_of_tasks_per_level_[level] == 0)
return;
// Compute max angular momentum.
const ldiffs_value ldiffs = prepare_get_ldiffs(func);
smem_parameters smem_params(ldiffs, lmax());
*lp_diff = smem_params.lp_diff();
init_constant_memory();
// kernel parameters
kernel_params params = set_kernel_parameters(level, smem_params);
params.func = func;
// Launch !
const dim3 threads_per_block(4, 4, 4);
if (func == GRID_FUNC_AB) {
calculate_coefficients<double, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
smem_params.smem_per_block(), level_streams[level]>>>(params);
} else {
calculate_coefficients<double, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
smem_params.smem_per_block(), level_streams[level]>>>(params);
}
if (grid_[level].is_distributed()) {
if (grid_[level].is_orthorhombic())
collocate_kernel<double, double3, true, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
else
collocate_kernel<double, double3, true, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
} else {
if (grid_[level].is_orthorhombic())
collocate_kernel<double, double3, false, true>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
else
collocate_kernel<double, double3, false, false>
<<<number_of_tasks_per_level_[level], threads_per_block,
ncoset(6) * sizeof(double), level_streams[level]>>>(params);
}
}
} // namespace rocm_backend
#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)

768
src/grid/hip/grid_hip_integrate.cu
View file

@ -1,768 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
/*
* inspirations from the gpu backend
* Authors :
- Dr Mathieu Taillefumier (ETH Zurich / CSCS)
- Advanced Micro Devices, Inc.
*/
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
#include <algorithm>
#include <assert.h>
#include <hip/hip_runtime.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "grid_hip_context.h"
#include "grid_hip_internal_header.h"
#include "grid_hip_process_vab.h"
#ifdef __HIP_PLATFORM_NVIDIA__
#include <cooperative_groups.h>
#if CUDA_VERSION >= 11000
#include <cooperative_groups/reduce.h>
#endif
namespace cg = cooperative_groups;
#endif
#if defined(_OMP_H)
#error "OpenMP should not be used in .cu files to accommodate HIP."
#endif
namespace rocm_backend {
// do a warp reduction and return the final sum to thread_id = 0
template <typename T>
__device__ __inline__ T warp_reduce(T *table, 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.
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
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[];
// T *smem_cab = &shared_memory[dev_.smem_cab_offset];
T *smem_alpha = &shared_memory[dev_.smem_alpha_offset];
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 number_of_tasks = dev_.num_tasks_per_block_dev[blockIdx.x];
if (number_of_tasks == 0)
return;
T fa[3], fb[3];
T virial[9];
fa[0] = 0.0;
fa[1] = 0.0;
fa[2] = 0.0;
fb[0] = 0.0;
fb[1] = 0.0;
fb[2] = 0.0;
virial[0] = 0.0;
virial[1] = 0.0;
virial[2] = 0.0;
virial[3] = 0.0;
virial[4] = 0.0;
virial[5] = 0.0;
virial[6] = 0.0;
virial[7] = 0.0;
virial[8] = 0.0;
for (int tk = 0; tk < number_of_tasks; tk++) {
__shared__ smem_task<T> task;
const int task_id = dev_.task_sorted_by_blocks_dev[offset + tk];
if (dev_.tasks[task_id].skip_task)
continue;
fill_smem_task_coef(dev_, task_id, task);
T *coef_ = &dev_.ptr_dev[2][dev_.tasks[task_id].coef_offset];
T *smem_cab = &dev_.ptr_dev[6][dev_.tasks[task_id].cab_offset];
compute_alpha(task, smem_alpha);
__syncthreads();
cxyz_to_cab(task, smem_alpha, coef_, smem_cab);
__syncthreads();
for (int jco = task.first_cosetb; jco < task.ncosetb; jco++) {
const auto &b = coset_inv[jco];
for (int ico = task.first_coseta; ico < task.ncoseta; ico++) {
const auto &a = coset_inv[ico];
const T hab = get_hab<COMPUTE_TAU, T>(a, b, task.zeta, task.zetb,
task.n1, smem_cab);
__shared__ T shared_forces_a[3];
__shared__ T shared_forces_b[3];
__shared__ T shared_virial[9];
if (CALCULATE_FORCES) {
if (tid < 3) {
shared_forces_a[tid] = get_force_a<COMPUTE_TAU, T>(
a, b, tid, task.zeta, task.zetb, task.n1, smem_cab);
shared_forces_b[tid] = get_force_b<COMPUTE_TAU, T>(
a, b, tid, task.zeta, task.zetb, task.rab, task.n1, smem_cab);
}
if ((tid < 9) && (dev_.ptr_dev[5] != nullptr)) {
shared_virial[tid] =
get_virial_a<COMPUTE_TAU, T>(a, b, tid / 3, tid % 3, task.zeta,
task.zetb, task.n1, smem_cab) +
get_virial_b<COMPUTE_TAU, T>(a, b, tid / 3, tid % 3, task.zeta,
task.zetb, task.rab, task.n1,
smem_cab);
}
}
__syncthreads();
for (int i = tid / 8; i < task.nsgf_setb; i += 8) {
const T sphib = task.sphib[i * task.maxcob + jco];
for (int j = tid % 8; j < task.nsgf_seta; j += 8) {
const T sphia_times_sphib =
task.sphia[j * task.maxcoa + ico] * sphib;
// T block_val = hab * sphia_times_sphib;
if (CALCULATE_FORCES) {
T block_val1 = 0.0;
if (task.block_transposed) {
block_val1 =
task.pab_block[j * task.nsgfb + i] * task.off_diag_twice;
} else {
block_val1 =
task.pab_block[i * task.nsgfa + j] * task.off_diag_twice;
}
for (int k = 0; k < 3; k++) {
fa[k] += block_val1 * sphia_times_sphib * shared_forces_a[k];
// get_force_a<COMPUTE_TAU, T>(
// a, b, k, task.zeta, task.zetb, task.n1, smem_cab);
fb[k] += block_val1 * sphia_times_sphib * shared_forces_b[k];
// get_force_b<COMPUTE_TAU, T>(a, b, k, task.zeta, task.zetb,
// task.rab, task.n1, smem_cab);
}
if (dev_.ptr_dev[5] != nullptr) {
for (int k = 0; k < 3; k++) {
for (int l = 0; l < 3; l++) {
virial[3 * k + l] += shared_virial[3 * k + l] * block_val1 *
sphia_times_sphib;
}
}
}
}
// we can use shuffle_down if it exists for T
// these atomic operations are not needed since the blocks are
// updated by one single block thread. However the grid_miniapp will
// fail if not there.
if (task.block_transposed) {
task.hab_block[j * task.nsgfb + i] += hab * sphia_times_sphib;
// atomicAdd(task.hab_block + j * task.nsgfb + i, hab
// * sphia_times_sphib);
} else {
task.hab_block[i * task.nsgfa + j] += hab * sphia_times_sphib;
// atomicAdd(task.hab_block + i * task.nsgfa + j, hab *
// sphia_times_sphib);
}
}
}
__syncthreads();
}
}
}
if (CALCULATE_FORCES) {
const int task_id = dev_.task_sorted_by_blocks_dev[offset];
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 *sum = (T *)shared_memory;
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();
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();
if (tid == 0)
atomicAdd(forces_a + i, fa[i]);
__syncthreads();
sum[tid] = fb[i];
__syncthreads();
fb[i] = warp_reduce<T>(sum, tid);
__syncthreads();
if (tid == 0)
atomicAdd(forces_b + i, fb[i]);
__syncthreads();
}
}
}
/*******************************************************************************
* Cuda kernel for calcualting the coefficients of a potential (density,
etc...) for a given pair of gaussian basis sets
We compute the discretized version of the following integral
$$\int _\infty ^\infty V_{ijk} P^\alpha_iP^beta_j P ^ \gamma _ k Exp(- \eta
|r_{ijk} - r_c|^2)$$
where in practice the summation is over a finite domain. The discrete form has
this shape
$$
\sum_{ijk < dmoain} V_{ijk} (x_i - x_c)^\alpha (y_j - y_c)^\beta (z_k - z_c) ^
\gamma Exp(- \eta |r_{ijk} - r_c|^2)
$$
where $0 \le \alpha + \beta + \gamma \le lmax$
It is formely the same operation than collocate (from a discrete point of view)
but the implementation differ because of technical reasons.
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_,
int lbatch = 10>
__global__
__launch_bounds__(64) void integrate_kernel(const kernel_params dev_) {
if (dev_.tasks[dev_.first_task + blockIdx.x].skip_task)
return;
const int tid =
threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z);
__shared__ T dh_inv_[9];
__shared__ T dh_[9];
for (int d = tid; d < 9; d += blockDim.x * blockDim.y * blockDim.z)
dh_inv_[d] = dev_.dh_inv_[d];
for (int d = tid; d < 9; d += blockDim.x * blockDim.y * blockDim.z)
dh_[d] = dev_.dh_[d];
__shared__ smem_task_reduced<T, T3> task;
fill_smem_task_reduced(dev_, dev_.first_task + blockIdx.x, 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);
}
}
}
__syncthreads();
__shared__ T accumulator[lbatch][64];
__shared__ T sum[lbatch];
// we use a multi pass algorithm here because shared memory usage (or
// register) would become too high for high angular momentum
const short int size_loop =
(ncoset(task.lp) / lbatch + ((ncoset(task.lp) % lbatch) != 0)) * lbatch;
const short int length = ncoset(task.lp);
for (int ico = 0; ico < size_loop; ico += lbatch) {
#pragma unroll lbatch
for (int i = 0; i < lbatch; i++)
accumulator[i][tid] = 0.0;
__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];
if (distributed__) {
// known at compile time. Will be stripped away
if (task.apply_border_mask) {
/* check if the point is within the window */
if ((z2 < task.window_shift.z) || (z2 > task.window_size.z)) {
continue;
}
}
}
/* compute the coordinates of the point in atomic coordinates */
int ymin = 0;
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;
}
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];
if (distributed__) {
/* check if the point is within the window */
if (task.apply_border_mask) {
if ((y2 < task.window_shift.y) || (y2 > task.window_size.y)) {
continue;
}
}
}
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;
}
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];
if (distributed__) {
/* check if the point is within the window */
if (task.apply_border_mask) {
if ((x2 < task.window_shift.x) || (x2 > task.window_size.x)) {
continue;
}
}
}
// I make no distinction between orthorhombic and non orthorhombic
// 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));
}
// 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 (distributed__) {
if (((task.radius * task.radius) <=
(r3.x * r3.x + r3.y * r3.y + r3.z * r3.z)) &&
(!orthorhombic_ || 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;
}
}
// read the next point assuming that reading is non blocking untill
// 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_[1] + y2) *
dev_.grid_local_size_[2] +
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;
grid_value *= exp(-(r3x2 + r3y2 + r3z2) * task.zetp);
switch (ico / lbatch) {
case 0: {
accumulator[0][tid] += grid_value;
if (task.lp >= 1) {
accumulator[1][tid] += grid_value * r3.x;
accumulator[2][tid] += grid_value * r3.y;
accumulator[3][tid] += grid_value * r3.z;
}
if (task.lp >= 2) {
accumulator[4][tid] += grid_value * r3x2;
accumulator[5][tid] += grid_value * r3xy;
accumulator[6][tid] += grid_value * r3xz;
accumulator[7][tid] += grid_value * r3y2;
accumulator[8][tid] += grid_value * r3yz;
accumulator[9][tid] += grid_value * r3z2;
}
} break;
case 1: {
if (task.lp >= 3) {
accumulator[0][tid] += grid_value * r3x2 * r3.x;
accumulator[1][tid] += grid_value * r3x2 * r3.y;
accumulator[2][tid] += grid_value * r3x2 * r3.z;
accumulator[3][tid] += grid_value * r3.x * r3y2;
accumulator[4][tid] += grid_value * r3xy * r3.z;
accumulator[5][tid] += grid_value * r3.x * r3z2;
accumulator[6][tid] += grid_value * r3y2 * r3.y;
accumulator[7][tid] += grid_value * r3y2 * r3.z;
accumulator[8][tid] += grid_value * r3.y * r3z2;
accumulator[9][tid] += grid_value * r3z2 * r3.z;
}
} break;
case 2: {
if (task.lp >= 4) {
accumulator[0][tid] += grid_value * r3x2 * r3x2;
accumulator[1][tid] += grid_value * r3x2 * r3xy;
accumulator[2][tid] += grid_value * r3x2 * r3xz;
accumulator[3][tid] += grid_value * r3x2 * r3y2;
accumulator[4][tid] += grid_value * r3x2 * r3yz;
accumulator[5][tid] += grid_value * r3x2 * r3z2;
accumulator[6][tid] += grid_value * r3xy * r3y2;
accumulator[7][tid] += grid_value * r3xz * r3y2;
accumulator[8][tid] += grid_value * r3xy * r3z2;
accumulator[9][tid] += grid_value * r3xz * r3z2;
}
} break;
case 3: {
if (task.lp >= 4) {
accumulator[0][tid] += grid_value * r3y2 * r3y2;
accumulator[1][tid] += grid_value * r3y2 * r3yz;
accumulator[2][tid] += grid_value * r3y2 * r3z2;
accumulator[3][tid] += grid_value * r3yz * r3z2;
accumulator[4][tid] += grid_value * r3z2 * r3z2;
}
if (task.lp >= 5) {
accumulator[5][tid] += grid_value * r3x2 * r3x2 * r3.x;
accumulator[6][tid] += grid_value * r3x2 * r3x2 * r3.y;
accumulator[7][tid] += grid_value * r3x2 * r3x2 * r3.z;
accumulator[8][tid] += grid_value * r3x2 * r3.x * r3y2;
accumulator[9][tid] += grid_value * r3x2 * r3xy * r3.z;
}
} break;
case 4: {
accumulator[0][tid] += grid_value * r3x2 * r3.x * r3z2;
accumulator[1][tid] += grid_value * r3x2 * r3y2 * r3.y;
accumulator[2][tid] += grid_value * r3x2 * r3y2 * r3.z;
accumulator[3][tid] += grid_value * r3x2 * r3.y * r3z2;
accumulator[4][tid] += grid_value * r3x2 * r3z2 * r3.z;
accumulator[5][tid] += grid_value * r3.x * r3y2 * r3y2;
accumulator[6][tid] += grid_value * r3.x * r3y2 * r3yz;
accumulator[7][tid] += grid_value * r3.x * r3y2 * r3z2;
accumulator[8][tid] += grid_value * r3xy * r3z2 * r3.z;
accumulator[9][tid] += grid_value * r3.x * r3z2 * r3z2;
} break;
case 5: {
accumulator[0][tid] += grid_value * r3y2 * r3y2 * r3.y;
accumulator[1][tid] += grid_value * r3y2 * r3y2 * r3.z;
accumulator[2][tid] += grid_value * r3y2 * r3.y * r3z2;
accumulator[3][tid] += grid_value * r3y2 * r3z2 * r3.z;
accumulator[4][tid] += grid_value * r3.y * r3z2 * r3z2;
accumulator[5][tid] += grid_value * r3z2 * r3z2 * r3.z;
if (task.lp >= 6) {
accumulator[6][tid] += grid_value * r3x2 * r3x2 * r3x2; // x^6
accumulator[7][tid] += grid_value * r3x2 * r3x2 * r3xy; // x^5 y
accumulator[8][tid] += grid_value * r3x2 * r3x2 * r3xz; // x^5 z
accumulator[9][tid] += grid_value * r3x2 * r3x2 * r3y2; // x^4 y^2
}
} break;
case 6: {
accumulator[0][tid] += grid_value * r3x2 * r3x2 * r3yz; // x^4 y z
accumulator[1][tid] += grid_value * r3x2 * r3x2 * r3z2; // x^4 z^2
accumulator[2][tid] += grid_value * r3x2 * r3y2 * r3xy; // x^3 y^3
accumulator[3][tid] += grid_value * r3x2 * r3y2 * r3xz; // x^3 y^2 z
accumulator[4][tid] += grid_value * r3x2 * r3xy * r3z2; // x^3 y z^2
accumulator[5][tid] += grid_value * r3x2 * r3z2 * r3xz; // x^3 z^3
accumulator[6][tid] += grid_value * r3x2 * r3y2 * r3y2; // x^2 y^4
accumulator[7][tid] += grid_value * r3x2 * r3y2 * r3yz; // x^3 y^2 z
accumulator[8][tid] +=
grid_value * r3x2 * r3y2 * r3z2; // x^2 y^2 z^2
accumulator[9][tid] += grid_value * r3x2 * r3z2 * r3yz; // x^2 y z^3
} break;
case 7: {
accumulator[0][tid] += grid_value * r3x2 * r3z2 * r3z2; // x^2 z^4
accumulator[1][tid] += grid_value * r3y2 * r3y2 * r3xy; // x y^5
accumulator[2][tid] += grid_value * r3y2 * r3y2 * r3xz; // x y^4 z
accumulator[3][tid] += grid_value * r3y2 * r3xy * r3z2; // x y^3 z^2
accumulator[4][tid] += grid_value * r3y2 * r3z2 * r3xz; // x y^2 z^3
accumulator[5][tid] += grid_value * r3xy * r3z2 * r3z2; // x y z^4
accumulator[6][tid] += grid_value * r3z2 * r3z2 * r3xz; // x z^5
accumulator[7][tid] += grid_value * r3y2 * r3y2 * r3y2; // y^6
accumulator[8][tid] += grid_value * r3y2 * r3y2 * r3yz; // y^5 z
accumulator[9][tid] += grid_value * r3y2 * r3y2 * r3z2; // y^4 z^2
} break;
default:
for (int ic = 0; (ic < lbatch) && ((ic + ico) < length); ic++) {
auto &co = coset_inv[ic + ico];
T tmp = 1.0;
for (int po = 0; po < (co.l[2] >> 1); po++)
tmp *= r3z2;
if (co.l[2] & 0x1)
tmp *= r3.z;
for (int po = 0; po < (co.l[1] >> 1); po++)
tmp *= r3y2;
if (co.l[1] & 0x1)
tmp *= r3.y;
for (int po = 0; po < (co.l[0] >> 1); po++)
tmp *= r3x2;
if (co.l[0] & 0x1)
tmp *= r3.x;
accumulator[ic][tid] += tmp * grid_value;
}
break;
}
}
}
}
__syncthreads();
// we know there is only 1 wavefront in each block
// lbatch threads could reduce the values saved by all threads of the warp
// and save results do a shuffle_down by hand
for (int i = 0; i < min(length - ico, lbatch); i++) {
if (tid < 32) {
accumulator[i][tid] += accumulator[i][tid + 32];
}
__syncthreads();
if (tid < 16) {
accumulator[i][tid] += accumulator[i][tid + 16];
}
__syncthreads();
if (tid < 8) {
accumulator[i][tid] += accumulator[i][tid + 8];
}
__syncthreads();
if (tid < 4) {
accumulator[i][tid] += accumulator[i][tid + 4];
}
__syncthreads();
if (tid < 2) {
accumulator[i][tid] += accumulator[i][tid + 2];
}
__syncthreads();
if (tid == 0)
sum[i] = accumulator[i][0] + accumulator[i][1];
}
__syncthreads();
for (int i = tid; i < min(length - ico, lbatch); i += lbatch)
dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset + i +
ico] = sum[i];
__syncthreads();
// if (tid == 0)
// printf("%.15lf\n", sum[0]);
}
}
kernel_params
context_info::set_kernel_parameters(const int level,
const smem_parameters &smem_params) {
kernel_params params;
params.smem_cab_offset = smem_params.smem_cab_offset();
params.smem_alpha_offset = smem_params.smem_alpha_offset();
params.first_task = 0;
params.la_min_diff = smem_params.ldiffs().la_min_diff;
params.lb_min_diff = smem_params.ldiffs().lb_min_diff;
params.la_max_diff = smem_params.ldiffs().la_max_diff;
params.lb_max_diff = smem_params.ldiffs().lb_max_diff;
params.first_task = 0;
params.tasks = this->tasks_dev.data();
params.task_sorted_by_blocks_dev = task_sorted_by_blocks_dev.data();
params.sorted_blocks_offset_dev = sorted_blocks_offset_dev.data();
params.num_tasks_per_block_dev = this->num_tasks_per_block_dev_.data();
params.block_offsets = this->block_offsets_dev.data();
params.la_min_diff = smem_params.ldiffs().la_min_diff;
params.lb_min_diff = smem_params.ldiffs().lb_min_diff;
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();
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);
}
params.first_task = first_task_per_level_[level];
}
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.sphi_dev = this->sphi_dev.data();
return params;
}
/*******************************************************************************
* \brief Launches the Cuda kernel that integrates all tasks of one grid level.
******************************************************************************/
void context_info::integrate_one_grid_level(const int level, int *lp_diff) {
if (number_of_tasks_per_level_[level] == 0)
return;
assert(!calculate_virial || calculate_forces);
// Compute max angular momentum.
const ldiffs_value ldiffs =
process_get_ldiffs(calculate_forces, calculate_virial, compute_tau);
smem_parameters smem_params(ldiffs, lmax());
*lp_diff = smem_params.lp_diff();
init_constant_memory();
kernel_params params = set_kernel_parameters(level, smem_params);
/* WARNING : if you change the block size please be aware of that the number
* of warps is hardcoded when we do the block reduction in the integrate
* kernel. The number of warps should be explicitly indicated in the
* templating parameters or simply adjust the switch statements inside the
* integrate kernels */
const dim3 threads_per_block(4, 4, 4);
if (grid_[level].is_distributed()) {
if (grid_[level].is_orthorhombic()) {
integrate_kernel<double, double3, true, true, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);
} else {
integrate_kernel<double, double3, true, false, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);
}
} else {
if (grid_[level].is_orthorhombic()) {
integrate_kernel<double, double3, false, true, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);
} else {
integrate_kernel<double, double3, false, false, 10>
<<<number_of_tasks_per_level_[level], threads_per_block, 0,
level_streams[level]>>>(params);
}
}
}
void context_info::compute_hab_coefficients() {
assert(!calculate_virial || calculate_forces);
// Compute max angular momentum.
const ldiffs_value ldiffs =
process_get_ldiffs(calculate_forces, calculate_virial, compute_tau);
smem_parameters smem_params(ldiffs, lmax());
init_constant_memory();
kernel_params params = set_kernel_parameters(-1, smem_params);
/* WARNING if you change the block size. The number
* of warps is hardcoded when we do the block reduction in the integrate
* kernel. */
const dim3 threads_per_block(4, 4, 4);
if (!compute_tau && !calculate_forces) {
compute_hab_v2<double, double3, false, false>
<<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
this->main_stream>>>(params);
return;
}
if (!compute_tau && calculate_forces) {
compute_hab_v2<double, double3, false, true>
<<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
this->main_stream>>>(params);
return;
}
if (compute_tau && calculate_forces) {
compute_hab_v2<double, double3, true, true>
<<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
this->main_stream>>>(params);
}
if (compute_tau && !calculate_forces) {
compute_hab_v2<double, double3, true, false>
<<<this->nblocks, threads_per_block, smem_params.smem_per_block(),
this->main_stream>>>(params);
}
}
}; // namespace rocm_backend
#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)

67
src/grid/hip/grid_hip_task_list.h
View file

@ -1,67 +0,0 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: BSD-3-Clause */
/*----------------------------------------------------------------------------*/
#ifndef GRID_HIP_TASK_LIST_H
#define GRID_HIP_TASK_LIST_H
#if defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
#include <stdbool.h>
#ifdef __cplusplus
extern "C" {
#endif
#include "../../offload/offload_buffer.h"
#include "../common/grid_basis_set.h"
#include "../common/grid_constants.h"
typedef void grid_hip_task_list;
/*******************************************************************************
* \brief Allocates a task list for the GPU backend.
* See grid_task_list.h for details.
******************************************************************************/
void grid_hip_create_task_list(
const bool ortho, const int ntasks, const int nlevels, const int natoms,
const int nkinds, const int nblocks, const int *block_offsets,
const double *atom_positions, const int *atom_kinds,
const grid_basis_set **basis_sets, const int *level_list,
const int *iatom_list, const int *jatom_list, const int *iset_list,
const int *jset_list, const int *ipgf_list, const int *jpgf_list,
const int *border_mask_list, const int *block_num_list,
const double *radius_list, const double *rab_list, const int *npts_global,
const int *npts_local, const int *shift_local, const int *border_width,
const double *dh, const double *dh_inv, void *ptr);
/*******************************************************************************
* \brief Deallocates given task list, basis_sets have to be freed separately.
******************************************************************************/
void grid_hip_free_task_list(void *ptr);
/*******************************************************************************
* \brief Collocate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
******************************************************************************/
void grid_hip_collocate_task_list(const void *ptr, const enum grid_func func,
const int nlevels,
const offload_buffer *pab_blocks,
offload_buffer **grids);
/*******************************************************************************
* \brief Integrate all tasks of in given list onto given grids.
* See grid_task_list.h for details.
******************************************************************************/
void grid_hip_integrate_task_list(const void *ptr, const bool compute_tau,
const int nlevels,
const offload_buffer *pab_blocks,
const offload_buffer **grids,
offload_buffer *hab_blocks, double *forces,
double *virial);
#ifdef __cplusplus
}
#endif
#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID)
#endif

8
src/input_cp2k_global.F
View file

@ -48,7 +48,6 @@ MODULE input_cp2k_global
GRID_BACKEND_CPU,&
GRID_BACKEND_DGEMM,&
GRID_BACKEND_GPU,&
GRID_BACKEND_HIP,&
GRID_BACKEND_REF
USE input_constants, ONLY: &
bsse_run, callgraph_all, callgraph_master, callgraph_none, cell_opt_run, debug_run, &
@ -806,14 +805,13 @@ CONTAINS
description="Selects the backed used by the grid library.", &
default_i_val=GRID_BACKEND_AUTO, &
enum_i_vals=[GRID_BACKEND_AUTO, GRID_BACKEND_REF, GRID_BACKEND_CPU, &
GRID_BACKEND_DGEMM, GRID_BACKEND_GPU, GRID_BACKEND_HIP], &
enum_c_vals=s2a("AUTO", "REFERENCE", "CPU", "DGEMM", "GPU", "HIP"), &
GRID_BACKEND_DGEMM, GRID_BACKEND_GPU], &
enum_c_vals=s2a("AUTO", "REFERENCE", "CPU", "DGEMM", "GPU"), &
enum_desc=s2a("Let the grid library pick the backend automatically", &
"Reference backend implementation", &
"Optimized CPU backend", &
"Alternative CPU backend based on DGEMM", &
"GPU backend optimized for CUDA that also supports HIP", &
"HIP backend optimized for ROCm"))
"GPU backend optimized for NVIDIA and AMD GPU"))
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)

1
tools/precommit/check_file_properties.py
View file

@ -30,6 +30,7 @@ FLAG_EXCEPTIONS = (
r"ACC_OPENCL_.+",
r"FD_DEBUG",
r"GRID_DO_COLLOCATE",
r"GRID_GPU.*_H",
r"INTEL_MKL_VERSION",
r"LIBINT2_MAX_AM_eri",
r"LIBGRPP",