Commit befd4c8b authored by Evan Weinberg's avatar Evan Weinberg
Browse files

Optimizations to Compute[Yi/Zi/Bi], switching over to an AoSoA data layout on... 

Optimizations to Compute[Yi/Zi/Bi], switching over to an AoSoA data layout on the GPU. CPU vs GPU code paths are now maximally divergent, will include some discussion of that in PR. Small performance tweaks in Compute[UiTot/FusedDeidrj].
parent 7e2f29bb

src/KOKKOS/pair_snap_kokkos.h

+51 −12
Original line number Diff line number Diff line
@@ -30,19 +30,31 @@ PairStyle(snap/kk/host,PairSNAPKokkos<LMPHostType>) 


namespace LAMMPS_NS {



// Routines for both the CPU and GPU backend

template<int NEIGHFLAG, int EVFLAG>

struct TagPairSNAPComputeForce{};



struct TagPairSNAPBeta{};

struct TagPairSNAPComputeNeigh{};



// GPU backend only

struct TagPairSNAPPreUi{};

struct TagPairSNAPComputeUi{};

struct TagPairSNAPComputeUiCPU{};

struct TagPairSNAPTransformUi{}; // re-order ulisttot from SoA to AoSoA, zero ylist

struct TagPairSNAPComputeZi{};

struct TagPairSNAPBeta{};

struct TagPairSNAPComputeBi{};

struct TagPairSNAPZeroYi{};

struct TagPairSNAPTransformBi{}; // re-order blist from AoSoA to AoS

struct TagPairSNAPComputeYi{};

struct TagPairSNAPTransformYi{}; // re-order ylist from AoSoA to AoS

struct TagPairSNAPComputeFusedDeidrj{};



// CPU backend only

struct TagPairSNAPPreUiCPU{};

struct TagPairSNAPComputeUiCPU{};

struct TagPairSNAPComputeZiCPU{};

struct TagPairSNAPBetaCPU{};

struct TagPairSNAPComputeBiCPU{};

struct TagPairSNAPZeroYiCPU{};

struct TagPairSNAPComputeYiCPU{};

struct TagPairSNAPComputeDuidrjCPU{};

struct TagPairSNAPComputeDeidrjCPU{};
@@ -81,6 +93,10 @@ public: 
  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeNeigh,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeNeigh>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPBetaCPU,const int& ii) const;



  // GPU backend only

  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPPreUi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPPreUi>::member_type& team) const;
@@ -88,31 +104,53 @@ public: 
  void operator() (TagPairSNAPComputeUi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeUiCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUiCPU>::member_type& team) const;

  void operator() (TagPairSNAPTransformUi,const int iatom_mod, const int idxu, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeZi,const int iatom_mod, const int idxz, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeZi,const int& ii) const;

  void operator() (TagPairSNAPBeta, const int& ii) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeBi,const int iatom_mod, const int idxb, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeBi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeBi>::member_type& team) const;

  void operator() (TagPairSNAPTransformBi,const int iatom_mod, const int idxb, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPZeroYi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPZeroYi>::member_type& team) const;

  void operator() (TagPairSNAPComputeYi,const int iatom_mod, const int idxz, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeYi,const int& ii) const;

  void operator() (TagPairSNAPTransformYi,const int iatom_mod, const int idxu, const int iatom_div) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeFusedDeidrj,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeFusedDeidrj>::member_type& team) const;



  // CPU backend only

  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeDuidrjCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrjCPU>::member_type& team) const;

  void operator() (TagPairSNAPPreUiCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPPreUiCPU>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeDeidrjCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrjCPU>::member_type& team) const;

  void operator() (TagPairSNAPComputeUiCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUiCPU>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPBeta,const int& ii) const;

  void operator() (TagPairSNAPComputeZiCPU,const int& ii) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeBiCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeBiCPU>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPZeroYiCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPZeroYiCPU>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeYiCPU,const int& ii) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeDuidrjCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrjCPU>::member_type& team) const;



  KOKKOS_INLINE_FUNCTION

  void operator() (TagPairSNAPComputeDeidrjCPU,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrjCPU>::member_type& team) const;



  template<int NEIGHFLAG>

  KOKKOS_INLINE_FUNCTION
@@ -175,6 +213,7 @@ inline double dist2(double* x,double* y); 
  Kokkos::View<T_INT*, DeviceType> d_map;                    // mapping from atom types to elements

  Kokkos::View<T_INT*, DeviceType> d_ninside;                // ninside for all atoms in list

  Kokkos::View<F_FLOAT**, DeviceType> d_beta;                // betas for all atoms in list

  Kokkos::View<F_FLOAT***, Kokkos::LayoutLeft, DeviceType> d_beta_pack;          // betas for all atoms in list, GPU

  Kokkos::View<F_FLOAT**, DeviceType> d_bispectrum;          // bispectrum components for all atoms in list



  typedef Kokkos::DualView<F_FLOAT**, DeviceType> tdual_fparams;

src/KOKKOS/pair_snap_kokkos_impl.h

+393 −131

File changed.

Preview size limit exceeded, changes collapsed.

