Skip to content

Commit

Permalink
Revert "get rid of single_bond_grad, attempted transpose (unhelpful),…
Browse files Browse the repository at this point in the history 
… added off-diagonal setflag, need cleaning"

This reverts commit 25e94b3.
  • Loading branch information
@anjohan
anjohan committed May 26, 2023
1 parent 25e94b3 commit 700f2c8
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 121 deletions.
147 changes: 33 additions & 114 deletions lammps_plugins/kokkos/pair_flare_kokkos.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,9 +193,9 @@ void PairFLAREKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
+ 2 // evdwls, B2_norm2s
+ 0.5 // numneigh_short
+ max_neighs * (
n_max*4 // g and gT
+ n_harmonics*4 // Y and YT
//+ n_max*n_harmonics*3 // single_bond_grad
n_max*4 // g
+ n_harmonics*4 // Y
+ n_max*n_harmonics*3 // single_bond_grad
+ 3 // partial_forces
+ 0.5 // neighs_short
)
Expand Down Expand Up @@ -258,12 +258,8 @@ void PairFLAREKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
g_ra = g;
Y_ra = Y;

//gT = View4D(); YT = View4D();
//gT = View4D(Kokkos::ViewAllocateWithoutInitializing("FLARE: gT"), batch_size, max_neighs, 4, n_max);
//YT = View4D(Kokkos::ViewAllocateWithoutInitializing("FLARE: YT"), batch_size, max_neighs, 4, n_harmonics);

// single_bond_grad = View5D();
// single_bond_grad = View5D(Kokkos::ViewAllocateWithoutInitializing("FLARE: single_bond_grad"), batch_size, max_neighs, 3, n_max, n_harmonics);
single_bond_grad = View5D();
single_bond_grad = View5D(Kokkos::ViewAllocateWithoutInitializing("FLARE: single_bond_grad"), batch_size, max_neighs, 3, n_max, n_harmonics);
partial_forces = View3D();
partial_forces = View3D(Kokkos::ViewAllocateWithoutInitializing("FLARE: partial forces"), batch_size, max_neighs, 3);

Expand All @@ -281,25 +277,14 @@ void PairFLAREKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
*this
);

// transpose R and Y for later use
int g_size = ScratchView3D::shmem_size(n_max, 4, max_neighs);
int Y_size = ScratchView3D::shmem_size(n_harmonics, 4, max_neighs);
/*
auto transpolicy = Kokkos::TeamPolicy<DeviceType, TagTransposeRY>(batch_size, SINGLE_BOND_TEAM_SIZE, vector_length).set_scratch_size(
0, Kokkos::PerTeam(g_size + Y_size));
Kokkos::parallel_for("FLARE: transpose R and Y",
transpolicy,
*this
);
*/

// compute single bond and its gradient
// dnlm, dnlmj
g_size = ScratchView1D::shmem_size(n_max);
Y_size = ScratchView1D::shmem_size(n_harmonics);
int g_size = ScratchView2D::shmem_size(n_max, 4);
int Y_size = ScratchView2D::shmem_size(n_harmonics, 4);
auto policy = Kokkos::TeamPolicy<DeviceType, TagSingleBond>(batch_size, SINGLE_BOND_TEAM_SIZE, vector_length).set_scratch_size(
0, Kokkos::PerThread(g_size + Y_size));
Kokkos::deep_copy(single_bond, 0.0);
//Kokkos::deep_copy(single_bond_grad, 0.0);
Kokkos::parallel_for("FLARE: single bond",
policy,
*this
Expand Down Expand Up @@ -343,15 +328,11 @@ void PairFLAREKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
);

// compute partial forces
g_size = ScratchView2D::shmem_size(4, n_max);
Y_size = ScratchView2D::shmem_size(4, n_harmonics);
int u_size = ScratchView2D::shmem_size(n_radial, n_harmonics);
Kokkos::parallel_for("FLARE: partial forces",
Kokkos::TeamPolicy<DeviceType, TagF>(batch_size, TEAM_SIZE, vector_length).set_scratch_size(
0, Kokkos::PerTeam(u_size)
)/*.set_scratch_size(
0, Kokkos::PerThread(g_size + Y_size)
)*/,
),
*this
);

