From 695da86d0cdae9080b95034b59edf65df6f95c10 Mon Sep 17 00:00:00 2001 From: Taillefumier Mathieu <29380261+mtaillefumier@users.noreply.github.com> Date: Sat, 16 May 2026 17:18:57 +0000 Subject: [PATCH] Grid gpu engines merge (#5206) Co-authored-by: Mathieu Taillefumier --- src/CMakeLists.txt | 11 +- src/environment.F | 4 - src/grid/common/grid_constants.h | 1 - src/grid/common/grid_library.c | 2 +- src/grid/gpu/grid_gpu_collint.h | 559 -------- src/grid/gpu/grid_gpu_collocate.cu | 742 ++++++----- src/grid/gpu/grid_gpu_collocate.h | 35 - .../grid_gpu_context.cu} | 36 +- .../grid_gpu_context.h} | 16 +- src/grid/gpu/grid_gpu_integrate.cu | 1168 ++++++++++------- src/grid/gpu/grid_gpu_integrate.h | 36 - .../grid_gpu_internal_header.h} | 20 +- .../grid_gpu_prepare_pab.h} | 4 + .../grid_gpu_process_vab.h} | 6 +- src/grid/gpu/grid_gpu_task_list.cu | 556 -------- src/grid/gpu/grid_gpu_task_list.h | 152 +-- src/grid/grid_api.F | 1 - src/grid/grid_task_list.c | 47 +- src/grid/grid_task_list.h | 4 - src/grid/hip/grid_hip_collocate.cu | 486 ------- src/grid/hip/grid_hip_integrate.cu | 768 ----------- src/grid/hip/grid_hip_task_list.h | 67 - src/input_cp2k_global.F | 8 +- tools/precommit/check_file_properties.py | 1 + 24 files changed, 1173 insertions(+), 3557 deletions(-) delete mode 100644 src/grid/gpu/grid_gpu_collint.h delete mode 100644 src/grid/gpu/grid_gpu_collocate.h rename src/grid/{hip/grid_hip_context.cu => gpu/grid_gpu_context.cu} (95%) rename src/grid/{hip/grid_hip_context.h => gpu/grid_gpu_context.h} (98%) delete mode 100644 src/grid/gpu/grid_gpu_integrate.h rename src/grid/{hip/grid_hip_internal_header.h => gpu/grid_gpu_internal_header.h} (98%) rename src/grid/{hip/grid_hip_prepare_pab.h => gpu/grid_gpu_prepare_pab.h} (99%) rename src/grid/{hip/grid_hip_process_vab.h => gpu/grid_gpu_process_vab.h} (99%) delete mode 100644 src/grid/gpu/grid_gpu_task_list.cu delete mode 100644 src/grid/hip/grid_hip_collocate.cu delete mode 100644 src/grid/hip/grid_hip_integrate.cu delete mode 100644 src/grid/hip/grid_hip_task_list.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index eae6055e39..8c8f473736 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -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}") diff --git a/src/environment.F b/src/environment.F index 15c2fe6193..d29badf30f 100644 --- a/src/environment.F +++ b/src/environment.F @@ -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" diff --git a/src/grid/common/grid_constants.h b/src/grid/common/grid_constants.h index db4b530e25..bff4a5db85 100644 --- a/src/grid/common/grid_constants.h +++ b/src/grid/common/grid_constants.h @@ -51,7 +51,6 @@ enum grid_backend { GRID_BACKEND_CPU = 12, GRID_BACKEND_DGEMM = 13, GRID_BACKEND_GPU = 14, - GRID_BACKEND_HIP = 15, }; #endif diff --git a/src/grid/common/grid_library.c b/src/grid/common/grid_library.c index b36bf61284..b0cf598082 100644 --- a/src/grid/common/grid_library.c +++ b/src/grid/common/grid_library.c @@ -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) diff --git a/src/grid/gpu/grid_gpu_collint.h b/src/grid/gpu/grid_gpu_collint.h deleted file mode 100644 index 227eb2b932..0000000000 --- a/src/grid/gpu/grid_gpu_collint.h +++ /dev/null @@ -1,559 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* SPDX-License-Identifier: BSD-3-Clause */ -/*----------------------------------------------------------------------------*/ - -#include "../../offload/offload_runtime.h" - -#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID) - -#include -#include -#include -#include -#include -#include -#include - -#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 = ¶ms->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 = ¶ms->pab_blocks[task->ab_block_offset]; - - // integrate -#if (!GRID_DO_COLLOCATE) - task->hab_block = ¶ms->hab_blocks[task->ab_block_offset]; - if (params->forces != NULL) { - task->forces_a = ¶ms->forces[3 * task->iatom]; - task->forces_b = ¶ms->forces[3 * task->jatom]; - } -#endif - } - __syncthreads(); // because of concurrent writes to task -} - -#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID) -// EOF diff --git a/src/grid/gpu/grid_gpu_collocate.cu b/src/grid/gpu/grid_gpu_collocate.cu index 4badf81739..49190ee2e4 100644 --- a/src/grid/gpu/grid_gpu_collocate.cu +++ b/src/grid/gpu/grid_gpu_collocate.cu @@ -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 #include -#include -#include -#include -#include -#include +#include +#include +#include +#include +#include -#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 -__device__ static void block_to_cab(const kernel_params *params, - const smem_task *task, cab_store *cab) { +template +__device__ __inline__ void block_to_cab(const kernel_params ¶ms, + const smem_task &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 -__device__ static void collocate_kernel(const kernel_params *params) { +template +__global__ +__launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) { + __shared__ smem_task 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(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 +__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 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(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(¶ms); -} - -/******************************************************************************* - * \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(¶ms); -} - /******************************************************************************* * \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<<>>(params); + calculate_coefficients + <<>>(params); } else { - collocate_kernel_anyfunc<<>>(params); + calculate_coefficients + <<>>(params); } - OFFLOAD_CHECK(offloadGetLastError()); -} -#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID) -// EOF + if (grid_[level].is_distributed()) { + if (grid_[level].is_orthorhombic()) + collocate_kernel + <<>>(params); + else + collocate_kernel + <<>>(params); + } else { + if (grid_[level].is_orthorhombic()) + collocate_kernel + <<>>(params); + else + collocate_kernel + <<>>(params); + } +} +} // namespace rocm_backend diff --git a/src/grid/gpu/grid_gpu_collocate.h b/src/grid/gpu/grid_gpu_collocate.h deleted file mode 100644 index 37828111e0..0000000000 --- a/src/grid/gpu/grid_gpu_collocate.h +++ /dev/null @@ -1,35 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 diff --git a/src/grid/hip/grid_hip_context.cu b/src/grid/gpu/grid_gpu_context.cu similarity index 95% rename from src/grid/hip/grid_hip_context.cu rename to src/grid/gpu/grid_gpu_context.cu index 544aba0179..8a6fae56ef 100644 --- a/src/grid/hip/grid_hip_context.cu +++ b/src/grid/gpu/grid_gpu_context.cu @@ -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 #include #include #include -#include #include #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) diff --git a/src/grid/hip/grid_hip_context.h b/src/grid/gpu/grid_gpu_context.h similarity index 98% rename from src/grid/hip/grid_hip_context.h rename to src/grid/gpu/grid_gpu_context.h index 5e5097a7be..2cfe469aa9 100644 --- a/src/grid/hip/grid_hip_context.h +++ b/src/grid/gpu/grid_gpu_context.h @@ -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 +#else +#include +#endif #include 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); diff --git a/src/grid/gpu/grid_gpu_integrate.cu b/src/grid/gpu/grid_gpu_integrate.cu index 718edc539b..2cf933f0e4 100644 --- a/src/grid/gpu/grid_gpu_integrate.cu +++ b/src/grid/gpu/grid_gpu_integrate.cu @@ -5,8 +5,13 @@ /* SPDX-License-Identifier: BSD-3-Clause */ /*----------------------------------------------------------------------------*/ -#include "../../offload/offload_runtime.h" -#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID) +/* + * inspirations from the gpu backend + * Authors : + - Mathieu Taillefumier (ETH Zurich / CSCS) + - Advanced Micro Devices, Inc. + - Ole Schuett +*/ #include #include @@ -16,371 +21,154 @@ #include #include -#define GRID_DO_COLLOCATE 0 -#include "../common/grid_common.h" -#include "grid_gpu_collint.h" -#include "grid_gpu_integrate.h" - -// This has to be included after grid_gpu_collint.h -#include "../common/grid_process_vab.h" +#include "grid_gpu_context.h" +#include "grid_gpu_internal_header.h" +#include "grid_gpu_process_vab.h" #if defined(_OMP_H) #error "OpenMP should not be used in .cu files to accommodate HIP." #endif -// Twenty registers are sufficient to integrate lp <= 3 with a single grid -// sweep. -#define GRID_N_CXYZ_REGISTERS 20 - -/******************************************************************************* - * \brief Add value to designated register without using dynamic indexing. - * Otherwise the array would be stored in local memory, which is slower. - * https://developer.nvidia.com/blog/fast-dynamic-indexing-private-arrays-cuda - * \author Ole Schuett - ******************************************************************************/ -__device__ static inline void -add_to_register(const double value, const int index, cxyz_store *store) { - switch (index) { - case 0: - store->regs[0] += value; - break; - case 1: - store->regs[1] += value; - break; - case 2: - store->regs[2] += value; - break; - case 3: - store->regs[3] += value; - break; - case 4: - store->regs[4] += value; - break; - case 5: - store->regs[5] += value; - break; - case 6: - store->regs[6] += value; - break; - case 7: - store->regs[7] += value; - break; - case 8: - store->regs[8] += value; - break; - case 9: - store->regs[9] += value; - break; - case 10: - store->regs[10] += value; - break; - case 11: - store->regs[11] += value; - break; - case 12: - store->regs[12] += value; - break; - case 13: - store->regs[13] += value; - break; - case 14: - store->regs[14] += value; - break; - case 15: - store->regs[15] += value; - break; - case 16: - store->regs[16] += value; - break; - case 17: - store->regs[17] += value; - break; - case 18: - store->regs[18] += value; - break; - case 19: - store->regs[19] += value; - break; +namespace rocm_backend { +// do a warp reduction and return the final sum to thread_id = 0 +template +__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]; } -} - -/******************************************************************************* - * \brief Integrate a single grid point with distance d{xyz} from center. - * \author Ole Schuett - ******************************************************************************/ -__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) { - - // Squared distance of point from center. - const double r2 = dx * dx + dy * dy + dz * dz; - const double gaussian = exp(-zetp * r2); - - // Loading throught read-only cache reduces register usage for some reason. - const double prefactor = __ldg(gridpoint) * gaussian; - - // Manually unrolled loops based on terms in coset_inv. - if (store->offset == 0) { - store->regs[0] += prefactor; - if (lp >= 1) { - store->regs[1] += prefactor * dx; - store->regs[2] += prefactor * dy; - store->regs[3] += prefactor * dz; - if (lp >= 2) { - store->regs[4] += prefactor * dx * dx; - store->regs[5] += prefactor * dx * dy; - store->regs[6] += prefactor * dx * dz; - store->regs[7] += prefactor * dy * dy; - store->regs[8] += prefactor * dy * dz; - store->regs[9] += prefactor * dz * dz; - if (lp >= 3) { - store->regs[10] += prefactor * dx * dx * dx; - store->regs[11] += prefactor * dx * dx * dy; - store->regs[12] += prefactor * dx * dx * dz; - store->regs[13] += prefactor * dx * dy * dy; - store->regs[14] += prefactor * dx * dy * dz; - store->regs[15] += prefactor * dx * dz * dz; - store->regs[16] += prefactor * dy * dy * dy; - store->regs[17] += prefactor * dy * dy * dz; - store->regs[18] += prefactor * dy * dz * dz; - store->regs[19] += prefactor * dz * dz * dz; - } - } - } - - } else if (store->offset == 20) { - store->regs[0] += prefactor * dx * dx * dx * dx; - store->regs[1] += prefactor * dx * dx * dx * dy; - store->regs[2] += prefactor * dx * dx * dx * dz; - store->regs[3] += prefactor * dx * dx * dy * dy; - store->regs[4] += prefactor * dx * dx * dy * dz; - store->regs[5] += prefactor * dx * dx * dz * dz; - store->regs[6] += prefactor * dx * dy * dy * dy; - store->regs[7] += prefactor * dx * dy * dy * dz; - store->regs[8] += prefactor * dx * dy * dz * dz; - store->regs[9] += prefactor * dx * dz * dz * dz; - store->regs[10] += prefactor * dy * dy * dy * dy; - store->regs[11] += prefactor * dy * dy * dy * dz; - store->regs[12] += prefactor * dy * dy * dz * dz; - store->regs[13] += prefactor * dy * dz * dz * dz; - store->regs[14] += prefactor * dz * dz * dz * dz; - if (lp >= 5) { - store->regs[15] += prefactor * dx * dx * dx * dx * dx; - store->regs[16] += prefactor * dx * dx * dx * dx * dy; - store->regs[17] += prefactor * dx * dx * dx * dx * dz; - store->regs[18] += prefactor * dx * dx * dx * dy * dy; - store->regs[19] += prefactor * dx * dx * dx * dy * dz; - } - } else if (store->offset == 40) { - store->regs[0] += prefactor * dx * dx * dx * dz * dz; - store->regs[1] += prefactor * dx * dx * dy * dy * dy; - store->regs[2] += prefactor * dx * dx * dy * dy * dz; - store->regs[3] += prefactor * dx * dx * dy * dz * dz; - store->regs[4] += prefactor * dx * dx * dz * dz * dz; - store->regs[5] += prefactor * dx * dy * dy * dy * dy; - store->regs[6] += prefactor * dx * dy * dy * dy * dz; - store->regs[7] += prefactor * dx * dy * dy * dz * dz; - store->regs[8] += prefactor * dx * dy * dz * dz * dz; - store->regs[9] += prefactor * dx * dz * dz * dz * dz; - store->regs[10] += prefactor * dy * dy * dy * dy * dy; - store->regs[11] += prefactor * dy * dy * dy * dy * dz; - store->regs[12] += prefactor * dy * dy * dy * dz * dz; - store->regs[13] += prefactor * dy * dy * dz * dz * dz; - store->regs[14] += prefactor * dy * dz * dz * dz * dz; - store->regs[15] += prefactor * dz * dz * dz * dz * dz; - if (lp >= 6) { - store->regs[16] += prefactor * dx * dx * dx * dx * dx * dx; - store->regs[17] += prefactor * dx * dx * dx * dx * dx * dy; - store->regs[18] += prefactor * dx * dx * dx * dx * dx * dz; - store->regs[19] += prefactor * dx * dx * dx * dx * dy * dy; - } - } else if (store->offset == 60) { - store->regs[0] += prefactor * dx * dx * dx * dx * dy * dz; - store->regs[1] += prefactor * dx * dx * dx * dx * dz * dz; - store->regs[2] += prefactor * dx * dx * dx * dy * dy * dy; - store->regs[3] += prefactor * dx * dx * dx * dy * dy * dz; - store->regs[4] += prefactor * dx * dx * dx * dy * dz * dz; - store->regs[5] += prefactor * dx * dx * dx * dz * dz * dz; - store->regs[6] += prefactor * dx * dx * dy * dy * dy * dy; - store->regs[7] += prefactor * dx * dx * dy * dy * dy * dz; - store->regs[8] += prefactor * dx * dx * dy * dy * dz * dz; - store->regs[9] += prefactor * dx * dx * dy * dz * dz * dz; - store->regs[10] += prefactor * dx * dx * dz * dz * dz * dz; - store->regs[11] += prefactor * dx * dy * dy * dy * dy * dy; - store->regs[12] += prefactor * dx * dy * dy * dy * dy * dz; - store->regs[13] += prefactor * dx * dy * dy * dy * dz * dz; - store->regs[14] += prefactor * dx * dy * dy * dz * dz * dz; - store->regs[15] += prefactor * dx * dy * dz * dz * dz * dz; - store->regs[16] += prefactor * dx * dz * dz * dz * dz * dz; - store->regs[17] += prefactor * dy * dy * dy * dy * dy * dy; - store->regs[18] += prefactor * dy * dy * dy * dy * dy * dz; - store->regs[19] += prefactor * dy * dy * dy * dy * dz * dz; - } else if (store->offset == 80) { - store->regs[0] += prefactor * dy * dy * dy * dz * dz * dz; - store->regs[1] += prefactor * dy * dy * dz * dz * dz * dz; - store->regs[2] += prefactor * dy * dz * dz * dz * dz * dz; - store->regs[3] += prefactor * dz * dz * dz * dz * dz * dz; - if (lp >= 7) { - store->regs[4] += prefactor * dx * dx * dx * dx * dx * dx * dx; - store->regs[5] += prefactor * dx * dx * dx * dx * dx * dx * dy; - store->regs[6] += prefactor * dx * dx * dx * dx * dx * dx * dz; - store->regs[7] += prefactor * dx * dx * dx * dx * dx * dy * dy; - store->regs[8] += prefactor * dx * dx * dx * dx * dx * dy * dz; - store->regs[9] += prefactor * dx * dx * dx * dx * dx * dz * dz; - store->regs[10] += prefactor * dx * dx * dx * dx * dy * dy * dy; - store->regs[11] += prefactor * dx * dx * dx * dx * dy * dy * dz; - store->regs[12] += prefactor * dx * dx * dx * dx * dy * dz * dz; - store->regs[13] += prefactor * dx * dx * dx * dx * dz * dz * dz; - store->regs[14] += prefactor * dx * dx * dx * dy * dy * dy * dy; - store->regs[15] += prefactor * dx * dx * dx * dy * dy * dy * dz; - store->regs[16] += prefactor * dx * dx * dx * dy * dy * dz * dz; - store->regs[17] += prefactor * dx * dx * dx * dy * dz * dz * dz; - store->regs[18] += prefactor * dx * dx * dx * dz * dz * dz * dz; - store->regs[19] += prefactor * dx * dx * dy * dy * dy * dy * dy; - } - - // Handle higher offsets, ie. values of lp. - } else { - for (int i = 0; i < GRID_N_CXYZ_REGISTERS; i++) { - double val = prefactor; - const orbital a = coset_inv[i + store->offset]; - 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; - } - add_to_register(val, i, store); - } + __syncthreads(); + if (tid < 16) { + table[tid] += table[tid + 16]; } -} - -/******************************************************************************* - * \brief Integrates the grid into coefficients C_xyz. - * \author Ole Schuett - ******************************************************************************/ -__device__ static void grid_to_cxyz(const kernel_params *params, - const smem_task *task, const double *grid, - double *cxyz) { - - // Atomics adds on shared memory are pretty slow. Hence, the coeffients are - // accumulated in registers while looping over the grid points. - // For larger values of lp we need to do multiple sweeps over the grid. - // Due to the higher register usage and the multiple sweeps, - // the integrate kernel runs about 70% slower than the collocate kernel. - for (int offset = 0; offset < ncoset(task->lp); - offset += GRID_N_CXYZ_REGISTERS) { - - double cxyz_regs[GRID_N_CXYZ_REGISTERS] = {0.0}; - cxyz_store store = {.regs = cxyz_regs, .offset = offset}; - - if (task->use_orthorhombic_kernel) { - ortho_cxyz_to_grid(params, task, &store, grid); - } else { - general_cxyz_to_grid(params, task, &store, grid); - } - - // Add register values to coefficients stored in shared memory. -#pragma unroll // avoid dynamic indexing of registers - for (int i = 0; i < GRID_N_CXYZ_REGISTERS; i++) { - if (i + offset < ncoset(task->lp)) { - atomicAddDouble(&cxyz[i + offset], cxyz_regs[i]); - } - } + __syncthreads(); + if (tid < 8) { + table[tid] += table[tid + 8]; } - __syncthreads(); // because of concurrent writes to cxyz -} - -/******************************************************************************* - * \brief Contracts the subblock, going from cartesian harmonics to spherical. - * \author Ole Schuett - ******************************************************************************/ -template -__device__ static void store_hab(const smem_task *task, const cab_store *cab) { - - // The spherical index runs over angular momentum and then over contractions. - // The carthesian index runs over exponents and then over angular momentum. - - // This is a double matrix product. Since the block can be quite large the - // two products are fused to conserve shared memory. - for (int i = threadIdx.x; i < task->nsgf_setb; i += blockDim.x) { - for (int j = threadIdx.y; j < task->nsgf_seta; j += blockDim.y) { - double block_val = 0.0; - const int jco_start = ncoset(task->lb_min_basis - 1) + threadIdx.z; - const int jco_end = ncoset(task->lb_max_basis); - for (int jco = jco_start; jco < jco_end; jco += blockDim.z) { - const orbital b = coset_inv[jco]; - const double sphib = task->sphib[i * task->maxcob + jco]; - const int ico_start = ncoset(task->la_min_basis - 1); - const int ico_end = ncoset(task->la_max_basis); - for (int ico = ico_start; ico < ico_end; ico++) { - const orbital a = coset_inv[ico]; - const double hab = - get_hab(a, b, task->zeta, task->zetb, cab, COMPUTE_TAU); - const double sphia = task->sphia[j * task->maxcoa + ico]; - block_val += hab * sphia * sphib; - } - } - if (task->block_transposed) { - atomicAddDouble(&task->hab_block[j * task->nsgfb + i], block_val); - } else { - atomicAddDouble(&task->hab_block[i * task->nsgfa + j], block_val); - } - } + __syncthreads(); + if (tid < 4) { + table[tid] += table[tid + 4]; } - __syncthreads(); // Not needed, but coalesced threads are nice. + __syncthreads(); + if (tid < 2) { + table[tid] += table[tid + 2]; + } + __syncthreads(); + return (table[0] + table[1]); } +// #endif -/******************************************************************************* - * \brief Adds contributions from cab to forces and virial. - * \author Ole Schuett - ******************************************************************************/ -template -__device__ static void store_forces_and_virial(const kernel_params *params, - const smem_task *task, - const cab_store *cab) { +template +__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]; - for (int i = threadIdx.x; i < task->nsgf_setb; i += blockDim.x) { - for (int j = threadIdx.y; j < task->nsgf_seta; j += blockDim.y) { - 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 int jco_start = ncoset(task->lb_min_basis - 1) + threadIdx.z; - const int jco_end = ncoset(task->lb_max_basis); - for (int jco = jco_start; jco < jco_end; jco += blockDim.z) { - const double sphib = task->sphib[i * task->maxcob + jco]; - const int ico_start = ncoset(task->la_min_basis - 1); - const int ico_end = ncoset(task->la_max_basis); - for (int ico = ico_start; ico < ico_end; ico++) { - const double sphia = task->sphia[j * task->maxcoa + ico]; - const double pabval = block_val * sphia * sphib; - const orbital b = coset_inv[jco]; - const orbital a = coset_inv[ico]; - for (int k = 0; k < 3; k++) { - const double force_a = - get_force_a(a, b, k, task->zeta, task->zetb, cab, COMPUTE_TAU); - atomicAddDouble(&task->forces_a[k], force_a * pabval); - const double force_b = get_force_b(a, b, k, task->zeta, task->zetb, - task->rab, cab, COMPUTE_TAU); - atomicAddDouble(&task->forces_b[k], force_b * pabval); + 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 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(a, b, task.zeta, task.zetb, + task.n1, smem_cab); + T *shared_forces_a = &shared_memory[0]; + T *shared_forces_b = &shared_memory[3]; + T *shared_virial = &shared_memory[6]; + + if (CALCULATE_FORCES) { + if (tid < 3) { + shared_forces_a[tid] = get_force_a( + a, b, tid, task.zeta, task.zetb, task.n1, smem_cab); + shared_forces_b[tid] = get_force_b( + a, b, tid, task.zeta, task.zetb, task.rab, task.n1, smem_cab); } - if (params->virial != NULL) { + if ((tid < 9) && (dev_.ptr_dev[5] != nullptr)) { + shared_virial[tid] = + get_virial_a(a, b, tid / 3, tid % 3, task.zeta, + task.zetb, task.n1, smem_cab) + + get_virial_b(a, b, tid / 3, tid % 3, task.zeta, + task.zetb, task.rab, task.n1, + smem_cab); + } + } + __syncthreads(); + T block_val1 = 0.0; + 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; + + if (CALCULATE_FORCES) { + if (task.block_transposed) { + block_val1 += task.pab_block[j * task.nsgfb + i] * + task.off_diag_twice * sphia_times_sphib; + } else { + block_val1 += task.pab_block[i * task.nsgfa + j] * + task.off_diag_twice * sphia_times_sphib; + } + } + + if (task.block_transposed) { + task.hab_block[j * task.nsgfb + i] += hab * sphia_times_sphib; + } else { + task.hab_block[i * task.nsgfa + j] += hab * sphia_times_sphib; + } + } + } + + if (CALCULATE_FORCES) { + for (int k = 0; k < 3; k++) { + fa[k] += block_val1 * shared_forces_a[k]; + fb[k] += block_val1 * shared_forces_b[k]; + } + if (dev_.ptr_dev[5] != nullptr) { for (int k = 0; k < 3; k++) { for (int l = 0; l < 3; l++) { - const double virial_a = get_virial_a( - a, b, k, l, task->zeta, task->zetb, cab, COMPUTE_TAU); - const double virial_b = - get_virial_b(a, b, k, l, task->zeta, task->zetb, task->rab, - cab, COMPUTE_TAU); - const double virial = pabval * (virial_a + virial_b); - atomicAddDouble(¶ms->virial[k * 3 + l], virial); + virial[3 * k + l] += shared_virial[3 * k + l] * block_val1; } } } @@ -388,183 +176,581 @@ __device__ static void store_forces_and_virial(const kernel_params *params, } } } - __syncthreads(); // Not needed, but coalesced threads are nice. -} -/******************************************************************************* - * \brief Initializes the cxyz matrix with zeros. - * \author Ole Schuett - ******************************************************************************/ -__device__ static void zero_cxyz(const smem_task *task, double *cxyz) { - if (threadIdx.z == 0 && threadIdx.y == 0) { - for (int i = threadIdx.x; i < ncoset(task->lp); i += blockDim.x) { - cxyz[i] = 0.0; + 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(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(sum, tid); + __syncthreads(); + if (tid == 0) + atomicAdd(forces_a + i, fa[i]); + __syncthreads(); + + sum[tid] = fb[i]; + __syncthreads(); + fb[i] = warp_reduce(sum, tid); + __syncthreads(); + if (tid == 0) + atomicAdd(forces_b + i, fb[i]); + __syncthreads(); } } - __syncthreads(); // because of concurrent writes to cxyz } /******************************************************************************* - * \brief Cuda kernel for integrating all tasks of one grid level. - * \author Ole Schuett - ******************************************************************************/ -template -__device__ static void integrate_kernel(const kernel_params *params) { + * Cuda kernel for calcualting the coefficients of a potential (density, + etc...) for a given pair of gaussian basis sets - // Copy task from global to shared memory and precompute some stuff. - __shared__ smem_task task; - load_task(params, &task); +We compute the discretized version of the following integral - // Check if radius is below the resolution of the grid. - if (2.0 * task.radius < task.dh_max) { - return; // nothing to do +$$\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 +__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 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(); - // 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]; + __shared__ T accumulator[lbatch][64]; + __shared__ T sum[lbatch]; - // 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; - } + // we use a multi pass algorithm here because shared memory usage (or + // register) would become too high for high angular momentum - zero_cab(&cab, task.n1 * task.n2); - compute_alpha(&task, smem_alpha); + 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; - zero_cxyz(&task, smem_cxyz); - grid_to_cxyz(params, &task, params->grid, smem_cxyz); - cab_to_cxyz(&task, smem_alpha, &cab, smem_cxyz); + __syncthreads(); - store_hab(&task, &cab); - if (CALCULATE_FORCES) { - store_forces_and_virial(params, &task, &cab); - } + 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 (params->smem_cab_length < task.n1 * task.n2) { - free_cab(cab.data); + 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 + if (tid < 32) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + accumulator[i][tid] += accumulator[i][tid + 32]; + } + } + __syncthreads(); + if (tid < 16) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + accumulator[i][tid] += accumulator[i][tid + 16]; + } + } + __syncthreads(); + if (tid < 8) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + accumulator[i][tid] += accumulator[i][tid + 8]; + } + } + __syncthreads(); + if (tid < 4) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + accumulator[i][tid] += accumulator[i][tid + 4]; + } + } + __syncthreads(); + if (tid < 2) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + accumulator[i][tid] += accumulator[i][tid + 2]; + } + } + __syncthreads(); + + if (tid == 0) { + for (int i = 0; i < min(length - ico, lbatch); i++) { + sum[i] = accumulator[i][0] + accumulator[i][1]; + } + } + __syncthreads(); + + if (tid < min(length - ico, lbatch)) + dev_.ptr_dev[2][dev_.tasks[dev_.first_task + blockIdx.x].coef_offset + + tid + ico] = sum[tid]; + __syncthreads(); } } -/******************************************************************************* - * \brief Specialized Cuda kernel for compute_tau=false & calculate_forces=false - * \author Ole Schuett - ******************************************************************************/ -__global__ static void grid_integrate_density(const kernel_params params) { - integrate_kernel(¶ms); -} +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; -/******************************************************************************* - * \brief Specialized Cuda kernel for compute_tau=true & calculate_forces=false. - * \author Ole Schuett - ******************************************************************************/ -__global__ static void grid_integrate_tau(const kernel_params params) { - integrate_kernel(¶ms); -} + params.la_min_diff = smem_params.ldiffs().la_min_diff; + params.lb_min_diff = smem_params.ldiffs().lb_min_diff; -/******************************************************************************* - * \brief Specialized Cuda kernel for compute_tau=false & calculate_forces=true. - * \author Ole Schuett - ******************************************************************************/ -__global__ static void -grid_integrate_density_forces(const kernel_params params) { - integrate_kernel(¶ms); -} + 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; -/******************************************************************************* - * \brief Specialized Cuda kernel for compute_tau=true & calculate_forces=true. - * \author Ole Schuett - ******************************************************************************/ -__global__ static void grid_integrate_tau_forces(const kernel_params params) { - integrate_kernel(¶ms); + 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. - * \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) { +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 bool calculate_forces = (forces_dev != NULL); - const bool calculate_virial = (virial_dev != NULL); - assert(!calculate_virial || calculate_forces); - const process_ldiffs ldiffs = + const ldiffs_value ldiffs = process_get_ldiffs(calculate_forces, calculate_virial, compute_tau); - *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. - } + 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 + kernel_params params = set_kernel_parameters(level, smem_params); - // 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); + /* 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 */ - // 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; - params.grid = grid_dev; - params.tasks = task_list->tasks_dev; - params.pab_blocks = pab_blocks_dev; - params.hab_blocks = hab_blocks_dev; - params.forces = forces_dev; - params.virial = virial_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; - 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)); + const dim3 threads_per_block(4, 4, 4); + if (grid_[level].is_distributed()) { + if (grid_[level].is_orthorhombic()) { + integrate_kernel + <<>>(params); + } else { + integrate_kernel + <<>>(params); + } + } else { + if (grid_[level].is_orthorhombic()) { + integrate_kernel + <<>>(params); + } else { + integrate_kernel + <<>>(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. */ - // Launch ! - const int nblocks = ntasks; const dim3 threads_per_block(4, 4, 4); if (!compute_tau && !calculate_forces) { - grid_integrate_density<<>>(params); - } else if (compute_tau && !calculate_forces) { - grid_integrate_tau<<>>( - params); - } else if (!compute_tau && calculate_forces) { - grid_integrate_density_forces<<>>(params); - } else if (compute_tau && calculate_forces) { - grid_integrate_tau_forces<<>>(params); + compute_hab_v2 + <<nblocks, threads_per_block, smem_params.smem_per_block(), + this->main_stream>>>(params); + return; } - OFFLOAD_CHECK(offloadGetLastError()); -} -#endif // defined(__OFFLOAD) && !defined(__NO_OFFLOAD_GRID) -// EOF + if (!compute_tau && calculate_forces) { + compute_hab_v2 + <<nblocks, threads_per_block, smem_params.smem_per_block(), + this->main_stream>>>(params); + return; + } + + if (compute_tau && calculate_forces) { + compute_hab_v2 + <<nblocks, threads_per_block, smem_params.smem_per_block(), + this->main_stream>>>(params); + } + + if (compute_tau && !calculate_forces) { + compute_hab_v2 + <<nblocks, threads_per_block, smem_params.smem_per_block(), + this->main_stream>>>(params); + } +} +}; // namespace rocm_backend diff --git a/src/grid/gpu/grid_gpu_integrate.h b/src/grid/gpu/grid_gpu_integrate.h deleted file mode 100644 index 0d304d082b..0000000000 --- a/src/grid/gpu/grid_gpu_integrate.h +++ /dev/null @@ -1,36 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 diff --git a/src/grid/hip/grid_hip_internal_header.h b/src/grid/gpu/grid_gpu_internal_header.h similarity index 98% rename from src/grid/hip/grid_hip_internal_header.h rename to src/grid/gpu/grid_gpu_internal_header.h index 978fa9065c..d1dd9b6c96 100644 --- a/src/grid/hip/grid_hip_internal_header.h +++ b/src/grid/gpu/grid_gpu_internal_header.h @@ -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 #include -#include #include #include #include @@ -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; } diff --git a/src/grid/hip/grid_hip_prepare_pab.h b/src/grid/gpu/grid_gpu_prepare_pab.h similarity index 99% rename from src/grid/hip/grid_hip_prepare_pab.h rename to src/grid/gpu/grid_gpu_prepare_pab.h index 3f1e26cec8..3752b2dced 100644 --- a/src/grid/hip/grid_hip_prepare_pab.h +++ b/src/grid/gpu/grid_gpu_prepare_pab.h @@ -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 #include @@ -404,3 +407,4 @@ inline ldiffs_value prepare_get_ldiffs(const enum grid_func func) { return ldiffs; } } // namespace rocm_backend +#endif diff --git a/src/grid/hip/grid_hip_process_vab.h b/src/grid/gpu/grid_gpu_process_vab.h similarity index 99% rename from src/grid/hip/grid_hip_process_vab.h rename to src/grid/gpu/grid_gpu_process_vab.h index 99916f9016..ce8124a63a 100644 --- a/src/grid/hip/grid_hip_process_vab.h +++ b/src/grid/gpu/grid_gpu_process_vab.h @@ -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 { diff --git a/src/grid/gpu/grid_gpu_task_list.cu b/src/grid/gpu/grid_gpu_task_list.cu deleted file mode 100644 index 4ad9df887b..0000000000 --- a/src/grid/gpu/grid_gpu_task_list.cu +++ /dev/null @@ -1,556 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 -#include -#include -#include -#include - -#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 diff --git a/src/grid/gpu/grid_gpu_task_list.h b/src/grid/gpu/grid_gpu_task_list.h index 85e89d82b7..13f6fe288f 100644 --- a/src/grid/gpu/grid_gpu_task_list.h +++ b/src/grid/gpu/grid_gpu_task_list.h @@ -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 #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 - -/******************************************************************************* - * \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 diff --git a/src/grid/grid_api.F b/src/grid/grid_api.F index 2548baf41a..f93b6072be 100644 --- a/src/grid/grid_api.F +++ b/src/grid/grid_api.F @@ -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 diff --git a/src/grid/grid_task_list.c b/src/grid/grid_task_list.c index 89e0af5df0..ddf136820f 100644 --- a/src/grid/grid_task_list.c +++ b/src/grid/grid_task_list.c @@ -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; diff --git a/src/grid/grid_task_list.h b/src/grid/grid_task_list.h index 4a393a3e6b..d714cf22e4 100644 --- a/src/grid/grid_task_list.h +++ b/src/grid/grid_task_list.h @@ -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; diff --git a/src/grid/hip/grid_hip_collocate.cu b/src/grid/hip/grid_hip_collocate.cu deleted file mode 100644 index 54e56d740b..0000000000 --- a/src/grid/hip/grid_hip_collocate.cu +++ /dev/null @@ -1,486 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 -#include -#include -#include -#include -#include -#include -#include - -#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 -__device__ __inline__ void block_to_cab(const kernel_params ¶ms, - const smem_task &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 -__global__ -__launch_bounds__(64) void calculate_coefficients(const kernel_params dev_) { - __shared__ smem_task 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(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 -__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 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 - <<>>(params); - } else { - calculate_coefficients - <<>>(params); - } - - if (grid_[level].is_distributed()) { - if (grid_[level].is_orthorhombic()) - collocate_kernel - <<>>(params); - else - collocate_kernel - <<>>(params); - } else { - if (grid_[level].is_orthorhombic()) - collocate_kernel - <<>>(params); - else - collocate_kernel - <<>>(params); - } -} -} // namespace rocm_backend -#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID) diff --git a/src/grid/hip/grid_hip_integrate.cu b/src/grid/hip/grid_hip_integrate.cu deleted file mode 100644 index 24ba0b96cd..0000000000 --- a/src/grid/hip/grid_hip_integrate.cu +++ /dev/null @@ -1,768 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 -#include -#include -#include -#include -#include -#include -#include - -#include "grid_hip_context.h" -#include "grid_hip_internal_header.h" -#include "grid_hip_process_vab.h" - -#ifdef __HIP_PLATFORM_NVIDIA__ -#include -#if CUDA_VERSION >= 11000 -#include -#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 -__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 -__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 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(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( - a, b, tid, task.zeta, task.zetb, task.n1, smem_cab); - shared_forces_b[tid] = get_force_b( - 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(a, b, tid / 3, tid % 3, task.zeta, - task.zetb, task.n1, smem_cab) + - get_virial_b(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( - // 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(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(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(sum, tid); - __syncthreads(); - if (tid == 0) - atomicAdd(forces_a + i, fa[i]); - __syncthreads(); - - sum[tid] = fb[i]; - __syncthreads(); - fb[i] = warp_reduce(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 -__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 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 - <<>>(params); - } else { - integrate_kernel - <<>>(params); - } - } else { - if (grid_[level].is_orthorhombic()) { - integrate_kernel - <<>>(params); - } else { - integrate_kernel - <<>>(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 - <<nblocks, threads_per_block, smem_params.smem_per_block(), - this->main_stream>>>(params); - return; - } - - if (!compute_tau && calculate_forces) { - compute_hab_v2 - <<nblocks, threads_per_block, smem_params.smem_per_block(), - this->main_stream>>>(params); - return; - } - - if (compute_tau && calculate_forces) { - compute_hab_v2 - <<nblocks, threads_per_block, smem_params.smem_per_block(), - this->main_stream>>>(params); - } - - if (compute_tau && !calculate_forces) { - compute_hab_v2 - <<nblocks, threads_per_block, smem_params.smem_per_block(), - this->main_stream>>>(params); - } -} -}; // namespace rocm_backend -#endif // defined(__OFFLOAD_HIP) && !defined(__NO_OFFLOAD_GRID) diff --git a/src/grid/hip/grid_hip_task_list.h b/src/grid/hip/grid_hip_task_list.h deleted file mode 100644 index d46615cb69..0000000000 --- a/src/grid/hip/grid_hip_task_list.h +++ /dev/null @@ -1,67 +0,0 @@ -/*----------------------------------------------------------------------------*/ -/* CP2K: A general program to perform molecular dynamics simulations */ -/* Copyright 2000-2026 CP2K developers group */ -/* */ -/* 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 - -#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 diff --git a/src/input_cp2k_global.F b/src/input_cp2k_global.F index 68dc899733..c77ae457c3 100644 --- a/src/input_cp2k_global.F +++ b/src/input_cp2k_global.F @@ -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) diff --git a/tools/precommit/check_file_properties.py b/tools/precommit/check_file_properties.py index 5a92c5d227..eb982de437 100755 --- a/tools/precommit/check_file_properties.py +++ b/tools/precommit/check_file_properties.py @@ -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",