Hi, I have been working on modifying our in house potential called bond valence model to be compatible with kokkos package. This in house potential requires two nested loops, the first nested loop(iterating over all the local atoms, for which then iterates over its neighbor atoms) helps us determine the bond valence for each atom. The second nested loop using the predetermined bond valence of the local and neighbor atom as well as the distance between the local and neighbor atom to determine their pairwise force, update the virial tensor, and update the energy.
However, the kokkos package does not support this. So I modified the some parts of the pair_kokkos.h file:
The first thing I have modified is adding a new functor structure named ‘PairComputeFunctorBv’ in which has the compute functions that are suitable to calculating the bond valence as well as pairwise forces, energy, and virial from the calculated bond valence. I have also adjusted alias definitions, constructor, and contribute function within the structure accordingly.
The second thing I have done is to change some part of the pair_compute_neighlist function.
Here is the relevant part in the pair_compute_neighlist function:
if (fpair->atom->ntypes > MAX_TYPES_STACKPARAMS) {
PairComputeFunctor<PairStyle,NEIGHFLAG,false,ZEROFLAG,Specialisation > ff(fpair,list);
GetMaxTeamSize<typename PairStyle::device_type>(ff, inum, teamsize_max_for, teamsize_max_reduce);
} else {
PairComputeFunctor<PairStyle,NEIGHFLAG,true,ZEROFLAG,Specialisation > ff(fpair,list);
GetMaxTeamSize<typename PairStyle::device_type>(ff, inum, teamsize_max_for, teamsize_max_reduce);
}
}
int teamsize_max = teamsize_max_for;
if (fpair->eflag || fpair->vflag)
teamsize_max = teamsize_max_reduce;
atoms_per_team = teamsize_max/vectorsize;
#else
vectorsize = 1;
atoms_per_team = 1;
#endif
const int num_teams = inum / atoms_per_team + (inum % atoms_per_team ? 1 : 0);
if (fpair->atom->ntypes > MAX_TYPES_STACKPARAMS) {
PairBvComputeFunctor<PairStyle,NEIGHFLAG,false,ZEROFLAG,Specialisation > ff(fpair,list);
Kokkos::TeamPolicy<typename PairStyle::device_type,Kokkos::IndexType<int> > policy(num_teams,atoms_per_team,vectorsize);
if (fpair->eflag || fpair->vflag) Kokkos::parallel_reduce(policy,ff,ev);
else Kokkos::parallel_for(policy,ff);
ff.contribute();
} else {
PairBvComputeFunctor<PairStyle,NEIGHFLAG,true,ZEROFLAG,Specialisation > ff(fpair,list);
Kokkos::TeamPolicy<typename PairStyle::device_type,Kokkos::IndexType<int> > policy(num_teams,atoms_per_team,vectorsize);
if (fpair->eflag || fpair->vflag) Kokkos::parallel_reduce(policy,ff,ev);
else Kokkos::parallel_for(policy,ff);
ff.contribute();
}
} else {
if (fpair->atom->ntypes > MAX_TYPES_STACKPARAMS) {
PairBvComputeFunctor<PairStyle,NEIGHFLAG,false,ZEROFLAG,Specialisation > ff(fpair,list);
if (fpair->eflag || fpair->vflag) Kokkos::parallel_reduce(inum,ff,ev);
else Kokkos::parallel_for(inum,ff);
ff.contribute();
} else {
PairBvComputeFunctor<PairStyle,NEIGHFLAG,true,ZEROFLAG,Specialisation > ff(fpair,list);
if (fpair->eflag || fpair->vflag) Kokkos::parallel_reduce(inum,ff,ev);
else Kokkos::parallel_for(inum,ff);
ff.contribute();
}
}
I used the original functor structure to initialize an instance of that structure that could be used to get the value for teamsize_max_for and teamsize_max_reduce, which helped me to determine the nums_team, atoms_per_team, and vectorsize. These help me to obtain the policy for parallel processing. But, the instance of the functor structure that I passed in as arguments for parallel_for and parallel_reduce to do the computation are from the PairComputeFunctorBv structure, which has the desired compute functions within it.
Those are the major changes I made on pair_kokkos.h file in order to accelerate my inhouse potential. I want to ask for some of your advice and cautions before I go straight to testing it.