diff --git a/src/KOKKOS/comm_kokkos.cpp b/src/KOKKOS/comm_kokkos.cpp index b586dca7a5..25045197b5 100644 --- a/src/KOKKOS/comm_kokkos.cpp +++ b/src/KOKKOS/comm_kokkos.cpp @@ -931,6 +931,7 @@ void CommKokkos::exchange_device() if (nextrarecv) { kkbase->unpack_exchange_kokkos( k_buf_recv,k_indices,nrecv/data_size, + nrecv1/data_size,nextrarecv1, ExecutionSpaceFromDevice::space); DeviceType().fence(); } diff --git a/src/KOKKOS/fix_neigh_history_kokkos.cpp b/src/KOKKOS/fix_neigh_history_kokkos.cpp index b4a852ba70..49fe3f1177 100644 --- a/src/KOKKOS/fix_neigh_history_kokkos.cpp +++ b/src/KOKKOS/fix_neigh_history_kokkos.cpp @@ -453,8 +453,12 @@ KOKKOS_INLINE_FUNCTION void FixNeighHistoryKokkos::operator()(TagFixNeighHistoryUnpackExchange, const int &i) const { int index = d_indices(i); + if (index > -1) { int m = (int) d_ubuf(d_buf(i)).i; + if (i >= nrecv1) + m = nextrarecv1 + (int) d_ubuf(d_buf(nextrarecv1 + i - nrecv1)).i; + int n = (int) d_ubuf(d_buf(m++)).i; d_npartner(index) = n; for (int p = 0; p < n; p++) { @@ -471,6 +475,7 @@ void FixNeighHistoryKokkos::operator()(TagFixNeighHistoryUnpackExcha template void FixNeighHistoryKokkos::unpack_exchange_kokkos( DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv, + int nrecv1, int nextrarecv1, ExecutionSpace /*space*/) { d_buf = typename AT::t_xfloat_1d_um( @@ -478,6 +483,9 @@ void FixNeighHistoryKokkos::unpack_exchange_kokkos( k_buf.extent(0)*k_buf.extent(1)); d_indices = k_indices.view(); + this->nrecv1 = nrecv1; + this->nextrarecv1 = nextrarecv1; + d_npartner = k_npartner.template view(); d_partner = k_partner.template view(); d_valuepartner = k_valuepartner.template view(); diff --git a/src/KOKKOS/fix_neigh_history_kokkos.h b/src/KOKKOS/fix_neigh_history_kokkos.h index 9c07a953c4..dd1ad769b8 100644 --- a/src/KOKKOS/fix_neigh_history_kokkos.h +++ b/src/KOKKOS/fix_neigh_history_kokkos.h @@ -72,12 +72,14 @@ class FixNeighHistoryKokkos : public FixNeighHistory, public KokkosBase { void unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &indices,int nrecv, + int nrecv1,int nrecv1extra, ExecutionSpace space) override; typename DAT::tdual_int_2d k_firstflag; typename DAT::tdual_float_2d k_firstvalue; private: + int nrecv1,nextrarecv1; int nlocal,nsend,beyond_contact; typename AT::t_tagint_1d tag; diff --git a/src/KOKKOS/fix_qeq_reaxff_kokkos.cpp b/src/KOKKOS/fix_qeq_reaxff_kokkos.cpp index a2a50d84bb..18d7af75a7 100644 --- a/src/KOKKOS/fix_qeq_reaxff_kokkos.cpp +++ b/src/KOKKOS/fix_qeq_reaxff_kokkos.cpp @@ -1416,6 +1416,7 @@ KOKKOS_INLINE_FUNCTION void FixQEqReaxFFKokkos::operator()(TagQEqUnpackExchange, const int &i) const { int index = d_indices(i); + if (index > -1) { for (int m = 0; m < nprev; m++) d_s_hist(index,m) = d_buf(i*nprev*2 + m); for (int m = 0; m < nprev; m++) d_t_hist(index,m) = d_buf(i*nprev*2 + nprev+m); @@ -1427,6 +1428,7 @@ void FixQEqReaxFFKokkos::operator()(TagQEqUnpackExchange, const int template void FixQEqReaxFFKokkos::unpack_exchange_kokkos( DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv, + int /*nrecv1*/, int /*nextrarecv1*/, ExecutionSpace /*space*/) { k_buf.sync(); diff --git a/src/KOKKOS/fix_qeq_reaxff_kokkos.h b/src/KOKKOS/fix_qeq_reaxff_kokkos.h index 9bc38b0492..6aa345fba6 100644 --- a/src/KOKKOS/fix_qeq_reaxff_kokkos.h +++ b/src/KOKKOS/fix_qeq_reaxff_kokkos.h @@ -143,6 +143,7 @@ class FixQEqReaxFFKokkos : public FixQEqReaxFF, public KokkosBase { void unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &indices,int nrecv, + int nrecv1,int nextrarecv1, ExecutionSpace space) override; struct params_qeq{ diff --git a/src/KOKKOS/fix_shake_kokkos.cpp b/src/KOKKOS/fix_shake_kokkos.cpp index dd6de8f9ec..7e9519da42 100644 --- a/src/KOKKOS/fix_shake_kokkos.cpp +++ b/src/KOKKOS/fix_shake_kokkos.cpp @@ -1581,8 +1581,8 @@ void FixShakeKokkos::pack_exchange_item(const int &mysend, int &offs else offset++; } else { - d_buf[mysend] = nsend + offset; int m = nsend + offset; + d_buf[mysend] = m; d_buf[m++] = flag; if (flag == 1) { d_buf[m++] = d_shake_atom(i,0); @@ -1703,6 +1703,8 @@ void FixShakeKokkos::operator()(TagFixShakeUnpackExchange, const int if (index > -1) { int m = d_buf[i]; + if (i >= nrecv1) + m = nextrarecv1 + d_buf[nextrarecv1 + i - nrecv1]; int flag = d_shake_flag[index] = static_cast (d_buf[m++]); if (flag == 1) { @@ -1739,6 +1741,7 @@ void FixShakeKokkos::operator()(TagFixShakeUnpackExchange, const int template void FixShakeKokkos::unpack_exchange_kokkos( DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv, + int nrecv1, int nextrarecv1, ExecutionSpace /*space*/) { k_buf.sync(); @@ -1749,6 +1752,9 @@ void FixShakeKokkos::unpack_exchange_kokkos( k_buf.extent(0)*k_buf.extent(1)); d_indices = k_indices.view(); + this->nrecv1 = nrecv1; + this->nextrarecv1 = nextrarecv1; + k_shake_flag.template sync(); k_shake_atom.template sync(); k_shake_type.template sync(); diff --git a/src/KOKKOS/fix_shake_kokkos.h b/src/KOKKOS/fix_shake_kokkos.h index 185e69ce86..efbecf5ac2 100644 --- a/src/KOKKOS/fix_shake_kokkos.h +++ b/src/KOKKOS/fix_shake_kokkos.h @@ -112,9 +112,12 @@ class FixShakeKokkos : public FixShake, public KokkosBase { void unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &indices,int nrecv, + int nrecv1,int nrecv1extra, ExecutionSpace space) override; protected: + int nrecv1,nextrarecv1; + typename AT::t_x_array d_x; typename AT::t_v_array d_v; typename AT::t_f_array d_f; @@ -259,4 +262,3 @@ struct FixShakeKokkosPackExchangeFunctor { #endif #endif - diff --git a/src/KOKKOS/fix_spring_self_kokkos.cpp b/src/KOKKOS/fix_spring_self_kokkos.cpp index efd8a652ff..b513d0bd37 100644 --- a/src/KOKKOS/fix_spring_self_kokkos.cpp +++ b/src/KOKKOS/fix_spring_self_kokkos.cpp @@ -188,8 +188,8 @@ void FixSpringSelfKokkos::pack_exchange_item(const int &mysend, int { const int i = d_exchange_sendlist(mysend); - d_buf[mysend] = nsend + offset; int m = nsend + offset; + d_buf[mysend] = m; d_buf[m++] = d_xoriginal(i,0); d_buf[m++] = d_xoriginal(i,1); d_buf[m++] = d_xoriginal(i,2); @@ -258,6 +258,8 @@ void FixSpringSelfKokkos::operator()(TagFixSpringSelfUnpackExchange, if (index > -1) { int m = d_buf[i]; + if (i >= nrecv1) + m = nextrarecv1 + d_buf[nextrarecv1 + i - nrecv1]; d_xoriginal(index,0) = static_cast (d_buf[m++]); d_xoriginal(index,1) = static_cast (d_buf[m++]); @@ -270,6 +272,7 @@ void FixSpringSelfKokkos::operator()(TagFixSpringSelfUnpackExchange, template void FixSpringSelfKokkos::unpack_exchange_kokkos( DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv, + int nrecv1, int nextrarecv1, ExecutionSpace /*space*/) { k_buf.sync(); @@ -280,6 +283,9 @@ void FixSpringSelfKokkos::unpack_exchange_kokkos( k_buf.extent(0)*k_buf.extent(1)); d_indices = k_indices.view(); + this->nrecv1 = nrecv1; + this->nextrarecv1 = nextrarecv1; + k_xoriginal.template sync(); copymode = 1; diff --git a/src/KOKKOS/fix_spring_self_kokkos.h b/src/KOKKOS/fix_spring_self_kokkos.h index b23e92249b..add5a80bc7 100644 --- a/src/KOKKOS/fix_spring_self_kokkos.h +++ b/src/KOKKOS/fix_spring_self_kokkos.h @@ -58,6 +58,7 @@ class FixSpringSelfKokkos : public FixSpringSelf, public KokkosBase { void unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &indices,int nrecv, + int nrecv1,int nrecv1extra, ExecutionSpace space) override; @@ -65,6 +66,8 @@ class FixSpringSelfKokkos : public FixSpringSelf, public KokkosBase { int unpack_exchange(int, double *) override; protected: + int nrecv1,nextrarecv1; + DAT::tdual_x_array k_xoriginal; typename AT::t_x_array d_xoriginal; diff --git a/src/KOKKOS/fix_wall_gran_kokkos.cpp b/src/KOKKOS/fix_wall_gran_kokkos.cpp index f870b0f240..25e405c798 100644 --- a/src/KOKKOS/fix_wall_gran_kokkos.cpp +++ b/src/KOKKOS/fix_wall_gran_kokkos.cpp @@ -419,6 +419,7 @@ void FixWallGranKokkos::operator()(TagFixWallGranUnpackExchange, con template void FixWallGranKokkos::unpack_exchange_kokkos( DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv, + int /*nrecv1*/, int /*nextrarecv1*/, ExecutionSpace /*space*/) { d_buf = typename ArrayTypes::t_xfloat_1d_um( @@ -430,7 +431,6 @@ void FixWallGranKokkos::unpack_exchange_kokkos( copymode = 1; - Kokkos::parallel_for(Kokkos::RangePolicy(0,nrecv),*this); copymode = 0; diff --git a/src/KOKKOS/fix_wall_gran_kokkos.h b/src/KOKKOS/fix_wall_gran_kokkos.h index c7d566ec72..ae54fdb085 100644 --- a/src/KOKKOS/fix_wall_gran_kokkos.h +++ b/src/KOKKOS/fix_wall_gran_kokkos.h @@ -62,12 +62,13 @@ class FixWallGranKokkos : public FixWallGranOld, public KokkosBase { void operator()(TagFixWallGranUnpackExchange, const int&) const; int pack_exchange_kokkos(const int &nsend,DAT::tdual_xfloat_2d &buf, - DAT::tdual_int_1d k_sendlist, - DAT::tdual_int_1d k_copylist, - ExecutionSpace space) override; + DAT::tdual_int_1d k_sendlist, + DAT::tdual_int_1d k_copylist, + ExecutionSpace space) override; void unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &indices,int nrecv, + int nrecv1,int nrecv1extra, ExecutionSpace space) override; private: @@ -91,6 +92,7 @@ class FixWallGranKokkos : public FixWallGranOld, public KokkosBase { typename AT::t_int_1d d_copylist; typename AT::t_int_1d d_indices; }; + } #endif diff --git a/src/KOKKOS/kokkos_base.h b/src/KOKKOS/kokkos_base.h index 1e22a38657..24fcc47579 100644 --- a/src/KOKKOS/kokkos_base.h +++ b/src/KOKKOS/kokkos_base.h @@ -47,6 +47,7 @@ class KokkosBase { ExecutionSpace /*space*/) { return 0; } virtual void unpack_exchange_kokkos(DAT::tdual_xfloat_2d & /*k_buf*/, DAT::tdual_int_1d & /*indices*/, int /*nrecv*/, + int /*nrecv1*/, int /*nextrarecv1*/, ExecutionSpace /*space*/) {} // Region