Commit 4e5da1c9 authored by Sievers's avatar Sievers
Browse files

Chem SNAP KOKKOS implementation draft

parent 53e3421b

src/KOKKOS/pair_snap_kokkos_impl.h

+11 −4
Original line number Diff line number Diff line
@@ -243,6 +243,7 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in) 
        typename Kokkos::TeamPolicy<DeviceType,TagPairSNAPComputeUiCPU> policy_ui_cpu(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);



        Kokkos::parallel_for("ComputeUiCPU",policy_ui_cpu,*this);



      } else { // GPU, vector parallelism, shared memory, separate ulist and ulisttot to avoid atomics



        vector_length = 32;
@@ -514,7 +515,7 @@ void PairSNAPKokkos<DeviceType>::coeff(int narg, char **arg) 
  Kokkos::deep_copy(d_map,h_map);



  snaKK = SNAKokkos<DeviceType>(rfac0,twojmax,

                  rmin0,switchflag,bzeroflag);

                  rmin0,switchflag,bzeroflag,alloyflag,wselfallflag,nelements);

  snaKK.grow_rij(0,0);

  snaKK.init();

}
@@ -554,7 +555,6 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeNeigh,const typen 


      const int jtype = type(j);

      const F_FLOAT rsq = dx*dx + dy*dy + dz*dz;

      const int elem_j = d_map[jtype];



      if ( rsq < rnd_cutsq(itype,jtype) )

       count++;
@@ -584,6 +584,10 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeNeigh,const typen 
        my_sna.inside(ii,offset) = j;

        my_sna.wj(ii,offset) = d_wjelem[elem_j];

        my_sna.rcutij(ii,offset) = (radi + d_radelem[elem_j])*rcutfac;

        if (alloyflag)

          my_sna.element(ii,offset) = elem_j;

        else

          my_sna.element(ii,offset) = 0;

      }

      offset++;

    }
@@ -598,8 +602,10 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPPreUi,const typename Kok 
  // Extract the atom number

  const 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;

  int itype = type(ii);

  int ielem = d_map[itype];



  my_sna.pre_ui(team,ii);

  my_sna.pre_ui(team,ii,ielem);

}



template<class DeviceType>
@@ -649,7 +655,8 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPZeroYi,const typename Ko 
  const int ii = team.league_rank() / ((my_sna.idxu_max+team.team_size()-1)/team.team_size());

  if (ii >= chunk_size) return;



  my_sna.zero_yi(idx,ii);

  for(int ielem = 0; ielem < nelements; ielem++)

    my_sna.zero_yi(idx,ii,ielem);

}



template<class DeviceType>

src/KOKKOS/sna_kokkos.h

+13 −4
Original line number Diff line number Diff line
@@ -58,7 +58,7 @@ inline 
  SNAKokkos(const SNAKokkos<DeviceType>& sna, const typename Kokkos::TeamPolicy<DeviceType>::member_type& team);



inline

  SNAKokkos(double, int, double, int, int);

  SNAKokkos(double, int, double, int, int, int, int, int);



  KOKKOS_INLINE_FUNCTION

  ~SNAKokkos();
@@ -75,7 +75,7 @@ inline 


  // functions for bispectrum coefficients

  KOKKOS_INLINE_FUNCTION

  void pre_ui(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team,const int&); // ForceSNAP

  void pre_ui(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team,const int&, int); // ForceSNAP

  KOKKOS_INLINE_FUNCTION

  void compute_ui(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, const int, const int); // ForceSNAP

  KOKKOS_INLINE_FUNCTION
@@ -83,7 +83,7 @@ inline 
  KOKKOS_INLINE_FUNCTION

  void compute_zi(const int&);    // ForceSNAP

  KOKKOS_INLINE_FUNCTION

  void zero_yi(const int&,const int&); // ForceSNAP

  void zero_yi(const int&, const int&, int); // ForceSNAP

  KOKKOS_INLINE_FUNCTION

  void compute_yi(int,

   const Kokkos::View<F_FLOAT**, DeviceType> &beta); // ForceSNAP
@@ -131,6 +131,7 @@ inline 
  t_sna_2i inside;

  t_sna_2d wj;

  t_sna_2d rcutij;

  t_sna_2i element;

  t_sna_3d dedr;

  int natom, nmax;
@@ -185,7 +186,7 @@ inline 
  void init_rootpqarray();    // init()



  KOKKOS_INLINE_FUNCTION

  void add_uarraytot(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int, int, double, double, double); // compute_ui

  void add_uarraytot(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int, int, double, double, double, int); // compute_ui



  KOKKOS_INLINE_FUNCTION

  void compute_uarray_cpu(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int, int,
@@ -208,8 +209,16 @@ inline 
  // 1 = cosine

  int switch_flag;



  // Chem snap flags

  int alloy_flag;

  int bnorm_flag;

  int nelements;

  int ndoubles;

  int ntriples;



  // Self-weight

  double wself;

  int wselfall_flag;



  int bzero_flag; // 1 if bzero subtracted from barray

  Kokkos::View<double*, DeviceType> bzero; // array of B values for isolated atoms

src/KOKKOS/sna_kokkos_impl.h

+294 −215

File changed.

Preview size limit exceeded, changes collapsed.

src/SNAP/pair_snap.cpp

+0 −4

File changed.

Preview size limit exceeded, changes collapsed.

src/SNAP/sna.cpp

+10 −10

File changed.

Preview size limit exceeded, changes collapsed.

CRA Git | Maintained and supported by SUSTech CRA and CCSE