util/memory: use mimalloc

include/faabric/util/memory.h

@@ -10,6 +10,13 @@
 
 namespace faabric::util {
 
+// We provide our own namespaced definitions for malloc/free to control the
+// memory allocator we use. For the moment, we just defer to off-the-shelve
+// malloc implementations.
+inline void* malloc(std::size_t size) { return std::malloc(size); }
+
+inline void free(void* ptr) { return std::free(ptr); }
+
 /*
  * Merges all the dirty page flags from the list of vectors into the first
  * vector in place.

src/mpi/MpiWorld.cpp

@@ -8,6 +8,7 @@
 #include <faabric/util/environment.h>
 #include <faabric/util/gids.h>
 #include <faabric/util/macros.h>
+#include <faabric/util/memory.h>
 #include <faabric/util/testing.h>
 
 // Each MPI rank runs in a separate thread, thus we use TLS to maintain the
@@ -526,13 +527,12 @@ void MpiWorld::send(int sendRank,
 
     // Dispatch the message locally or globally
     if (isLocal) {
-        void* bufferPtr = malloc(count * dataType->size);
-        std::memcpy(bufferPtr, buffer, count* dataType->size);
-
         if (mustSendData) {
+            void* bufferPtr = faabric::util::malloc(count * dataType->size);
+            std::memcpy(bufferPtr, buffer, count* dataType->size);
+
             m->set_bufferptr((uint64_t)bufferPtr);
         }
-        SPDLOG_INFO("Send (Ptr: {} - Size: {} - Data as int: {})", m->bufferptr(), count * dataType->size, ((int*)m->bufferptr())[0]);
 
         SPDLOG_TRACE(
           "MPI - send {} -> {} ({})", sendRank, recvRank, messageType);
@@ -611,9 +611,8 @@ void MpiWorld::doRecv(std::shared_ptr<MPIMessage>& m,
 
     if (m->count() > 0) {
         if (isLocal) {
-            SPDLOG_INFO("Recv (Ptr: {} - Size: {} - Data as int: {})", m->bufferptr(), count * dataType->size, ((int*)m->bufferptr())[0]);
             std::memcpy(buffer, (void*)m->bufferptr(), count * dataType->size);
-            free((void*)m->bufferptr());
+            faabric::util::free((void*)m->bufferptr());
         } else {
             // TODO - avoid copy here
             std::move(m->buffer().begin(), m->buffer().end(), buffer);

tests/dist/mpi/mpi_native.cpp

@@ -14,6 +14,7 @@
 #include <faabric/util/compare.h>
 #include <faabric/util/config.h>
 #include <faabric/util/logging.h>
+#include <faabric/util/memory.h>
 
 using namespace faabric::mpi;
 
@@ -508,7 +509,7 @@ int MPI_Alloc_mem(MPI_Aint size, MPI_Info info, void* baseptr)
         throw std::runtime_error("Non-null info not supported");
     }
 
-    *((void**)baseptr) = malloc(size);
+    *((void**)baseptr) = faabric::util::malloc(size);
 
     return MPI_SUCCESS;
 }
@@ -641,7 +642,7 @@ int MPI_Isend(const void* buf,
     SPDLOG_TRACE("MPI - MPI_Isend {} -> {}", executingContext.getRank(), dest);
 
     MpiWorld& world = getExecutingWorld();
-    (*request) = (faabric_request_t*)malloc(sizeof(faabric_request_t));
+    (*request) = (faabric_request_t*)faabric::util::malloc(sizeof(faabric_request_t));
     int requestId = world.isend(
       executingContext.getRank(), dest, (uint8_t*)buf, datatype, count);
     (*request)->id = requestId;
@@ -661,7 +662,7 @@ int MPI_Irecv(void* buf,
       "MPI - MPI_Irecv {} <- {}", executingContext.getRank(), source);
 
     MpiWorld& world = getExecutingWorld();
-    (*request) = (faabric_request_t*)malloc(sizeof(faabric_request_t));
+    (*request) = (faabric_request_t*)faabric::util::malloc(sizeof(faabric_request_t));
     int requestId = world.irecv(
       source, executingContext.getRank(), (uint8_t*)buf, datatype, count);
     (*request)->id = requestId;