From c37b8147bf9c70ea9f7ce39f4a9c2a8554863bc0 Mon Sep 17 00:00:00 2001
From: abkein <rickbatra0z@gmail.com>
Date: Mon, 20 Jan 2025 19:38:42 +0300
Subject: [PATCH] Made special type for cldata for less errors

---
 src/compute_cluster_size_ext.cpp | 46 +++++++++++++++++++++-----------
 src/compute_cluster_size_ext.h   | 13 +++++++--
 2 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/src/compute_cluster_size_ext.cpp b/src/compute_cluster_size_ext.cpp
index 0dad89b05a8..1b7878c28ef 100644
--- a/src/compute_cluster_size_ext.cpp
+++ b/src/compute_cluster_size_ext.cpp
@@ -82,12 +82,26 @@ ComputeClusterSizeExt::ComputeClusterSizeExt(LAMMPS* lmp, int narg, char** arg)
 
   cIDs_by_size.reserve(size_cutoff);
   cIDs_by_size_all.reserve(size_cutoff);
+
+
+  MPI_Datatype type[2] = {MPI_INT, MPI_INT};
+  int blocklen[2] = {1, 1};
+  MPI_Aint disp[2];
+
+  // Calculate displacements
+  disp[0] = offsetof(cldata, id);
+  disp[1] = offsetof(cldata, sz);
+
+  MPI_Type_create_struct(2, blocklen, disp, type, &MPI_CLDATA);
+  MPI_Type_commit(&MPI_CLDATA);
 }
 
 /* ---------------------------------------------------------------------- */
 
 ComputeClusterSizeExt::~ComputeClusterSizeExt() noexcept(true)
 {
+  MPI_Type_free(&MPI_CLDATA);
+
   dist.destroy(memory);
   dist_local.destroy(memory);
   counts_global.destroy(memory);
@@ -131,8 +145,8 @@ void ComputeClusterSizeExt::init()
 
   if (ns.empty() || clusters.empty() || monomers.empty()) { error->one(FLERR, "{}: Inconsistent arrays state", style); }
 
-  if ((gathered.empty()) || (natom_loc < 2 * atom->natoms)) {
-    natom_loc = static_cast<bigint>(static_cast<long double>(2 * atom->natoms) * LMP_NUCC_ALLOC_COEFF);
+  if ((gathered.empty()) || (natom_loc < atom->natoms)) {
+    natom_loc = static_cast<bigint>(static_cast<long double>(atom->natoms) * LMP_NUCC_ALLOC_COEFF);
     gathered.grow(memory, natom_loc, "size/cluster/ext:gathered");
   }
 
@@ -182,19 +196,19 @@ void ComputeClusterSizeExt::compute_vector()
       if (cmap.count(clid) == 0) {
         const int clidx = cmap.size();
         cmap[clid] = clidx;
-        ns[2 * clidx] = clid;
-        ns[2 * clidx + 1] = 0;
+        ns[clidx].id = clid;
+        ns[clidx].sz = 0;
         clusters[clidx] = cluster_data(clid);
       }
       // possible segfault if actual cluster size exceeds LMP_NUCC_CLUSTER_MAX_SIZE + LMP_NUCC_CLUSTER_MAX_GHOST
       const int clidx = cmap[clid];
-      clusters[clidx].atoms<false>()[ns[2 * clidx + 1]++] = i;
+      clusters[clidx].atoms<false>()[ns[clidx].sz++] = i;
     }
   }
   for (const auto& [clid, clidx] : cmap) {
     cluster_data& clstr = clusters[clidx];
     const auto clatoms = clstr.atoms();
-    for (int i = 0; i < ns[2 * clidx + 1]; ++i) {
+    for (int i = 0; i < ns[clidx].sz; ++i) {
       if (clatoms[i] >= atom->nlocal) { error->one(FLERR, "{}/compute_vector_1:{}: particle index exceeds nlocal", style, comm->me); }
     }
   }
@@ -220,8 +234,8 @@ void ComputeClusterSizeExt::compute_vector()
   // }
 
   // communicate about number of unique clusters
-  const int ncluster_local = 2 * cmap.size();
-  ::MPI_Allgather(&ncluster_local, 1, MPI_INT, counts_global.data(), 1, MPI_INT, world);
+  const int ncluster_local = cmap.size();
+  ::MPI_Allgather(&ncluster_local, 1, MPI_CLDATA, counts_global.data(), 1, MPI_CLDATA, world);
 
   int tcon = counts_global[0];
   for (int i = 1; i < comm->nprocs; ++i) {
@@ -235,19 +249,19 @@ void ComputeClusterSizeExt::compute_vector()
   }
 
   // communicate about local cluster sizes
-  ::MPI_Allgatherv(ns.data(), ncluster_local, MPI_INT, gathered.data(), counts_global.data(), displs.data(), MPI_INT, world);
+  ::MPI_Allgatherv(ns.data(), ncluster_local, MPI_CLDATA, gathered.data(), counts_global.data(), displs.data(), MPI_CLDATA, world);
 
   // fill local data
   for (int i = 0; i < comm->nprocs; ++i) {
-    for (int j = 0; j < counts_global[i] - 1; j += 2) {
+    for (int j = 0; j < counts_global[i]; ++j) {
       int const k = displs[i] + j;
-      if (cmap.count(gathered[k]) > 0) {
-        cluster_data& clstr = clusters[cmap[gathered[k]]];
+      if (cmap.count(gathered[k].id) > 0) {
+        cluster_data& clstr = clusters[cmap[gathered[k].id]];
         if (i != comm->me) { clstr.owners<false>()[clstr.nowners++] = i; }
-        clstr.g_size += gathered[k + 1];
-        if (gathered[k + 1] > clstr.nhost) {
+        clstr.g_size += gathered[k].sz;
+        if (gathered[k].sz > clstr.nhost) {
           clstr.host = i;
-          clstr.nhost = gathered[k + 1];
+          clstr.nhost = gathered[k].sz;
         }
       }
     }
@@ -267,7 +281,7 @@ void ComputeClusterSizeExt::compute_vector()
   for (const auto& [clid, clidx] : cmap) {
     cluster_data& clstr = clusters[clidx];
     const auto clatoms = clstr.atoms();
-    clstr.l_size = ns[2 * clidx + 1];
+    clstr.l_size = ns[clidx].sz;
     for (int i = 0; i < clstr.l_size; ++i) {
       if (clatoms[i] >= atom->nlocal) { error->one(FLERR, "{}/compute_vector_3:{}: particle index exceeds nlocal", style, comm->me); }
     }
diff --git a/src/compute_cluster_size_ext.h b/src/compute_cluster_size_ext.h
index 1e34e0a4088..86b90a5b4f6 100644
--- a/src/compute_cluster_size_ext.h
+++ b/src/compute_cluster_size_ext.h
@@ -34,6 +34,13 @@ ComputeStyle(size/cluster/ext,ComputeClusterSizeExt);
 #    include <unordered_map>
 #    include <vector>
 
+namespace NUCC {
+  struct cldata {
+    int id = 0;
+    int sz = 0;
+  };
+}
+
 namespace LAMMPS_NS {
 
 class ComputeClusterSizeExt : public Compute {
@@ -81,8 +88,8 @@ class ComputeClusterSizeExt : public Compute {
   NUCC::cspan<int> counts_global;
   NUCC::cspan<int> displs;
   NUCC::cspan<NUCC::cluster_data> clusters;
-  NUCC::cspan<int> ns;
-  NUCC::cspan<int> gathered;
+  NUCC::cspan<NUCC::cldata> ns;
+  NUCC::cspan<NUCC::cldata> gathered;
   NUCC::cspan<double> peratom_size;
   bigint natom_loc = 0;
   int nonexclusive = 0;
@@ -92,6 +99,8 @@ class ComputeClusterSizeExt : public Compute {
   int nmono = 0;
 
   Compute *compute_cluster_atom = nullptr;
+
+  MPI_Datatype MPI_CLDATA;
 };
 
 }    // namespace LAMMPS_NS