Skip to content

Commit

Permalink
Fix memory corruption in comm_tiled_kokkos
Browse files Browse the repository at this point in the history
  • Loading branch information
@stanmoore1
stanmoore1 committed Dec 18, 2024
1 parent aeb2190 commit 8cf03b0
Show file tree
Hide file tree
Showing 5 changed files with 57 additions and 43 deletions.
44 changes: 24 additions & 20 deletions src/KOKKOS/comm_tiled_kokkos.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ static constexpr int BUFEXTRA = 1000;
CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp) : CommTiled(_lmp)
{
sendlist = nullptr;
maxsendlist = nullptr;
nprocmaxtot = 0;
}

/* ---------------------------------------------------------------------- */
Expand All @@ -49,14 +51,18 @@ CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp) : CommTiled(_lmp)
CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp, Comm *oldcomm) : CommTiled(_lmp,oldcomm)
{
sendlist = nullptr;
maxsendlist = nullptr;
nprocmaxtot = 0;
}

/* ---------------------------------------------------------------------- */

CommTiledKokkos::~CommTiledKokkos()
{
memoryKK->destroy_kokkos(k_sendlist,sendlist);
memory->destroy(maxsendlist);
sendlist = nullptr;
maxsendlist = nullptr;
buf_send = nullptr;
buf_recv = nullptr;
}
Expand Down Expand Up @@ -657,12 +663,11 @@ void CommTiledKokkos::grow_list(int iswap, int iwhich, int n)
k_sendlist.sync<LMPHostType>();
k_sendlist.modify<LMPHostType>();

if (size > (int)k_sendlist.extent(2)) {
memoryKK->grow_kokkos(k_sendlist,sendlist,maxswap,maxsend,size,"comm:sendlist");
memoryKK->grow_kokkos(k_sendlist,sendlist,maxswap,nprocmaxtot,size,"comm:sendlist");

for (int i = 0; i < maxswap; i++)
maxsendlist[iswap][iwhich] = size;
}
for (int i = 0; i < maxswap; i++)
for (int j = 0; j < nprocmaxtot; j++)
maxsendlist[i][j] = size;
}

/* ----------------------------------------------------------------------
Expand Down Expand Up @@ -692,24 +697,23 @@ void CommTiledKokkos::grow_swap_send(int i, int n, int /*nold*/)
memory->destroy(sendbox_multiold[i]);
memory->create(sendbox_multiold[i],n,atom->ntypes+1,6,"comm:sendbox_multiold");

delete [] maxsendlist[i];
maxsendlist[i] = new int[n];

for (int j = 0; j < n; j++)
maxsendlist[i][j] = BUFMIN;

