Merge pull request #2 from simonepassera/main

Introduced optimizer threshold for GatherMPI. Fixed memory leaks on All-to-all generic implementation
ParaGroup · Nov 2, 2023 · 94b232b · 94b232b
2 parents f8a97fb + 7798daa
commit 94b232b
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 5 deletions.
diff --git a/include/collectives/collectiveImpl.hpp b/include/collectives/collectiveImpl.hpp
@@ -832,8 +832,12 @@ class AlltoallGeneric : public CollectiveImpl {
                         errno = ECONNRESET;
                         return -1;
                     }
+
+                    delete [] chunkbuff;
                 }
             }
+
+            delete [] allsendbuff;
 
             return selfrecvcount;
         } else {

diff --git a/include/collectives/mpiImpl.hpp b/include/collectives/mpiImpl.hpp
@@ -2,6 +2,7 @@
 #define MPICOLLIMPL_HPP
 
 #include "collectiveImpl.hpp"
+#include "../config.hpp"
 #include <mpi.h>
 #include <cassert>
 
@@ -289,14 +290,29 @@ class GatherMPI : public MPICollective {
         }
 
         size_t datacount = recvsize / datasize;
+
+        int recvcount = (datacount / nparticipants) * datasize;
+        int rcount = datacount % nparticipants;
+
+        if ((rcount == 0) && (recvcount >= GATHER_THRESHOLD_MSG_SIZE)) {
+            if ((size_t)recvcount > sendsize) {
+                MTCL_ERROR("[internal]:\t","sending buffer too small %ld instead of %ld\n", sendsize, recvcount);
+                errno = EINVAL;
+                return -1;
+            }
+
+            if (MPI_Gather(sendbuff, recvcount, MPI_BYTE, recvbuff, recvcount, MPI_BYTE, 0, comm) != MPI_SUCCESS) {
+                errno = ECOMM;
+                return -1;
+            }
+
+            return recvcount;
+        }
 
         int *recvcounts = new int[nparticipants]();
         int *displs = new int[nparticipants]();
 
         int displ = 0;
-
-        int recvcount = (datacount / nparticipants) * datasize;
-        int rcount = datacount % nparticipants;
 
         for (int i = 0; i < nparticipants; i++) {
             recvcounts[i] = recvcount;

diff --git a/include/config.hpp b/include/config.hpp
@@ -45,7 +45,7 @@ const unsigned UCX_POLL_TIMEOUT        = 10;
 
 // -------- COLLECTIVES ------
 const int CCONNECTION_RETRY            = 10;
-const unsigned CCONNECTION_TIMEOUT     = 100;  // milliseconds
-
+const unsigned CCONNECTION_TIMEOUT     = 100;     // milliseconds
+const int GATHER_THRESHOLD_MSG_SIZE    = (1<<18); // bytes
 
 #endif