Line data    Source code

       1             : /*----------------------------------------------------------------------------*/
       2             : /*  CP2K: A general program to perform molecular dynamics simulations         */
       3             : /*  Copyright 2000-2024 CP2K developers group <https://cp2k.org>              */
       4             : /*                                                                            */
       5             : /*  SPDX-License-Identifier: BSD-3-Clause                                     */
       6             : /*----------------------------------------------------------------------------*/
       7             : 
       8             : #ifndef GRID_DGEMM_PRIVATE_HEADER_H
       9             : #define GRID_DGEMM_PRIVATE_HEADER_H
      10             : 
      11             : #include "grid_dgemm_tensor_local.h"
      12             : #include <assert.h>
      13             : #include <stdbool.h>
      14             : /* everything here is specific to the dgemm and gpu backends*/
      15             : #include "../../offload/offload_buffer.h"
      16             : #include "../common/grid_basis_set.h"
      17             : #include "../common/grid_common.h"
      18             : #include "../common/grid_constants.h"
      19             : enum checksum_ { task_checksum = 0x2384989, ctx_checksum = 0x2356734 };
      20             : 
      21             : typedef struct {
      22             :   int xmin, xmax;
      23             : } Interval;
      24             : 
      25             : typedef struct {
      26             :   int level;
      27             :   int iatom;
      28             :   int jatom;
      29             :   int iset;
      30             :   int jset;
      31             :   int ipgf;
      32             :   int jpgf;
      33             :   int border_mask;
      34             :   int block_num;
      35             :   double radius;
      36             :   double zetp;
      37             :   double zeta[2];
      38             :   double ra[3];
      39             :   double rb[3];
      40             :   double rp[3];
      41             :   int lmax[2];
      42             :   int lmin[2];
      43             :   int l1_plus_l2_;
      44             :   int offset[2];
      45             :   bool update_block_;
      46             :   double rab[3];
      47             :   double prefactor;
      48             :   enum checksum_ checksum;
      49             : } _task;
      50             : 
      51             : typedef struct {
      52             :   int npts_global[3];
      53             :   int npts_local[3];
      54             :   int shift_local[3];
      55             :   int border_width[3];
      56             :   double dh[3][3];
      57             :   double dh_inv[3][3];
      58             : } _layout;
      59             : 
      60             : typedef struct grid_context_ {
      61             :   int ntasks;  // total number of tasks
      62             :   int nlevels; // number of different grid
      63             :   int natoms;
      64             :   int nkinds;
      65             :   int nblocks;
      66             :   int nblocks_total;
      67             :   int nkinds_total;
      68             :   int nlevels_total;
      69             :   int ntasks_total;
      70             :   int *block_offsets;
      71             :   double *atom_positions;
      72             :   int *atom_kinds;
      73             :   grid_basis_set **basis_sets;
      74             :   _task **tasks;
      75             :   _layout *layouts;
      76             :   int *tasks_per_level;
      77             :   int maxco;
      78             :   bool apply_cutoff;
      79             :   bool work_on_gpu;
      80             :   int number_of_devices;
      81             :   int *device_id;
      82             :   int queue_length;
      83             :   struct collocation_integration_ **handler;
      84             :   int number_of_handler;
      85             :   tensor *grid;
      86             :   void *scratch;
      87             :   bool orthorhombic;
      88             :   enum checksum_ checksum;
      89             : } grid_context;
      90             : 
      91     1287557 : static inline void update_loop_index(const int global_grid_size, int x1,
      92             :                                      int *const x) {
      93     1287557 :   *x += global_grid_size - x1 - 1;
      94     1287557 : }
      95             : 
      96             : static inline Interval create_interval(const int xmin, const int xmax) {
      97             :   assert(xmax >= xmin);
      98             : 
      99             :   Interval t = {.xmin = xmin, .xmax = xmax};
     100             :   return t;
     101             : }
     102             : 
     103     1287557 : static inline bool is_point_in_interval(const int value, Interval x) {
     104     1287557 :   return (value >= x.xmin) && (value <= x.xmax);
     105             : }
     106             : 
     107             : static inline bool intersection_interval_is_empty(const Interval x,
     108             :                                                   const Interval y) {
     109             :   /* return true if the intersection is empty */
     110             :   if ((x.xmin > y.xmax) || (x.xmax < y.xmin))
     111             :     return true;
     112             :   else
     113             :     return false;
     114             : }
     115             : 
     116             : static inline Interval intersection_interval(const Interval x,
     117             :                                              const Interval y) {
     118             :   Interval z;
     119             :   z.xmin = imax(x.xmin, y.xmin);
     120             :   z.xmax = imin(x.xmax, y.xmax);
     121             :   return z;
     122             : }
     123             : 
     124     1287557 : static inline int compute_next_boundaries(const int y1, const int y,
     125             :                                           const int grid_size,
     126             :                                           const int cube_size) {
     127     1287557 :   return y1 + imin(cube_size - y, grid_size - y1);
     128             : }
     129             : 
     130             : extern void grid_transform_coef_jik_to_yxz(const double dh[3][3],
     131             :                                            const tensor *coef_xyz);
     132             : extern void grid_transform_coef_xzy_to_ikj(const double dh[3][3],
     133             :                                            const tensor *coef_xyz);
     134             : extern void compute_block_boundaries(
     135             :     const int *blockDim, const int *lb_grid, const int *grid_size,
     136             :     const int *blocked_grid_size, const int *period, const int *cube_center,
     137             :     const int *cube_size, const int *lower_boundaries_cube,
     138             :     int *lower_block_corner, int *upper_block_corner, int *pol_offsets);
     139             : 
     140             : extern void grid_fill_pol_dgemm(const bool transpose, const double dr,
     141             :                                 const double roffset, const int pol_offset,
     142             :                                 const int xmin, const int xmax, const int lp,
     143             :                                 const int cmax, const double zetp,
     144             :                                 double *pol_);
     145             : 
     146             : /* this function is not exported outside. Should move */
     147             : extern void tensor_reduction_for_collocate_integrate(
     148             :     double *scratch, const double alpha, const bool *const orthogonal,
     149             :     const struct tensor_ *Exp, const struct tensor_ *co,
     150             :     const struct tensor_ *p_alpha_beta_reduced_, struct tensor_ *cube);
     151             : 
     152             : extern void set_grid_parameters(
     153             :     tensor *grid, /* tensor describing the grid */
     154             :     const bool orthorhombic,
     155             :     const int grid_full_size[3],  /* size of the full grid */
     156             :     const int grid_local_size[3], /* size of the local grid block */
     157             :     const int shift_local[3],     /* coordinates of the lower coordinates of the
     158             :                                      local grid window */
     159             :     const int border_width[3],    /* width of the borders */
     160             :     const double
     161             :         dh[3][3], /* displacement vectors of the grid (cartesian) -> (ijk) */
     162             :     const double dh_inv[3][3], /* (ijk) -> (x,y,z) */
     163             :     offload_buffer *grid_);
     164             : 
     165             : extern void collocate_one_grid_level_dgemm(grid_context *const ctx,
     166             :                                            const int *const, const int *const,
     167             :                                            const enum grid_func func,
     168             :                                            const int level,
     169             :                                            const offload_buffer *pab_blocks);
     170             : 
     171             : extern void integrate_one_grid_level_dgemm(
     172             :     grid_context *const ctx, const int level, const bool calculate_tau,
     173             :     const bool calculate_forces, const bool calculate_virial,
     174             :     const int *const shift_local, const int *const border_width,
     175             :     const offload_buffer *const pab_blocks, offload_buffer *const hab_blocks,
     176             :     tensor *forces_, tensor *virial_);
     177             : 
     178             : extern void compute_coefficients(grid_context *const ctx,
     179             :                                  struct collocation_integration_ *handler,
     180             :                                  const _task *previous_task, const _task *task,
     181             :                                  const offload_buffer *pab_blocks,
     182             :                                  tensor *const pab, tensor *const work,
     183             :                                  tensor *const pab_prep);
     184             : 
     185             : extern void extract_blocks(grid_context *const ctx, const _task *const task,
     186             :                            const offload_buffer *pab_blocks, tensor *const work,
     187             :                            tensor *const pab);
     188             : 
     189             : #endif