src/KOKKOS/sna_kokkos.h

+45 −19
Original line number Diff line number Diff line
@@ -36,7 +36,9 @@ public: 
  typedef Kokkos::View<double**, DeviceType> t_sna_2d;

  typedef Kokkos::View<double**, Kokkos::LayoutLeft, DeviceType> t_sna_2d_ll;

  typedef Kokkos::View<double***, DeviceType> t_sna_3d;

  typedef Kokkos::View<double***, Kokkos::LayoutLeft, DeviceType> t_sna_3d_ll;

  typedef Kokkos::View<double***[3], DeviceType> t_sna_4d;

  typedef Kokkos::View<double****, Kokkos::LayoutLeft, DeviceType> t_sna_4d_ll;

  typedef Kokkos::View<double**[3], DeviceType> t_sna_3d3;

  typedef Kokkos::View<double*****, DeviceType> t_sna_5d;
@@ -48,7 +50,8 @@ public: 
  typedef Kokkos::View<SNAcomplex***, DeviceType> t_sna_3c;

  typedef Kokkos::View<SNAcomplex***, Kokkos::LayoutLeft, DeviceType> t_sna_3c_ll;

  typedef Kokkos::View<SNAcomplex***[3], DeviceType> t_sna_4c;

  typedef Kokkos::View<SNAcomplex***[3], Kokkos::LayoutLeft, DeviceType> t_sna_4c_ll;

  typedef Kokkos::View<SNAcomplex***[3], Kokkos::LayoutLeft, DeviceType> t_sna_4c3_ll;

  typedef Kokkos::View<SNAcomplex****, Kokkos::LayoutLeft, DeviceType> t_sna_4c_ll;

  typedef Kokkos::View<SNAcomplex**[3], DeviceType> t_sna_3c3;

  typedef Kokkos::View<SNAcomplex*****, DeviceType> t_sna_5c;
@@ -73,27 +76,39 @@ inline 


  int ncoeff;



  // functions for bispectrum coefficients

  // functions for bispectrum coefficients, GPU only

  KOKKOS_INLINE_FUNCTION

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

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

  KOKKOS_INLINE_FUNCTION

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

  KOKKOS_INLINE_FUNCTION

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

  KOKKOS_INLINE_FUNCTION

  void compute_yi(int,int,int,

   const Kokkos::View<F_FLOAT***, Kokkos::LayoutLeft, DeviceType> &beta_pack); // ForceSNAP

  KOKKOS_INLINE_FUNCTION

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



  // functions for bispectrum coefficients, CPU only

  KOKKOS_INLINE_FUNCTION

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

  KOKKOS_INLINE_FUNCTION

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

  KOKKOS_INLINE_FUNCTION

  void compute_zi(const int&);    // ForceSNAP

  void compute_zi_cpu(const int&);    // ForceSNAP

  KOKKOS_INLINE_FUNCTION

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

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

  KOKKOS_INLINE_FUNCTION

  void compute_yi(int,

  void compute_yi_cpu(int,

   const Kokkos::View<F_FLOAT**, DeviceType> &beta); // ForceSNAP

    KOKKOS_INLINE_FUNCTION

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



  // functions for derivatives

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



  // functions for derivatives, GPU only

  KOKKOS_INLINE_FUNCTION

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



  // functions for derivatives, CPU only

  KOKKOS_INLINE_FUNCTION

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

  KOKKOS_INLINE_FUNCTION
@@ -139,18 +154,32 @@ inline 


  int twojmax, diagonalstyle;



  t_sna_2d_ll blist;

  t_sna_2c_ll ulisttot;

  t_sna_2c_ll zlist;

  t_sna_3d_ll blist;

  t_sna_3c_ll ulisttot;

  t_sna_3c_ll zlist;



  t_sna_3c_ll ulist;

  t_sna_2c_ll ylist;

  t_sna_3c_ll ylist;

  

  // derivatives of data

  t_sna_4c_ll dulist;

  t_sna_4c3_ll dulist;

  

  // Modified structures for GPU backend

  t_sna_3d_ll ulisttot_re; // split real,

  t_sna_3d_ll ulisttot_im; // imag

  t_sna_4c_ll ulisttot_pack; // AoSoA layout

  t_sna_4c_ll zlist_pack; // AoSoA layout

  t_sna_4d_ll blist_pack;

  t_sna_4d_ll ylist_pack_re; // split real, 

  t_sna_4d_ll ylist_pack_im; // imag AoSoA layout



  int idxcg_max, idxu_max, idxz_max, idxb_max;



  // Chem snap counts

  int nelements;

  int ndoubles;

  int ntriples;



private:

  double rmin0, rfac0;
@@ -212,9 +241,6 @@ inline 
  // Chem snap flags

  int chem_flag;

  int bnorm_flag;

  int nelements;

  int ndoubles;

  int ntriples;



  // Self-weight

  double wself;

src/KOKKOS/sna_kokkos_impl.h

+665 −331

File changed.

Preview size limit exceeded, changes collapsed.

CRA Git | Maintained and supported by SUSTech CRA and CCSE