1 // clang-format off 2 /* -*- c++ -*- ---------------------------------------------------------- 3 LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator 4 https://www.lammps.org/, Sandia National Laboratories 5 Steve Plimpton, sjplimp@sandia.gov 6 7 Copyright (2003) Sandia Corporation. Under the terms of Contract 8 DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains 9 certain rights in this software. This software is distributed under 10 the GNU General Public License. 11 12 See the README file in the top-level LAMMPS directory. 13 ------------------------------------------------------------------------- */ 14 15 #include "atom.h" // IWYU pragma: export 16 #include "kokkos_type.h" 17 18 #ifndef LMP_ATOM_KOKKOS_H 19 #define LMP_ATOM_KOKKOS_H 20 21 namespace LAMMPS_NS { 22 23 class AtomKokkos : public Atom { 24 public: 25 DAT::tdual_tagint_1d k_tag; 26 DAT::tdual_int_1d k_type, k_mask; 27 DAT::tdual_imageint_1d k_image; 28 DAT::tdual_x_array k_x; 29 DAT::tdual_v_array k_v; 30 DAT::tdual_f_array k_f; 31 32 DAT::tdual_float_1d k_mass; 33 34 DAT::tdual_float_1d k_q; 35 DAT::tdual_float_1d k_radius; 36 DAT::tdual_float_1d k_rmass; 37 DAT::tdual_v_array k_omega; 38 DAT::tdual_v_array k_angmom; 39 DAT::tdual_f_array k_torque; 40 DAT::tdual_tagint_1d k_molecule; 41 DAT::tdual_int_2d k_nspecial; 42 DAT::tdual_tagint_2d k_special; 43 DAT::tdual_int_1d k_num_bond; 44 DAT::tdual_int_2d k_bond_type; 45 DAT::tdual_tagint_2d k_bond_atom; 46 DAT::tdual_int_1d k_num_angle; 47 DAT::tdual_int_2d k_angle_type; 48 DAT::tdual_tagint_2d k_angle_atom1, k_angle_atom2, k_angle_atom3; 49 DAT::tdual_int_1d k_num_dihedral; 50 DAT::tdual_int_2d k_dihedral_type; 51 DAT::tdual_tagint_2d k_dihedral_atom1, k_dihedral_atom2, k_dihedral_atom3, k_dihedral_atom4; 52 DAT::tdual_int_1d k_num_improper; 53 DAT::tdual_int_2d k_improper_type; 54 DAT::tdual_tagint_2d k_improper_atom1, k_improper_atom2, k_improper_atom3, k_improper_atom4; 55 56 DAT::tdual_float_2d k_dvector; 57 58 // SPIN package 59 60 DAT::tdual_float_1d_4 k_sp; 61 DAT::tdual_f_array k_fm; 62 DAT::tdual_f_array k_fm_long; 63 64 // DPD-REACT package 65 DAT::tdual_efloat_1d k_uCond, k_uMech, k_uChem, k_uCG, k_uCGnew, 66 k_rho,k_dpdTheta,k_duChem; 67 68 69 AtomKokkos(class LAMMPS *); 70 virtual ~AtomKokkos(); 71 72 void map_init(int check = 1); 73 void map_set(); 74 void map_delete(); 75 76 DAT::tdual_int_1d k_sametag; 77 DAT::tdual_int_1d k_map_array; 78 DAT::tdual_int_scalar k_error_flag; 79 dual_hash_type k_map_hash; 80 81 // map lookup function inlined for efficiency 82 // return -1 if no map defined 83 84 template<class DeviceType> 85 KOKKOS_INLINE_FUNCTION map_kokkos(tagint global,int map_style,DAT::tdual_int_1d k_map_array,dual_hash_type k_map_hash)86 static int map_kokkos(tagint global, int map_style, DAT::tdual_int_1d k_map_array, dual_hash_type k_map_hash) 87 { 88 if (map_style == 1) 89 return k_map_array.view<DeviceType>()(global); 90 else if (map_style == 2) 91 return AtomKokkos::map_find_hash_kokkos<DeviceType>(global,k_map_hash); 92 else 93 return -1; 94 } 95 96 template<class DeviceType> 97 KOKKOS_INLINE_FUNCTION map_find_hash_kokkos(tagint global,dual_hash_type & k_map_hash)98 static int map_find_hash_kokkos(tagint global, dual_hash_type &k_map_hash) 99 { 100 int local = -1; 101 auto d_map_hash = k_map_hash.view<DeviceType>(); 102 auto index = d_map_hash.find(global); 103 if (d_map_hash.valid_at(index)) 104 local = d_map_hash.value_at(index); 105 return local; 106 } 107 108 virtual void allocate_type_arrays(); 109 void sync(const ExecutionSpace space, unsigned int mask); 110 void modified(const ExecutionSpace space, unsigned int mask); 111 void sync_overlapping_device(const ExecutionSpace space, unsigned int mask); 112 virtual void sort(); 113 virtual void grow(unsigned int mask); 114 int add_custom(const char *, int, int); 115 void remove_custom(int, int, int); 116 virtual void deallocate_topology(); 117 void sync_modify(ExecutionSpace, unsigned int, unsigned int); 118 private: 119 class AtomVec *new_avec(const std::string &, int, int &); 120 }; 121 122 template<class ViewType, class IndexView> 123 struct SortFunctor { 124 typedef typename ViewType::device_type device_type; 125 ViewType source; 126 Kokkos::View<typename ViewType::non_const_data_type,typename ViewType::array_type,device_type> dest; 127 IndexView index; SortFunctorSortFunctor128 SortFunctor(ViewType src, typename std::enable_if<ViewType::dynamic_rank==1,IndexView>::type ind):source(src),index(ind) { 129 dest = Kokkos::View<typename ViewType::non_const_data_type,typename ViewType::array_type,device_type>("",src.extent(0)); 130 } SortFunctorSortFunctor131 SortFunctor(ViewType src, typename std::enable_if<ViewType::dynamic_rank==2,IndexView>::type ind):source(src),index(ind) { 132 dest = Kokkos::View<typename ViewType::non_const_data_type,typename ViewType::array_type,device_type>("",src.extent(0),src.extent(1)); 133 } SortFunctorSortFunctor134 SortFunctor(ViewType src, typename std::enable_if<ViewType::dynamic_rank==3,IndexView>::type ind):source(src),index(ind) { 135 dest = Kokkos::View<typename ViewType::non_const_data_type,typename ViewType::array_type,device_type>("",src.extent(0),src.extent(1),src.extent(2)); 136 } SortFunctorSortFunctor137 SortFunctor(ViewType src, typename std::enable_if<ViewType::dynamic_rank==4,IndexView>::type ind):source(src),index(ind) { 138 dest = Kokkos::View<typename ViewType::non_const_data_type,typename ViewType::array_type,device_type>("",src.extent(0),src.extent(1),src.extent(2),src.extent(3)); 139 } 140 KOKKOS_INLINE_FUNCTION operatorSortFunctor141 void operator()(const typename std::enable_if<ViewType::rank==1, int>::type& i) { 142 dest(i) = source(index(i)); 143 } operatorSortFunctor144 void operator()(const typename std::enable_if<ViewType::rank==2, int>::type& i) { 145 for (int j=0; j < (int)source.extent(1); j++) 146 dest(i,j) = source(index(i),j); 147 } operatorSortFunctor148 void operator()(const typename std::enable_if<ViewType::rank==3, int>::type& i) { 149 for (int j=0; j < (int)source.extent(1); j++) 150 for (int k=0; k < (int)source.extent(2); k++) 151 dest(i,j,k) = source(index(i),j,k); 152 } operatorSortFunctor153 void operator()(const typename std::enable_if<ViewType::rank==4, int>::type& i) { 154 for (int j=0; j < (int)source.extent(1); j++) 155 for (int k=0; k < (int)source.extent(2); k++) 156 for (int l=0; l < (int)source.extent(3); l++) 157 dest(i,j,k,l) = source(index(i),j,k,l); 158 } 159 }; 160 161 } 162 163 #endif 164 165 /* ERROR/WARNING messages: 166 167 E: KOKKOS package requires a kokkos enabled atom_style 168 169 Self-explanatory. 170 171 */ 172