Optimize MPI transfers using the TrivialCopyTraits

Author: fwyzard

HeterogeneousCore/MPICore/plugins/api.cc

@@ -192,4 +192,63 @@ void MPIChannel::receiveTrivialProduct_(int instance, edm::ObjectWithDict& produ
   MPI_Get_count(&status, MPI_BYTE, &size);
   assert(static_cast<int>(product.typeOf().size()) == size);
   MPI_Mrecv(product.address(), size, MPI_BYTE, &message, MPI_STATUS_IGNORE);
-}
+}
+
+// transfer a wrapped object using its TrivialCopyTraits
+void MPIChannel::sendTrivialCopyProduct_(int instance, edm::WrapperBase const* wrapper) {
+  int tag = EDM_MPI_SendTrivialCopyProduct | instance * EDM_MPI_MessageTagWidth_;
+
+  // if the wrapped type requires it, send the properties required toinitialise the remote copy
+  if (wrapper->hasTrivialCopyProperties()) {
+    edm::AnyBuffer buffer = wrapper->trivialCopyParameters();
+    MPI_Send(buffer.data(), buffer.size_bytes(), MPI_BYTE, dest_, tag, comm_);
+  }
+
+  // transfer the memory regions
+  auto regions = wrapper->trivialCopyRegions();
+  // TODO send the number of regions ?
+  for (size_t i = 0; i < regions.size(); ++i) {
+    assert(regions[i].data() != nullptr);
+    MPI_Send(regions[i].data(), regions[i].size_bytes(), MPI_BYTE, dest_, tag, comm_);
+  }
+}
+
+// receive a wrapped object using its TrivialCopyTraits
+void MPIChannel::receiveTrivialCopyProduct_(int instance, edm::WrapperBase* wrapper) {
+  int tag = EDM_MPI_SendTrivialCopyProduct | instance * EDM_MPI_MessageTagWidth_;
+
+  MPI_Message message;
+  MPI_Status status;
+  int size;
+
+  // mark the wrapped object as present
+  wrapper->markAsPresent();
+
+  // if the wrapped type requires it, send the properties required toinitialise the remote copy
+  if (wrapper->hasTrivialCopyProperties()) {
+    edm::AnyBuffer buffer = wrapper->trivialCopyParameters();
+    MPI_Mprobe(dest_, tag, comm_, &message, &status);
+    // check that the message size matches the expected buffer size
+    MPI_Get_count(&status, MPI_BYTE, &size);
+    assert(static_cast<int>(buffer.size_bytes()) == size);
+    // receive the properties
+    MPI_Mrecv(buffer.data(), buffer.size_bytes(), MPI_BYTE, &message, &status);
+    wrapper->trivialCopyInitialize(buffer);
+  }
+
+  // receive the memory regions
+  auto regions = wrapper->trivialCopyRegions();
+  // TODO receive and validate the number of regions ?
+  for (size_t i = 0; i < regions.size(); ++i) {
+    assert(regions[i].data() != nullptr);
+    MPI_Mprobe(dest_, tag, comm_, &message, &status);
+    // check that the message size matches the expected region size
+    MPI_Get_count(&status, MPI_BYTE, &size);
+    assert(static_cast<int>(regions[i].size_bytes()) == size);
+    // receive the data region
+    MPI_Mrecv(regions[i].data(), regions[i].size_bytes(), MPI_BYTE, &message, &status);
+  }
+
+  // finalize the clone after the trivialCopy, if the type requires it
+  wrapper->trivialCopyFinalize();
+}

HeterogeneousCore/MPICore/plugins/api.h

@@ -93,9 +93,13 @@ class MPIChannel {
     }
   }
 
-  // serialize a wrapped object using its ROOT dictionary, and transmit it
+  // transfer a wrapped object using the TrivialCopyTraits or its ROOT dictionary
   void sendProduct(int instance, edm::TypeWithDict const& type, edm::WrapperBase const& wrapper) {
-    sendSerializedProduct_(instance, type.getClass(), &wrapper);
+    if (wrapper.hasTrivialCopyTraits()) {
+      sendTrivialCopyProduct_(instance, &wrapper);
+    } else {
+      sendSerializedProduct_(instance, type.getClass(), &wrapper);
+    }
   }
 
   // signal that an expected product will not be transmitted
@@ -130,9 +134,13 @@ class MPIChannel {
     }
   }
 
-  // receive a wrapped object, and deserialize it using its ROOT dictionary
+  // receive a wrapped object using the TrivialCopyTraits or its ROOT dictionary
   void receiveProduct(int instance, edm::TypeWithDict const& type, edm::WrapperBase& wrapper) {
-    receiveSerializedProduct_(instance, type.getClass(), &wrapper);
+    if (wrapper.hasTrivialCopyTraits()) {
+      receiveTrivialCopyProduct_(instance, &wrapper);
+    } else {
+      receiveSerializedProduct_(instance, type.getClass(), &wrapper);
+    }
   }
 
 private:
@@ -194,6 +202,12 @@ class MPIChannel {
   // receive a binary blob, and deserialize an object of generic type using its ROOT dictionary
   void receiveSerializedProduct_(int instance, TClass const* type, void* product);
 
+  // transfer a wrapped object using its TrivialCopyTraits
+  void sendTrivialCopyProduct_(int instance, edm::WrapperBase const* wrapper);
+
+  // receive a wrapped object using its TrivialCopyTraits
+  void receiveTrivialCopyProduct_(int instance, edm::WrapperBase* wrapper);
+
   // MPI intercommunicator
   MPI_Comm comm_ = MPI_COMM_NULL;
 

HeterogeneousCore/MPICore/plugins/messages.h

@@ -23,6 +23,7 @@ enum EDM_MPI_MessageTag {
   EDM_MPI_ProcessEvent,
   EDM_MPI_SendSerializedProduct,
   EDM_MPI_SendTrivialProduct,
+  EDM_MPI_SendTrivialCopyProduct,
   EDM_MPI_SkipProduct,
   EDM_MPI_SendComplete,
   EDM_MPI_MessageTagCount_