Fixed up kokkos version

This commit is contained in:
Aidan Thompson
2023-06-01 18:51:31 -06:00
parent 710597a353
commit 0d875d4f8f
2 changed files with 22 additions and 12 deletions

View File

@ -83,10 +83,11 @@ void ComputeLocalCompAtomKokkos<DeviceType>::compute_peratom()
// grow result array if necessary // grow result array if necessary
int size_peratom_cols = 1 + atom->ntypes;
if (atom->nmax > nmax) { if (atom->nmax > nmax) {
memoryKK->destroy_kokkos(k_result,result); memoryKK->destroy_kokkos(k_result,result);
nmax = atom->nmax; nmax = atom->nmax;
memoryKK->create_kokkos(k_result,result,nmax,2,"local/comp/atom:result"); memoryKK->create_kokkos(k_result,result,nmax,size_peratom_cols,"local/comp/atom:result");
d_result = k_result.view<DeviceType>(); d_result = k_result.view<DeviceType>();
array_atom = result; array_atom = result;
} }
@ -107,13 +108,12 @@ void ComputeLocalCompAtomKokkos<DeviceType>::compute_peratom()
atomKK->sync(execution_space,X_MASK|V_MASK|RMASS_MASK|TYPE_MASK|MASK_MASK); atomKK->sync(execution_space,X_MASK|V_MASK|RMASS_MASK|TYPE_MASK|MASK_MASK);
x = atomKK->k_x.view<DeviceType>(); x = atomKK->k_x.view<DeviceType>();
type = atomKK->k_type.view<DeviceType>(); type = atomKK->k_type.view<DeviceType>();
ntypes = atomKK->k_ntypes.view<DeviceType>();
mask = atomKK->k_mask.view<DeviceType>(); mask = atomKK->k_mask.view<DeviceType>();
Kokkos::deep_copy(d_result,0.0); Kokkos::deep_copy(d_result,0.0);
copymode = 1; copymode = 1;
typename Kokkos::RangePolicy<DeviceType, TagComputeLocalComp> policy(0,inum); typename Kokkos::RangePolicy<DeviceType, TagComputeLocalCompAtom> policy(0,inum);
Kokkos::parallel_for("ComputeLocalComp",policy,*this); Kokkos::parallel_for("ComputeLocalComp",policy,*this);
copymode = 0; copymode = 0;
@ -123,15 +123,17 @@ void ComputeLocalCompAtomKokkos<DeviceType>::compute_peratom()
template<class DeviceType> template<class DeviceType>
KOKKOS_INLINE_FUNCTION KOKKOS_INLINE_FUNCTION
void ComputeLocalCompAtomKokkos<DeviceType>::operator()(TagComputeLocalComp, const int &ii) const void ComputeLocalCompAtomKokkos<DeviceType>::operator()(TagComputeLocalCompAtom, const int &ii) const
{ {
double typeone_i, typeone_j; double typeone_i, typeone_j;
int ntypes = atom->ntypes;
int lcomp[ntypes]; int lcomp[ntypes];
// get per-atom local compositions // get per-atom local compositions
for (ii = 0; ii < inum; ii++) { int inum = list->inum;
for (int ii = 0; ii < inum; ii++) {
for (int i = 0; i < ntypes; i++) lcomp[i] = 0; for (int i = 0; i < ntypes; i++) lcomp[i] = 0;
@ -148,13 +150,14 @@ void ComputeLocalCompAtomKokkos<DeviceType>::operator()(TagComputeLocalComp, con
int count = 1; int count = 1;
itype = type[i]; int itype = type[i];
lcomp[itype-1]++;
for (int jj = 0; jj < jnum; jj++) { for (int jj = 0; jj < jnum; jj++) {
int j = d_neighbors(i,jj); int j = d_neighbors(i,jj);
j &= NEIGHMASK; j &= NEIGHMASK;
jtype = type[j]; int jtype = type[j];
const F_FLOAT delx = x(j,0) - xtmp; const F_FLOAT delx = x(j,0) - xtmp;
const F_FLOAT dely = x(j,1) - ytmp; const F_FLOAT dely = x(j,1) - ytmp;
@ -162,12 +165,20 @@ void ComputeLocalCompAtomKokkos<DeviceType>::operator()(TagComputeLocalComp, con
const F_FLOAT rsq = delx*delx + dely*dely + delz*delz; const F_FLOAT rsq = delx*delx + dely*dely + delz*delz;
if (rsq < cutsq) { if (rsq < cutsq) {
count++; count++;
lcomp[jtype]++; lcomp[jtype-1]++;
} }
} }
for (int n = 0 n < 4; n++) {
d_result(i,n+1) = lcomp[n] / density; // total count of atoms found in sampled radius range
}
d_result(i,0) = count;
// local comp fractions per element
double lfac = 1.0 / count;
for (int n = 0; n < ntypes; n++)
d_result(i,n+1) = lcomp[n] * lfac;
} }
} }
} }

View File

@ -48,7 +48,6 @@ template <class DeviceType> class ComputeLocalCompAtomKokkos : public ComputeLoc
private: private:
typename AT::t_x_array x; typename AT::t_x_array x;
typename AT::t_v_array v;
typename ArrayTypes<DeviceType>::t_int_1d type; typename ArrayTypes<DeviceType>::t_int_1d type;
typename ArrayTypes<DeviceType>::t_int_1d mask; typename ArrayTypes<DeviceType>::t_int_1d mask;