Merge pull request #1894 from stanmoore1/kk_snap_chunk

Make Kokkos pair SNAP chunksize variable user-settable

doc/src/compute_sna_atom.rst

@@ -46,8 +46,6 @@ Syntax
           *0* = do not generate quadratic terms
           *1* = generate quadratic terms
 
-
-
 Examples
 """"""""
 

doc/src/pair_snap.rst

@@ -139,7 +139,7 @@ The SNAP parameter file can contain blank and comment lines (start
 with #) anywhere. Each non-blank non-comment line must contain one
 keyword/value pair. The required keywords are *rcutfac* and
 *twojmax*\ . Optional keywords are *rfac0*\ , *rmin0*\ ,
-*switchflag*\ , and *bzeroflag*\ .
+*switchflag*\ , *bzeroflag*\, and *chunksize*\.
 
 The default values for these keywords are
 
@@ -148,8 +148,20 @@ The default values for these keywords are
 * *switchflag* = 0
 * *bzeroflag* = 1
 * *quadraticflag* = 1
+* *chunksize* = 2000
+
+The keyword *chunksize* is only applicable when using the
+pair style *snap* with the KOKKOS package and is ignored otherwise. 
+This keyword controls
+the number of atoms in each pass used to compute the bispectrum
+components and is used to avoid running out of memory. For example
+if there are 4000 atoms in the simulation and the *chunksize*
+is set to 2000, the bispectrum calculation will be broken up
+into two passes.
+
+Detailed definitions for all the other keywords 
+are given on the :doc:`compute sna/atom <compute_sna_atom>` doc page. 
 
-Detailed definitions for all the keywords are given on the :doc:`compute sna/atom <compute_sna_atom>` doc page.
 If *quadraticflag* is set to 1, then the SNAP energy expression includes the quadratic term,
 0.5\*B\^t.alpha.B, where alpha is a symmetric *K* by *K* matrix.
 The SNAP element file should contain *K*\ (\ *K*\ +1)/2 additional coefficients

src/KOKKOS/pair_snap_kokkos.h

@@ -122,7 +122,7 @@ protected:
   t_dbvec dbvec;
   SNAKokkos<DeviceType> snaKK;
 
-  int inum,max_neighs,chunk_offset;
+  int inum,max_neighs,chunk_size,chunk_offset;
 
   int eflag,vflag;
 

src/KOKKOS/pair_snap_kokkos_impl.h

@@ -198,7 +198,7 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
     d_ninside = Kokkos::View<int*, DeviceType>("PairSNAPKokkos:ninside",inum);
   }
 
-  int chunk_size = MIN(2000,inum);
+  chunk_size = MIN(chunksize,inum); // "chunksize" variable is set by user
   chunk_offset = 0;
 
   snaKK.grow_rij(chunk_size,max_neighs);
@@ -221,7 +221,7 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
     Kokkos::parallel_for("PreUi",policy_preui,*this);
 
     //ComputeUi
-    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi> policy_ui(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
+    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi> policy_ui(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
     Kokkos::parallel_for("ComputeUi",policy_ui,*this);
 
     //Ulisttot transpose
@@ -253,11 +253,11 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
     Kokkos::parallel_for("ComputeYi",policy_yi,*this);
 
     //ComputeDuidrj
-    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrj> policy_duidrj(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
+    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrj> policy_duidrj(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
     Kokkos::parallel_for("ComputeDuidrj",policy_duidrj,*this);
 
     //ComputeDeidrj
-    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrj> policy_deidrj(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
+    typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrj> policy_deidrj(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
     Kokkos::parallel_for("ComputeDeidrj",policy_deidrj,*this);
 
     //ComputeForce
@@ -514,11 +514,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeUi,const typename
   SNAKokkos<DeviceType> my_sna = snaKK;
 
   // Extract the atom number
-  int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
-  if (ii >= inum) return;
+  int ii = team.team_rank() + team.team_size() * (team.league_rank() %
+           ((chunk_size+team.team_size()-1)/team.team_size()));
+  if (ii >= chunk_size) return;
 
   // Extract the neighbor number
-  const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
+  const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
   const int ninside = d_ninside(ii);
   if (jj >= ninside) return;
 
@@ -560,11 +561,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeDuidrj,const type
   SNAKokkos<DeviceType> my_sna = snaKK;
 
   // Extract the atom number
-  int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
-  if (ii >= inum) return;
+  int ii = team.team_rank() + team.team_size() * (team.league_rank() %
+           ((chunk_size+team.team_size()-1)/team.team_size()));
+  if (ii >= chunk_size) return;
 
   // Extract the neighbor number
-  const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
+  const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
   const int ninside = d_ninside(ii);
   if (jj >= ninside) return;
 
@@ -577,11 +579,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeDeidrj,const type
   SNAKokkos<DeviceType> my_sna = snaKK;
 
   // Extract the atom number
-  int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
-  if (ii >= inum) return;
+  int ii = team.team_rank() + team.team_size() * (team.league_rank() %
+           ((chunk_size+team.team_size()-1)/team.team_size()));
+  if (ii >= chunk_size) return;
 
   // Extract the neighbor number
-  const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
+  const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
   const int ninside = d_ninside(ii);
   if (jj >= ninside) return;
 

src/SNAP/pair_snap.cpp

@@ -635,6 +635,7 @@ void PairSNAP::read_files(char *coefffilename, char *paramfilename)
   switchflag = 1;
   bzeroflag = 1;
   quadraticflag = 0;
+  chunksize = 2000;
 
   // open SNAP parameter file on proc 0
 
@@ -698,6 +699,8 @@ void PairSNAP::read_files(char *coefffilename, char *paramfilename)
       bzeroflag = atoi(keyval);
     else if (strcmp(keywd,"quadraticflag") == 0)
       quadraticflag = atoi(keyval);
+    else if (strcmp(keywd,"chunksize") == 0)
+      chunksize = atoi(keyval);
     else
       error->all(FLERR,"Incorrect SNAP parameter file");
   }

src/SNAP/pair_snap.h

@@ -59,6 +59,7 @@ protected:
   double** bispectrum;          // bispectrum components for all atoms in list
   int *map;                     // mapping from atom types to elements
   int twojmax, switchflag, bzeroflag;
+  int chunksize;
   double rfac0, rmin0, wj1, wj2;
   int rcutfacflag, twojmaxflag; // flags for required parameters
   int beta_max;                 // length of beta