Line data    Source code

       1             : /*----------------------------------------------------------------------------*/
       2             : /*  CP2K: A general program to perform molecular dynamics simulations         */
       3             : /*  Copyright 2000-2025 CP2K developers group <https://cp2k.org>              */
       4             : /*                                                                            */
       5             : /*  SPDX-License-Identifier: BSD-3-Clause                                     */
       6             : /*----------------------------------------------------------------------------*/
       7             : 
       8             : #include <assert.h>
       9             : #include <math.h>
      10             : #include <omp.h>
      11             : #include <stdint.h>
      12             : #include <stdio.h>
      13             : #include <stdlib.h>
      14             : #include <string.h>
      15             : 
      16             : #include "../common/grid_common.h"
      17             : #include "grid_ref_collocate.h"
      18             : #include "grid_ref_integrate.h"
      19             : #include "grid_ref_task_list.h"
      20             : 
      21             : /*******************************************************************************
      22             :  * \brief Comperator passed to qsort to compare two tasks.
      23             :  * \author Ole Schuett
      24             :  ******************************************************************************/
      25    49740006 : static int compare_tasks(const void *a, const void *b) {
      26    49740006 :   const grid_ref_task *task_a = a, *task_b = b;
      27    49740006 :   if (task_a->level != task_b->level) {
      28     3708882 :     return task_a->level - task_b->level;
      29    46031124 :   } else if (task_a->block_num != task_b->block_num) {
      30    21392224 :     return task_a->block_num - task_b->block_num;
      31    24638900 :   } else if (task_a->iset != task_b->iset) {
      32     2932155 :     return task_a->iset - task_b->iset;
      33             :   } else {
      34    21706745 :     return task_a->jset - task_b->jset;
      35             :   }
      36             : }
      37             : /*******************************************************************************
      38             :  * \brief Allocates a task list for the reference backend.
      39             :  *        See grid_task_list.h for details.
      40             :  * \author Ole Schuett
      41             :  ******************************************************************************/
      42       13752 : void grid_ref_create_task_list(
      43             :     const bool orthorhombic, const int ntasks, const int nlevels,
      44             :     const int natoms, const int nkinds, const int nblocks,
      45             :     const int block_offsets[nblocks], const double atom_positions[natoms][3],
      46             :     const int atom_kinds[natoms], const grid_basis_set *basis_sets[nkinds],
      47             :     const int level_list[ntasks], const int iatom_list[ntasks],
      48             :     const int jatom_list[ntasks], const int iset_list[ntasks],
      49             :     const int jset_list[ntasks], const int ipgf_list[ntasks],
      50             :     const int jpgf_list[ntasks], const int border_mask_list[ntasks],
      51             :     const int block_num_list[ntasks], const double radius_list[ntasks],
      52             :     const double rab_list[ntasks][3], const int npts_global[nlevels][3],
      53             :     const int npts_local[nlevels][3], const int shift_local[nlevels][3],
      54             :     const int border_width[nlevels][3], const double dh[nlevels][3][3],
      55             :     const double dh_inv[nlevels][3][3], grid_ref_task_list **task_list_out) {
      56             : 
      57       13752 :   if (*task_list_out != NULL) {
      58             :     // This is actually an opportunity to reuse some buffers.
      59        5580 :     grid_ref_free_task_list(*task_list_out);
      60             :   }
      61             : 
      62       13752 :   grid_ref_task_list *task_list = malloc(sizeof(grid_ref_task_list));
      63       13752 :   assert(task_list != NULL);
      64             : 
      65       13752 :   task_list->orthorhombic = orthorhombic;
      66       13752 :   task_list->ntasks = ntasks;
      67       13752 :   task_list->nlevels = nlevels;
      68       13752 :   task_list->natoms = natoms;
      69       13752 :   task_list->nkinds = nkinds;
      70       13752 :   task_list->nblocks = nblocks;
      71             : 
      72       13752 :   size_t size = nblocks * sizeof(int);
      73       13752 :   task_list->block_offsets = malloc(size);
      74       13752 :   assert(task_list->block_offsets != NULL || size == 0);
      75       13752 :   if (size != 0) {
      76       13521 :     memcpy(task_list->block_offsets, block_offsets, size);
      77             :   }
      78             : 
      79       13752 :   size = 3 * natoms * sizeof(double);
      80       13752 :   task_list->atom_positions = malloc(size);
      81       13752 :   assert(task_list->atom_positions != NULL || size == 0);
      82       13752 :   if (size != 0) {
      83       13752 :     memcpy(task_list->atom_positions, atom_positions, size);
      84             :   }
      85             : 
      86       13752 :   size = natoms * sizeof(int);
      87       13752 :   task_list->atom_kinds = malloc(size);
      88       13752 :   assert(task_list->atom_kinds != NULL || size == 0);
      89       13752 :   if (size != 0) {
      90       13752 :     memcpy(task_list->atom_kinds, atom_kinds, size);
      91             :   }
      92             : 
      93       13752 :   size = nkinds * sizeof(grid_basis_set *);
      94       13752 :   task_list->basis_sets = malloc(size);
      95       13752 :   assert(task_list->basis_sets != NULL || size == 0);
      96       13752 :   if (size != 0) {
      97       13752 :     memcpy(task_list->basis_sets, basis_sets, size);
      98             :   }
      99             : 
     100       13752 :   size = ntasks * sizeof(grid_ref_task);
     101       13752 :   task_list->tasks = malloc(size);
     102       13752 :   assert(task_list->tasks != NULL || size == 0);
     103     7947994 :   for (int i = 0; i < ntasks; i++) {
     104     7934242 :     task_list->tasks[i].level = level_list[i];
     105     7934242 :     task_list->tasks[i].iatom = iatom_list[i];
     106     7934242 :     task_list->tasks[i].jatom = jatom_list[i];
     107     7934242 :     task_list->tasks[i].iset = iset_list[i];
     108     7934242 :     task_list->tasks[i].jset = jset_list[i];
     109     7934242 :     task_list->tasks[i].ipgf = ipgf_list[i];
     110     7934242 :     task_list->tasks[i].jpgf = jpgf_list[i];
     111     7934242 :     task_list->tasks[i].border_mask = border_mask_list[i];
     112     7934242 :     task_list->tasks[i].block_num = block_num_list[i];
     113     7934242 :     task_list->tasks[i].radius = radius_list[i];
     114     7934242 :     task_list->tasks[i].rab[0] = rab_list[i][0];
     115     7934242 :     task_list->tasks[i].rab[1] = rab_list[i][1];
     116     7934242 :     task_list->tasks[i].rab[2] = rab_list[i][2];
     117             :   }
     118             : 
     119             :   // Store grid layouts.
     120       13752 :   size = nlevels * sizeof(grid_ref_layout);
     121       13752 :   task_list->layouts = malloc(size);
     122       13752 :   assert(task_list->layouts != NULL || size == 0);
     123       68102 :   for (int level = 0; level < nlevels; level++) {
     124      217400 :     for (int i = 0; i < 3; i++) {
     125      163050 :       task_list->layouts[level].npts_global[i] = npts_global[level][i];
     126      163050 :       task_list->layouts[level].npts_local[i] = npts_local[level][i];
     127      163050 :       task_list->layouts[level].shift_local[i] = shift_local[level][i];
     128      163050 :       task_list->layouts[level].border_width[i] = border_width[level][i];
     129      652200 :       for (int j = 0; j < 3; j++) {
     130      489150 :         task_list->layouts[level].dh[i][j] = dh[level][i][j];
     131      489150 :         task_list->layouts[level].dh_inv[i][j] = dh_inv[level][i][j];
     132             :       }
     133             :     }
     134             :   }
     135             : 
     136             :   // Sort tasks by level, block_num, iset, and jset.
     137       13752 :   qsort(task_list->tasks, ntasks, sizeof(grid_ref_task), &compare_tasks);
     138             : 
     139             :   // Find first and last task for each level and block.
     140       13752 :   size = nlevels * nblocks * sizeof(int);
     141       13752 :   task_list->first_level_block_task = malloc(size);
     142       13752 :   assert(task_list->first_level_block_task != NULL || size == 0);
     143       13752 :   task_list->last_level_block_task = malloc(size);
     144       13752 :   assert(task_list->last_level_block_task != NULL || size == 0);
     145     1105725 :   for (int i = 0; i < nlevels * nblocks; i++) {
     146     1091973 :     task_list->first_level_block_task[i] = 0;
     147     1091973 :     task_list->last_level_block_task[i] = -1; // last < first means no tasks
     148             :   }
     149     7947994 :   for (int itask = 0; itask < ntasks; itask++) {
     150     7934242 :     const int level = task_list->tasks[itask].level - 1;
     151     7934242 :     const int block_num = task_list->tasks[itask].block_num - 1;
     152     7934242 :     if (itask == 0 || task_list->tasks[itask - 1].level - 1 != level ||
     153     7897071 :         task_list->tasks[itask - 1].block_num - 1 != block_num) {
     154      551967 :       task_list->first_level_block_task[level * nblocks + block_num] = itask;
     155             :     }
     156     7934242 :     task_list->last_level_block_task[level * nblocks + block_num] = itask;
     157             :   }
     158             : 
     159             :   // Find largest Cartesian subblock size.
     160       13752 :   task_list->maxco = 0;
     161       38497 :   for (int i = 0; i < nkinds; i++) {
     162       24745 :     task_list->maxco = imax(task_list->maxco, task_list->basis_sets[i]->maxco);
     163             :   }
     164             : 
     165             :   // Initialize thread-local storage.
     166       13752 :   size = omp_get_max_threads() * sizeof(double *);
     167       13752 :   task_list->threadlocals = malloc(size);
     168       13752 :   assert(task_list->threadlocals != NULL);
     169       13752 :   memset(task_list->threadlocals, 0, size);
     170       13752 :   size = omp_get_max_threads() * sizeof(size_t);
     171       13752 :   task_list->threadlocal_sizes = malloc(size);
     172       13752 :   assert(task_list->threadlocal_sizes != NULL);
     173       13752 :   memset(task_list->threadlocal_sizes, 0, size);
     174             : 
     175       13752 :   *task_list_out = task_list;
     176       13752 : }
     177             : 
     178             : /*******************************************************************************
     179             :  * \brief Deallocates given task list, basis_sets have to be freed separately.
     180             :  * \author Ole Schuett
     181             :  ******************************************************************************/
     182       13752 : void grid_ref_free_task_list(grid_ref_task_list *task_list) {
     183       13752 :   free(task_list->block_offsets);
     184       13752 :   free(task_list->atom_positions);
     185       13752 :   free(task_list->atom_kinds);
     186       13752 :   free(task_list->basis_sets);
     187       13752 :   free(task_list->tasks);
     188       13752 :   free(task_list->layouts);
     189       13752 :   free(task_list->first_level_block_task);
     190       13752 :   free(task_list->last_level_block_task);
     191       27504 :   for (int i = 0; i < omp_get_max_threads(); i++) {
     192       13752 :     if (task_list->threadlocals[i] != NULL) {
     193           8 :       free(task_list->threadlocals[i]);
     194             :     }
     195             :   }
     196       13752 :   free(task_list->threadlocals);
     197       13752 :   free(task_list->threadlocal_sizes);
     198       13752 :   free(task_list);
     199       13752 : }
     200             : 
     201             : /*******************************************************************************
     202             :  * \brief Prototype for BLAS dgemm.
     203             :  * \author Ole Schuett
     204             :  ******************************************************************************/
     205             : void dgemm_(const char *transa, const char *transb, const int *m, const int *n,
     206             :             const int *k, const double *alpha, const double *a, const int *lda,
     207             :             const double *b, const int *ldb, const double *beta, double *c,
     208             :             const int *ldc);
     209             : 
     210             : /*******************************************************************************
     211             :  * \brief Convenient wrapper to hide Fortran nature of dgemm_, swapping a and b.
     212             :  * \author Ole Schuett
     213             :  ******************************************************************************/
     214        1928 : static void dgemm(const char transa, const char transb, const int m,
     215             :                   const int n, const int k, const double alpha, const double *a,
     216             :                   const int lda, const double *b, const int ldb,
     217             :                   const double beta, double *c, const int ldc) {
     218        1928 :   dgemm_(&transb, &transa, &n, &m, &k, &alpha, b, &ldb, a, &lda, &beta, c,
     219             :          &ldc);
     220        1928 : }
     221             : 
     222             : /*******************************************************************************
     223             :  * \brief Transforms pab from contracted spherical to prim. cartesian basis.
     224             :  * \author Ole Schuett
     225             :  ******************************************************************************/
     226         580 : static void load_pab(const grid_basis_set *ibasis, const grid_basis_set *jbasis,
     227             :                      const int iset, const int jset, const bool transpose,
     228         580 :                      const double *block, double *pab) {
     229             : 
     230             :   // Define some more convenient aliases.
     231         580 :   const int ncoseta = ncoset(ibasis->lmax[iset]);
     232         580 :   const int ncosetb = ncoset(jbasis->lmax[jset]);
     233         580 :   const int ncoa = ibasis->npgf[iset] * ncoseta; // size of carthesian set
     234         580 :   const int ncob = jbasis->npgf[jset] * ncosetb;
     235             : 
     236         580 :   const int nsgf_seta = ibasis->nsgf_set[iset]; // size of spherical set
     237         580 :   const int nsgf_setb = jbasis->nsgf_set[jset];
     238         580 :   const int nsgfa = ibasis->nsgf; // size of entire spherical basis
     239         580 :   const int nsgfb = jbasis->nsgf;
     240         580 :   const int sgfa = ibasis->first_sgf[iset] - 1; // start of spherical set
     241         580 :   const int sgfb = jbasis->first_sgf[jset] - 1;
     242         580 :   const int maxcoa = ibasis->maxco;
     243         580 :   const int maxcob = jbasis->maxco;
     244             : 
     245         580 :   double work[nsgf_setb * ncoa];
     246         580 :   if (transpose) {
     247             :     // work[nsgf_setb][ncoa] = MATMUL(subblock, ibasis->sphi)
     248         412 :     dgemm('N', 'N', nsgf_setb, ncoa, nsgf_seta, 1.0,
     249         412 :           &block[sgfb * nsgfa + sgfa], nsgfa, &ibasis->sphi[sgfa * maxcoa],
     250             :           maxcoa, 0.0, work, ncoa);
     251             :   } else {
     252             :     // work[nsgf_setb][ncoa] = MATMUL(TRANSPOSE(subblock), ibasis->sphi)
     253         168 :     dgemm('T', 'N', nsgf_setb, ncoa, nsgf_seta, 1.0,
     254         168 :           &block[sgfa * nsgfb + sgfb], nsgfb, &ibasis->sphi[sgfa * maxcoa],
     255             :           maxcoa, 0.0, work, ncoa);
     256             :   }
     257             :   // pab[ncob][ncoa] = MATMUL(TRANSPOSE(jbasis->sphi), work)
     258         580 :   dgemm('T', 'N', ncob, ncoa, nsgf_setb, 1.0, &jbasis->sphi[sgfb * maxcob],
     259             :         maxcob, work, ncoa, 0.0, pab, ncoa);
     260         580 : }
     261             : 
     262             : /*******************************************************************************
     263             :  * \brief Collocate a range of tasks which are destined for the same grid level.
     264             :  * \author Ole Schuett
     265             :  ******************************************************************************/
     266         240 : static void collocate_one_grid_level(
     267             :     const grid_ref_task_list *task_list, const int *first_block_task,
     268             :     const int *last_block_task, const enum grid_func func,
     269             :     const int npts_global[3], const int npts_local[3], const int shift_local[3],
     270             :     const int border_width[3], const double dh[3][3], const double dh_inv[3][3],
     271             :     const double *pab_blocks, offload_buffer *grid) {
     272             : 
     273             : // Using default(shared) because with GCC 9 the behavior around const changed:
     274             : // https://www.gnu.org/software/gcc/gcc-9/porting_to.html
     275         240 : #pragma omp parallel default(shared)
     276             :   {
     277             :     const int ithread = omp_get_thread_num();
     278             :     const int nthreads = omp_get_num_threads();
     279             : 
     280             :     // Initialize variables to detect when a new subblock has to be fetched.
     281             :     int old_offset = -1, old_iset = -1, old_jset = -1;
     282             : 
     283             :     // Matrix pab is re-used across tasks.
     284             :     double pab[imax(task_list->maxco * task_list->maxco, 1)];
     285             : 
     286             :     // Ensure that grid can fit into thread-local storage, reallocate if needed.
     287             :     const int npts_local_total = npts_local[0] * npts_local[1] * npts_local[2];
     288             :     const size_t grid_size = npts_local_total * sizeof(double);
     289             :     if (task_list->threadlocal_sizes[ithread] < grid_size) {
     290             :       if (task_list->threadlocals[ithread] != NULL) {
     291             :         free(task_list->threadlocals[ithread]);
     292             :       }
     293             :       task_list->threadlocals[ithread] = malloc(grid_size);
     294             :       assert(task_list->threadlocals[ithread] != NULL);
     295             :       task_list->threadlocal_sizes[ithread] = grid_size;
     296             :     }
     297             : 
     298             :     // Zero thread-local copy of the grid.
     299             :     double *const my_grid = task_list->threadlocals[ithread];
     300             :     memset(my_grid, 0, grid_size);
     301             : 
     302             :     // Parallelize over blocks to avoid unnecessary calls to load_pab.
     303             :     const int chunk_size = imax(1, task_list->nblocks / (nthreads * 50));
     304             : #pragma omp for schedule(dynamic, chunk_size)
     305             :     for (int block_num = 0; block_num < task_list->nblocks; block_num++) {
     306             :       const int first_task = first_block_task[block_num];
     307             :       const int last_task = last_block_task[block_num];
     308             : 
     309             :       for (int itask = first_task; itask <= last_task; itask++) {
     310             :         // Define some convenient aliases.
     311             :         const grid_ref_task *task = &task_list->tasks[itask];
     312             :         const int iatom = task->iatom - 1;
     313             :         const int jatom = task->jatom - 1;
     314             :         const int iset = task->iset - 1;
     315             :         const int jset = task->jset - 1;
     316             :         const int ipgf = task->ipgf - 1;
     317             :         const int jpgf = task->jpgf - 1;
     318             :         const int ikind = task_list->atom_kinds[iatom] - 1;
     319             :         const int jkind = task_list->atom_kinds[jatom] - 1;
     320             :         const grid_basis_set *ibasis = task_list->basis_sets[ikind];
     321             :         const grid_basis_set *jbasis = task_list->basis_sets[jkind];
     322             :         const double zeta = ibasis->zet[iset * ibasis->maxpgf + ipgf];
     323             :         const double zetb = jbasis->zet[jset * jbasis->maxpgf + jpgf];
     324             :         const int ncoseta = ncoset(ibasis->lmax[iset]);
     325             :         const int ncosetb = ncoset(jbasis->lmax[jset]);
     326             :         const int ncoa = ibasis->npgf[iset] * ncoseta; // size of carthesian set
     327             :         const int ncob = jbasis->npgf[jset] * ncosetb;
     328             :         const int block_num = task->block_num - 1;
     329             :         const int block_offset = task_list->block_offsets[block_num];
     330             :         const double *block = &pab_blocks[block_offset];
     331             :         const bool transpose = (iatom <= jatom);
     332             : 
     333             :         // Load subblock from buffer and decontract into Cartesian sublock pab.
     334             :         // The previous pab can be reused when only ipgf or jpgf has changed.
     335             :         if (block_offset != old_offset || iset != old_iset ||
     336             :             jset != old_jset) {
     337             :           old_offset = block_offset;
     338             :           old_iset = iset;
     339             :           old_jset = jset;
     340             :           load_pab(ibasis, jbasis, iset, jset, transpose, block, pab);
     341             :         }
     342             : 
     343             :         grid_ref_collocate_pgf_product(
     344             :             /*orthorhombic=*/task_list->orthorhombic,
     345             :             /*border_mask=*/task->border_mask,
     346             :             /*func=*/func,
     347             :             /*la_max=*/ibasis->lmax[iset],
     348             :             /*la_min=*/ibasis->lmin[iset],
     349             :             /*lb_max=*/jbasis->lmax[jset],
     350             :             /*lb_min=*/jbasis->lmin[jset],
     351             :             /*zeta=*/zeta,
     352             :             /*zetb=*/zetb,
     353             :             /*rscale=*/(iatom == jatom) ? 1 : 2,
     354             :             /*dh=*/dh,
     355             :             /*dh_inv=*/dh_inv,
     356             :             /*ra=*/&task_list->atom_positions[3 * iatom],
     357             :             /*rab=*/task->rab,
     358             :             /*npts_global=*/npts_global,
     359             :             /*npts_local=*/npts_local,
     360             :             /*shift_local=*/shift_local,
     361             :             /*border_width=*/border_width,
     362             :             /*radius=*/task->radius,
     363             :             /*o1=*/ipgf * ncoseta,
     364             :             /*o2=*/jpgf * ncosetb,
     365             :             /*n1=*/ncoa,
     366             :             /*n2=*/ncob,
     367             :             /*pab=*/(const double(*)[ncoa])pab,
     368             :             /*grid=*/my_grid);
     369             : 
     370             :       } // end of task loop
     371             :     } // end of block loop
     372             : 
     373             : // While there should be an implicit barrier at the end of the block loop, this
     374             : // explicit barrier eliminates occasional seg faults with icc compiled binaries.
     375             : #pragma omp barrier
     376             : 
     377             :     // Merge thread-local grids via an efficient tree reduction.
     378             :     const int nreduction_cycles = ceil(log(nthreads) / log(2)); // tree depth
     379             :     for (int icycle = 1; icycle <= nreduction_cycles; icycle++) {
     380             :       // Threads are divided into groups, whose size doubles with each cycle.
     381             :       // After a cycle the reduced data is stored at first thread of each group.
     382             :       const int group_size = 1 << icycle; // 2**icycle
     383             :       const int igroup = ithread / group_size;
     384             :       const int dest_thread = igroup * group_size;
     385             :       const int src_thread = dest_thread + group_size / 2;
     386             :       // The last group might actually be a bit smaller.
     387             :       const int actual_group_size = imin(group_size, nthreads - dest_thread);
     388             :       // Parallelize summation by dividing grid points across group members.
     389             :       const int rank = modulo(ithread, group_size); // position within the group
     390             :       const int lb = (npts_local_total * rank) / actual_group_size;
     391             :       const int ub = (npts_local_total * (rank + 1)) / actual_group_size;
     392             :       if (src_thread < nthreads) {
     393             :         for (int i = lb; i < ub; i++) {
     394             :           task_list->threadlocals[dest_thread][i] +=
     395             :               task_list->threadlocals[src_thread][i];
     396             :         }
     397             :       }
     398             : #pragma omp barrier
     399             :     }
     400             : 
     401             :     // Copy final result from first thread into shared grid.
     402             :     const int lb = (npts_local_total * ithread) / nthreads;
     403             :     const int ub = (npts_local_total * (ithread + 1)) / nthreads;
     404             :     for (int i = lb; i < ub; i++) {
     405             :       grid->host_buffer[i] = task_list->threadlocals[0][i];
     406             :     }
     407             : 
     408             :   } // end of omp parallel region
     409         240 : }
     410             : 
     411             : /*******************************************************************************
     412             :  * \brief Collocate all tasks of in given list onto given grids.
     413             :  *        See grid_task_list.h for details.
     414             :  * \author Ole Schuett
     415             :  ******************************************************************************/
     416          60 : void grid_ref_collocate_task_list(const grid_ref_task_list *task_list,
     417             :                                   const enum grid_func func, const int nlevels,
     418             :                                   const offload_buffer *pab_blocks,
     419             :                                   offload_buffer *grids[nlevels]) {
     420             : 
     421          60 :   assert(task_list->nlevels == nlevels);
     422             : 
     423         300 :   for (int level = 0; level < task_list->nlevels; level++) {
     424         240 :     const int idx = level * task_list->nblocks;
     425         240 :     const int *first_block_task = &task_list->first_level_block_task[idx];
     426         240 :     const int *last_block_task = &task_list->last_level_block_task[idx];
     427         240 :     const grid_ref_layout *layout = &task_list->layouts[level];
     428         240 :     collocate_one_grid_level(
     429         240 :         task_list, first_block_task, last_block_task, func, layout->npts_global,
     430         240 :         layout->npts_local, layout->shift_local, layout->border_width,
     431         240 :         layout->dh, layout->dh_inv, pab_blocks->host_buffer, grids[level]);
     432             :   }
     433          60 : }
     434             : 
     435             : /*******************************************************************************
     436             :  * \brief Transforms hab from prim. cartesian to contracted spherical basis.
     437             :  * \author Ole Schuett
     438             :  ******************************************************************************/
     439         384 : static inline void store_hab(const grid_basis_set *ibasis,
     440             :                              const grid_basis_set *jbasis, const int iset,
     441             :                              const int jset, const bool transpose,
     442         384 :                              const double *hab, double *block) {
     443             : 
     444             :   // Define some more convenient aliases.
     445         384 :   const int ncoseta = ncoset(ibasis->lmax[iset]);
     446         384 :   const int ncosetb = ncoset(jbasis->lmax[jset]);
     447         384 :   const int ncoa = ibasis->npgf[iset] * ncoseta; // size of carthesian set
     448         384 :   const int ncob = jbasis->npgf[jset] * ncosetb;
     449             : 
     450         384 :   const int nsgf_seta = ibasis->nsgf_set[iset]; // size of spherical set
     451         384 :   const int nsgf_setb = jbasis->nsgf_set[jset];
     452         384 :   const int nsgfa = ibasis->nsgf; // size of entire spherical basis
     453         384 :   const int nsgfb = jbasis->nsgf;
     454         384 :   const int sgfa = ibasis->first_sgf[iset] - 1; // start of spherical set
     455         384 :   const int sgfb = jbasis->first_sgf[jset] - 1;
     456         384 :   const int maxcoa = ibasis->maxco;
     457         384 :   const int maxcob = jbasis->maxco;
     458             : 
     459         384 :   double work[nsgf_setb * ncoa];
     460             : 
     461             :   // work[nsgf_setb][ncoa] = MATMUL(jbasis->sphi, hab)
     462         384 :   dgemm('N', 'N', nsgf_setb, ncoa, ncob, 1.0, &jbasis->sphi[sgfb * maxcob],
     463             :         maxcob, hab, ncoa, 0.0, work, ncoa);
     464             : 
     465         384 :   if (transpose) {
     466             :     // subblock[nsgf_setb][nsgf_seta] += MATMUL(work, TRANSPOSE(ibasis->sphi))
     467         288 :     dgemm('N', 'T', nsgf_setb, nsgf_seta, ncoa, 1.0, work, ncoa,
     468         288 :           &ibasis->sphi[sgfa * maxcoa], maxcoa, 1.0,
     469         288 :           &block[sgfb * nsgfa + sgfa], nsgfa);
     470             :   } else {
     471             :     // subblock[nsgf_seta][nsgf_setb] += MATMUL(ibasis->sphi, TRANSPOSE(work))
     472          96 :     dgemm('N', 'T', nsgf_seta, nsgf_setb, ncoa, 1.0,
     473          96 :           &ibasis->sphi[sgfa * maxcoa], maxcoa, work, ncoa, 1.0,
     474          96 :           &block[sgfa * nsgfb + sgfb], nsgfb);
     475             :   }
     476         384 : }
     477             : 
     478             : /*******************************************************************************
     479             :  * \brief Integrate a range of tasks that belong to the same grid level.
     480             :  * \author Ole Schuett
     481             :  ******************************************************************************/
     482         208 : static void integrate_one_grid_level(
     483             :     const grid_ref_task_list *task_list, const int *first_block_task,
     484             :     const int *last_block_task, const bool compute_tau, const int natoms,
     485             :     const int npts_global[3], const int npts_local[3], const int shift_local[3],
     486             :     const int border_width[3], const double dh[3][3], const double dh_inv[3][3],
     487             :     const offload_buffer *pab_blocks, const offload_buffer *grid,
     488             :     offload_buffer *hab_blocks, double forces[natoms][3], double virial[3][3]) {
     489             : 
     490             : // Using default(shared) because with GCC 9 the behavior around const changed:
     491             : // https://www.gnu.org/software/gcc/gcc-9/porting_to.html
     492         208 : #pragma omp parallel default(shared)
     493             :   {
     494             :     // Initialize variables to detect when a new subblock has to be fetched.
     495             :     int old_offset = -1, old_iset = -1, old_jset = -1;
     496             :     grid_basis_set *old_ibasis = NULL, *old_jbasis = NULL;
     497             :     bool old_transpose = false;
     498             : 
     499             :     // Matrix pab and hab are re-used across tasks.
     500             :     double pab[imax(task_list->maxco * task_list->maxco, 1)];
     501             :     double hab[imax(task_list->maxco * task_list->maxco, 1)];
     502             : 
     503             :     // Parallelize over blocks to avoid concurred access to hab_blocks.
     504             :     const int nthreads = omp_get_num_threads();
     505             :     const int chunk_size = imax(1, task_list->nblocks / (nthreads * 50));
     506             : #pragma omp for schedule(dynamic, chunk_size)
     507             :     for (int block_num = 0; block_num < task_list->nblocks; block_num++) {
     508             :       const int first_task = first_block_task[block_num];
     509             :       const int last_task = last_block_task[block_num];
     510             : 
     511             :       // Accumulate forces per block as it corresponds to a pair of atoms.
     512             :       const int iatom = task_list->tasks[first_task].iatom - 1;
     513             :       const int jatom = task_list->tasks[first_task].jatom - 1;
     514             :       double my_forces[2][3] = {0};
     515             :       double my_virials[2][3][3] = {0};
     516             : 
     517             :       for (int itask = first_task; itask <= last_task; itask++) {
     518             :         // Define some convenient aliases.
     519             :         const grid_ref_task *task = &task_list->tasks[itask];
     520             :         assert(task->block_num - 1 == block_num);
     521             :         assert(task->iatom - 1 == iatom && task->jatom - 1 == jatom);
     522             :         const int ikind = task_list->atom_kinds[iatom] - 1;
     523             :         const int jkind = task_list->atom_kinds[jatom] - 1;
     524             :         grid_basis_set *ibasis = task_list->basis_sets[ikind];
     525             :         grid_basis_set *jbasis = task_list->basis_sets[jkind];
     526             :         const int iset = task->iset - 1;
     527             :         const int jset = task->jset - 1;
     528             :         const int ipgf = task->ipgf - 1;
     529             :         const int jpgf = task->jpgf - 1;
     530             :         const double zeta = ibasis->zet[iset * ibasis->maxpgf + ipgf];
     531             :         const double zetb = jbasis->zet[jset * jbasis->maxpgf + jpgf];
     532             :         const int ncoseta = ncoset(ibasis->lmax[iset]);
     533             :         const int ncosetb = ncoset(jbasis->lmax[jset]);
     534             :         const int ncoa = ibasis->npgf[iset] * ncoseta; // size of carthesian set
     535             :         const int ncob = jbasis->npgf[jset] * ncosetb;
     536             :         const int block_offset = task_list->block_offsets[block_num];
     537             :         const bool transpose = (iatom <= jatom);
     538             :         const bool pab_required = (forces != NULL || virial != NULL);
     539             : 
     540             :         // Load pab and store hab subblocks when needed.
     541             :         // Previous hab and pab can be reused when only ipgf or jpgf changed.
     542             :         if (block_offset != old_offset || iset != old_iset ||
     543             :             jset != old_jset) {
     544             :           if (pab_required) {
     545             :             load_pab(ibasis, jbasis, iset, jset, transpose,
     546             :                      &pab_blocks->host_buffer[block_offset], pab);
     547             :           }
     548             :           if (old_offset >= 0) { // skip first iteration
     549             :             store_hab(old_ibasis, old_jbasis, old_iset, old_jset, old_transpose,
     550             :                       hab, &hab_blocks->host_buffer[old_offset]);
     551             :           }
     552             :           memset(hab, 0, ncoa * ncob * sizeof(double));
     553             :           old_offset = block_offset;
     554             :           old_iset = iset;
     555             :           old_jset = jset;
     556             :           old_ibasis = ibasis;
     557             :           old_jbasis = jbasis;
     558             :           old_transpose = transpose;
     559             :         }
     560             : 
     561             :         grid_ref_integrate_pgf_product(
     562             :             /*orthorhombic=*/task_list->orthorhombic,
     563             :             /*compute_tau=*/compute_tau,
     564             :             /*border_mask=*/task->border_mask,
     565             :             /*la_max=*/ibasis->lmax[iset],
     566             :             /*la_min=*/ibasis->lmin[iset],
     567             :             /*lb_max=*/jbasis->lmax[jset],
     568             :             /*lb_min=*/jbasis->lmin[jset],
     569             :             /*zeta=*/zeta,
     570             :             /*zetb=*/zetb,
     571             :             /*dh=*/dh,
     572             :             /*dh_inv=*/dh_inv,
     573             :             /*ra=*/&task_list->atom_positions[3 * iatom],
     574             :             /*rab=*/task->rab,
     575             :             /*npts_global=*/npts_global,
     576             :             /*npts_local=*/npts_local,
     577             :             /*shift_local=*/shift_local,
     578             :             /*border_width=*/border_width,
     579             :             /*radius=*/task->radius,
     580             :             /*o1=*/ipgf * ncoseta,
     581             :             /*o2=*/jpgf * ncosetb,
     582             :             /*n1=*/ncoa,
     583             :             /*n2=*/ncob,
     584             :             /*grid=*/grid->host_buffer,
     585             :             /*hab=*/(double(*)[ncoa])hab,
     586             :             /*pab=*/(pab_required) ? (const double(*)[ncoa])pab : NULL,
     587             :             /*forces=*/(forces != NULL) ? my_forces : NULL,
     588             :             /*virials=*/(virial != NULL) ? my_virials : NULL,
     589             :             /*hdab=*/NULL,
     590             :             /*hadb=*/NULL,
     591             :             /*a_hdab=*/NULL);
     592             : 
     593             :       } // end of task loop
     594             : 
     595             :       // Merge thread-local forces and virial into shared ones.
     596             :       // It does not seem worth the trouble to accumulate them thread-locally.
     597             :       const double scalef = (iatom == jatom) ? 1.0 : 2.0;
     598             :       if (forces != NULL) {
     599             : #pragma omp critical(forces)
     600             :         for (int i = 0; i < 3; i++) {
     601             :           forces[iatom][i] += scalef * my_forces[0][i];
     602             :           forces[jatom][i] += scalef * my_forces[1][i];
     603             :         }
     604             :       }
     605             :       if (virial != NULL) {
     606             : #pragma omp critical(virial)
     607             :         for (int i = 0; i < 3; i++) {
     608             :           for (int j = 0; j < 3; j++) {
     609             :             virial[i][j] += scalef * my_virials[0][i][j];
     610             :             virial[i][j] += scalef * my_virials[1][i][j];
     611             :           }
     612             :         }
     613             :       }
     614             : 
     615             :     } // end of block loop
     616             : 
     617             :     // store final hab
     618             :     if (old_offset >= 0) {
     619             :       store_hab(old_ibasis, old_jbasis, old_iset, old_jset, old_transpose, hab,
     620             :                 &hab_blocks->host_buffer[old_offset]);
     621             :     }
     622             : 
     623             :   } // end of omp parallel region
     624         208 : }
     625             : 
     626             : /*******************************************************************************
     627             :  * \brief Integrate all tasks of in given list from given grids.
     628             :  *        See grid_task_list.h for details.
     629             :  * \author Ole Schuett
     630             :  ******************************************************************************/
     631          52 : void grid_ref_integrate_task_list(
     632             :     const grid_ref_task_list *task_list, const bool compute_tau,
     633             :     const int natoms, const int nlevels, const offload_buffer *pab_blocks,
     634             :     const offload_buffer *grids[nlevels], offload_buffer *hab_blocks,
     635             :     double forces[natoms][3], double virial[3][3]) {
     636             : 
     637          52 :   assert(task_list->nlevels == nlevels);
     638          52 :   assert(task_list->natoms == natoms);
     639             : 
     640             :   // Zero result arrays.
     641          52 :   memset(hab_blocks->host_buffer, 0, hab_blocks->size);
     642          52 :   if (forces != NULL) {
     643           8 :     memset(forces, 0, natoms * 3 * sizeof(double));
     644             :   }
     645          52 :   if (virial != NULL) {
     646           0 :     memset(virial, 0, 9 * sizeof(double));
     647             :   }
     648             : 
     649         260 :   for (int level = 0; level < task_list->nlevels; level++) {
     650         208 :     const int idx = level * task_list->nblocks;
     651         208 :     const int *first_block_task = &task_list->first_level_block_task[idx];
     652         208 :     const int *last_block_task = &task_list->last_level_block_task[idx];
     653         208 :     const grid_ref_layout *layout = &task_list->layouts[level];
     654         208 :     integrate_one_grid_level(
     655             :         task_list, first_block_task, last_block_task, compute_tau, natoms,
     656         208 :         layout->npts_global, layout->npts_local, layout->shift_local,
     657         208 :         layout->border_width, layout->dh, layout->dh_inv, pab_blocks,
     658         208 :         grids[level], hab_blocks, forces, virial);
     659             :   }
     660          52 : }
     661             : 
     662             : // EOF