if (sendlist && !k_sendlist.d_view.data()) {
for (int ii = 0; ii < maxswap; ii++) {
if (sendlist[ii]) {
for (int jj = 0; jj < nprocmax[ii]; jj++)
memory->destroy(sendlist[ii][jj]);
delete [] sendlist[ii];
}
}
if (sendlist && !k_sendlist.h_view.data()) {
delete [] sendlist;
delete [] maxsendlist;

sendlist = nullptr;
maxsendlist = nullptr;
} else {
memoryKK->destroy_kokkos(k_sendlist,sendlist);
memory->destroy(maxsendlist);
}

memoryKK->create_kokkos(k_sendlist,sendlist,maxswap,n,BUFMIN,"comm:sendlist");
nprocmaxtot = MAX(nprocmaxtot,n);

memoryKK->create_kokkos(k_sendlist,sendlist,maxswap,nprocmaxtot,BUFMIN,"comm:sendlist");
memory->create(maxsendlist,maxswap,nprocmaxtot,"comm:maxsendlist");

for (int i = 0; i < maxswap; i++)
for (int j = 0; j < nprocmaxtot; j++)
maxsendlist[i][j] = BUFMIN;
}
11 changes: 5 additions & 6 deletions src/KOKKOS/comm_tiled_kokkos.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,18 +64,17 @@ class CommTiledKokkos : public CommTiled {
template<class DeviceType> void reverse_comm_device();

protected:
int nprocmaxtot;

DAT::tdual_int_3d k_sendlist;
//DAT::tdual_int_scalar k_total_send;
DAT::tdual_xfloat_2d k_buf_send,k_buf_recv;
//DAT::tdual_int_scalar k_count;

void grow_send(int, int) override;
void grow_recv(int, int flag = 0) override;
void grow_send(int, int) override; // reallocate send buffer
void grow_recv(int, int flag = 0) override; // free/allocate recv buffer
void grow_send_kokkos(int, int, ExecutionSpace space = Host);
void grow_recv_kokkos(int, int, ExecutionSpace space = Host);
void grow_list(int, int, int) override;
void grow_swap_send(int, int, int) override; // grow swap arrays for send and recv
void grow_list(int, int, int) override; // reallocate sendlist for one swap/proc
void grow_swap_send(int, int, int) override; // grow swap arrays for send and recv
};

} // namespace LAMMPS_NS
Expand Down
40 changes: 25 additions & 15 deletions src/KOKKOS/memory_kokkos.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,15 +221,19 @@ TYPE create_kokkos(TYPE &data, typename TYPE::value_type ***&array,
int n1, int n2, int n3, const char *name)
{
data = TYPE(std::string(name),n1,n2,n3);
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n1 * n2;
typename TYPE::value_type **plane = (typename TYPE::value_type **) smalloc(nbytes,name);
nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
array = (typename TYPE::value_type ***) smalloc(nbytes,name);

bigint m;
for (int i = 0; i < n1; i++) {
if (n2 == 0) {
array[i] = nullptr;
} else {
nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n2;
array[i] = (typename TYPE::value_type **) smalloc(nbytes,name);
m = ((bigint) i) * n2;
array[i] = &plane[m];

for (int j = 0; j < n2; j++) {
if (n3 == 0)
array[i][j] = nullptr;
Expand All @@ -248,15 +252,19 @@ template <typename TYPE, typename HTYPE>
{
data = TYPE(std::string(name),n1,n2);
h_data = Kokkos::create_mirror_view(data);
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n1 * n2;
typename TYPE::value_type **plane = (typename TYPE::value_type **) smalloc(nbytes,name);
nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
array = (typename TYPE::value_type ***) smalloc(nbytes,name);

bigint m;
for (int i = 0; i < n1; i++) {
if (n2 == 0) {
array[i] = nullptr;
} else {
nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n2;
array[i] = (typename TYPE::value_type **) smalloc(nbytes,name);
m = ((bigint) i) * n2;
array[i] = &plane[m];

for (int j = 0; j < n2; j++) {
if (n3 == 0)
array[i][j] = nullptr;
Expand Down Expand Up @@ -288,15 +296,19 @@ TYPE grow_kokkos(TYPE &data, typename TYPE::value_type ***&array,
{
if (array == nullptr) return create_kokkos(data,array,n1,n2,n3,name);
data.resize(n1,n2,n3);
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
array = (typename TYPE::value_type ***) smalloc(nbytes,name);
bigint nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n1 * n2;
typename TYPE::value_type **plane = (typename TYPE::value_type **) srealloc(array[0],nbytes,name);
nbytes = ((bigint) sizeof(typename TYPE::value_type **)) * n1;
array = (typename TYPE::value_type ***) srealloc(array,nbytes,name);

bigint m;
for (int i = 0; i < n1; i++) {
if (n2 == 0) {
array[i] = nullptr;
} else {
nbytes = ((bigint) sizeof(typename TYPE::value_type *)) * n2;
array[i] = (typename TYPE::value_type **) smalloc(nbytes,name);
m = ((bigint) i) * n2;
array[i] = &plane[m];

for (int j = 0; j < n2; j++) {
if (n3 == 0)
array[i][j] = nullptr;
Expand All @@ -316,10 +328,9 @@ template <typename TYPE>
void destroy_kokkos(TYPE data, typename TYPE::value_type*** &array)
{
if (array == nullptr) return;
int n1 = data.extent(0);
for (int i = 0; i < n1; ++i)
sfree(array[i]);
data = TYPE();

sfree(array[0]);
sfree(array);
array = nullptr;
}
Expand Down Expand Up @@ -411,7 +422,7 @@ template <typename TYPE>
TYPE create_kokkos(TYPE &data, int n1, int n2, int n3, int n4, int n5 , int n6 ,const char *name)
{
data = TYPE();
data = TYPE(std::string(name) ,n1,n2,n3,n4,n5,n6);
data = TYPE(std::string(name),n1,n2,n3,n4,n5,n6);
return data;
}

Expand All @@ -420,4 +431,3 @@ TYPE create_kokkos(TYPE &data, int n1, int n2, int n3, int n4, int n5 , int n6 ,
}

#endif

3 changes: 2 additions & 1 deletion src/comm_tiled.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2507,7 +2507,8 @@ void CommTiled::deallocate_swap(int n)
memory->destroy(sendbox_multi[i]);
memory->destroy(sendbox_multiold[i]);

delete [] maxsendlist[i];
if (maxsendlist)
delete [] maxsendlist[i];

if (sendlist && sendlist[i]) {
for (int j = 0; j < nprocmax[i]; j++) memory->destroy(sendlist[i][j]);
Expand Down
2 changes: 1 addition & 1 deletion src/comm_tiled.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ class CommTiled : public Comm {
virtual void grow_swap_send(int, int, int); // grow swap arrays for send and recv
void grow_swap_send_multi(int, int); // grow multi swap arrays for send and recv
void grow_swap_recv(int, int);
void deallocate_swap(int); // deallocate swap arrays
void deallocate_swap(int); // deallocate swap arrays
};

} // namespace LAMMPS_NS
Expand Down

0 comments on commit 8cf03b0

Please sign in to comment.