Expand Down Expand Up @@ -430,113 +411,63 @@ void PairFLAREKokkos<DeviceType>::operator()(const int ii, const int jj) const {

template <class DeviceType>
KOKKOS_INLINE_FUNCTION
void PairFLAREKokkos<DeviceType>::operator()(TagTransposeRY, const MemberType team_member) const{
void PairFLAREKokkos<DeviceType>::operator()(TagSingleBond, const MemberType team_member) const{
int ii = team_member.league_rank();

const int jnum = d_numneigh_short(ii);

ScratchView3D gscratch(team_member.team_scratch(0), 4, n_max, max_neighs);
ScratchView3D Yscratch(team_member.team_scratch(0), 4, n_harmonics, max_neighs);
ScratchView2D gscratch(team_member.thread_scratch(0), 4, n_max);
ScratchView2D Yscratch(team_member.thread_scratch(0), 4, n_harmonics);

Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, jnum), [&] (int jj){

Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, 4*n_max), [&] (int nc){
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, jnum), [&] (int jj){
//int n = nc / 4;
//int c = nc -4*n;
int c = nc / n_max;
int n = nc - c*n_max;
gscratch(c, n, jj) = g(ii, jj, n, c);
});
});
Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, 4*n_harmonics), [&] (int lmc){
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, jnum), [&] (int jj){
//int lm = lmc / 4;
//int c = lmc - 4 * lm;
int c = lmc / n_harmonics;
int lm = lmc - c*n_harmonics;
Yscratch(c, lm, jj) = Y(ii, jj, lm, c);
});
});
team_member.team_barrier();
int j = d_neighbors_short(ii,jj);
j &= NEIGHMASK;
int s = type[j] - 1;

Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, jnum), [&] (int jj){

Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, 4*n_max), [&] (int nc){
//int n = nc / 4;
//int c = nc -4*n;
int c = nc / n_max;
int n = nc - c*n_max;
gT(ii, jj, c, n) = gscratch(c, n, jj);
gscratch(c, n) = g_ra(ii, jj, n, c);
});
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, 4*n_harmonics), [&] (int lmc){
//int lm = lmc / 4;
//int c = lmc - 4 * lm;
int c = lmc / n_harmonics;
int lm = lmc - c*n_harmonics;
YT(ii, jj, c, lm) = Yscratch(c, lm, jj);
});
});
}

template <class DeviceType>
KOKKOS_INLINE_FUNCTION
void PairFLAREKokkos<DeviceType>::operator()(TagSingleBond, const MemberType team_member) const{
int ii = team_member.league_rank();

const int jnum = d_numneigh_short(ii);

ScratchView1D gscratch(team_member.thread_scratch(0), n_max);
ScratchView1D Yscratch(team_member.thread_scratch(0), n_harmonics);

Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, jnum), [&] (int jj){

int j = d_neighbors_short(ii,jj);
j &= NEIGHMASK;
int s = type[j] - 1;


Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, n_max), [&] (int n){
//int n = nc / 4;
//int c = nc -4*n;
// int c = nc / n_max;
// int n = nc - c*n_max;
gscratch(n) = g(ii, jj, n, 0);
});
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, n_harmonics), [&] (int lm){
//int lm = lmc / 4;
//int c = lmc - 4 * lm;
// int c = lmc / n_harmonics;
// int lm = lmc - c*n_harmonics;
Yscratch(lm) = Y(ii, jj, lm, 0);
Yscratch(c, lm) = Y_ra(ii, jj, lm, c);
});

Kokkos::parallel_for(Kokkos::ThreadVectorRange(team_member, n_max*n_harmonics), [&] (int nlm){
int n = nlm / n_harmonics;
int lm = nlm - n_harmonics*n;

int radial_index = s*n_max + n;
double g_val = gscratch(n);
// double gx_val = gscratch(1,n);
// double gy_val = gscratch(2,n);
// double gz_val = gscratch(3,n);
double g_val = gscratch(0,n);
double gx_val = gscratch(1,n);
double gy_val = gscratch(2,n);
double gz_val = gscratch(3,n);


double h_val = Yscratch(lm);
// double hx_val = Yscratch(1,lm);
// double hy_val = Yscratch(2,lm);
// double hz_val = Yscratch(3,lm);
double h_val = Yscratch(0,lm);
double hx_val = Yscratch(1,lm);
double hy_val = Yscratch(2,lm);
double hz_val = Yscratch(3,lm);

double bond = g_val * h_val;
// double bond_x = gx_val * h_val + g_val * hx_val;
// double bond_y = gy_val * h_val + g_val * hy_val;
// double bond_z = gz_val * h_val + g_val * hz_val;
double bond_x = gx_val * h_val + g_val * hx_val;
double bond_y = gy_val * h_val + g_val * hy_val;
double bond_z = gz_val * h_val + g_val * hz_val;

// Update single bond basis arrays.
Kokkos::atomic_add(&single_bond(ii, radial_index, lm),bond); // TODO: bad?

// single_bond_grad(ii,jj,0,n,lm) = bond_x;
// single_bond_grad(ii,jj,1,n,lm) = bond_y;
// single_bond_grad(ii,jj,2,n,lm) = bond_z;
single_bond_grad(ii,jj,0,n,lm) = bond_x;
single_bond_grad(ii,jj,1,n,lm) = bond_y;
single_bond_grad(ii,jj,2,n,lm) = bond_z;
});
});
}
Expand Down Expand Up @@ -659,9 +590,6 @@ void PairFLAREKokkos<DeviceType>::operator()(TagF, const MemberType team_member)
const int i = ilist_curr_type[ii+startatom];
const int jnum = d_numneigh_short(ii);

//ScratchView2D gscratch(team_member.thread_scratch(0), 4, n_max);
//ScratchView2D Yscratch(team_member.thread_scratch(0), 4, n_harmonics);

ScratchView2D uscratch(team_member.team_scratch(0), n_radial, n_harmonics);
Kokkos::parallel_for(Kokkos::TeamVectorRange(team_member, n_bond), [&] (int nlm){
int n = nlm / n_harmonics;
Expand All @@ -683,16 +611,7 @@ void PairFLAREKokkos<DeviceType>::operator()(TagF, const MemberType team_member)
int n = nlm / n_harmonics;
int lm = nlm - n*n_harmonics;
int radial_index = s*n_max + n;

double gval = g(ii, jj, n, 0);
double gg = g(ii, jj, n, c+1);

double Yval = Y(ii, jj, lm, 0);
double Yg = Y(ii, jj, lm, c+1);

tmp += (gg*Yval + gval*Yg) * uscratch(radial_index, lm);

// tmp += single_bond_grad(ii, jj, c, n, lm)*uscratch(radial_index, lm);
tmp += single_bond_grad(ii, jj, c, n, lm)*uscratch(radial_index, lm);
}, tmp);
partial_forces(ii,jj,c) = tmp;
});
Expand Down
5 changes: 0 additions & 5 deletions lammps_plugins/kokkos/pair_flare_kokkos.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ struct Tagw{};
struct Tagu{};
struct TagF{};
struct TagStoreF{};
struct TagTransposeRY{};

namespace LAMMPS_NS {

Expand All @@ -58,9 +57,6 @@ class PairFLAREKokkos : public PairFLARE {
KOKKOS_INLINE_FUNCTION
void operator()(TagFindCurrType, const int) const;

KOKKOS_INLINE_FUNCTION
void operator()(TagTransposeRY, const MemberType) const;

KOKKOS_INLINE_FUNCTION
void operator()(TagSingleBond, const MemberType) const;

Expand Down Expand Up @@ -140,7 +136,6 @@ class PairFLAREKokkos : public PairFLARE {
View3D beta, single_bond, u, partial_forces;
gYView4D g, Y;
gYView4DRA g_ra, Y_ra;
View4D gT, YT;
View5D single_bond_grad;

int B2_chunk_size;
Expand Down
4 changes: 2 additions & 2 deletions lammps_plugins/pair_flare.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,7 @@ void PairFLARE::allocate() {

// Set the diagonal of setflag to 1 (otherwise pair.cpp will throw an error)
for (int i = 1; i <= n; i++)
for(int j = 1; j <= n; j++)
setflag[i][j] = 1;
setflag[i][i] = 1;
}

/* ----------------------------------------------------------------------
Expand All @@ -211,6 +210,7 @@ void PairFLARE::settings(int narg, char ** /*arg*/) {

/* ----------------------------------------------------------------------
set coeffs for one or more type pairs
read DYNAMO funcfl file
------------------------------------------------------------------------- */

void PairFLARE::coeff(int narg, char **arg) {
Expand Down

0 comments on commit 700f2c8

Please sign in to comment.