Track Message Type in Execution Graph

include/faabric/scheduler/MpiWorld.h

@@ -12,6 +12,13 @@
 #include <atomic>
 #include <unordered_map>
 
+// Constants for profiling MPI parameters like number of messages sent or
+// message breakdown by type in the execution graph. Remember to increase the
+// counter if you add another one
+#define NUM_MPI_EXEC_GRAPH_DETAILS 2
+#define MPI_MSG_COUNT_PREFIX "mpi-msgcount-torank-"
+#define MPI_MSGTYPE_COUNT_PREFIX "mpi-msgtype-torank"
+
 namespace faabric::scheduler {
 typedef faabric::util::Queue<std::shared_ptr<faabric::MPIMessage>>
   InMemoryMpiQueue;

include/faabric/util/exec_graph.h

@@ -7,13 +7,12 @@
 #include <map>
 
 namespace faabric::util::exec_graph {
+
 void addDetail(faabric::Message& msg,
                const std::string& key,
                const std::string& value);
 
 void incrementCounter(faabric::Message& msg,
                       const std::string& key,
                       const int valueToIncrement = 1);
-
-static inline std::string const mpiMsgCountPrefix = "mpi-msgcount-torank-";
 }

src/scheduler/MpiWorld.cpp

@@ -588,10 +588,16 @@ void MpiWorld::send(int sendRank,
 
     // If the message is set and recording on, track we have sent this message
     if (thisRankMsg != nullptr && thisRankMsg->recordexecgraph()) {
+        faabric::util::exec_graph::incrementCounter(
+          *thisRankMsg, MPI_MSG_COUNT_PREFIX + std::to_string(recvRank));
+
+        // Work out the message type breakdown
         faabric::util::exec_graph::incrementCounter(
           *thisRankMsg,
-          faabric::util::exec_graph::mpiMsgCountPrefix +
-            std::to_string(recvRank));
+          fmt::format("{}-{}-{}",
+                      MPI_MSGTYPE_COUNT_PREFIX,
+                      std::to_string(messageType),
+                      std::to_string(recvRank)));
     }
 }
 

tests/test/scheduler/test_mpi_exec_graph.cpp

@@ -23,8 +23,7 @@ TEST_CASE_METHOD(MpiTestFixture,
     auto buffer = new int[messageData.size()];
 
     int numToSend = 10;
-    std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rankA2);
+    std::string expectedKey = MPI_MSG_COUNT_PREFIX + std::to_string(rankA2);
 
     for (int i = 0; i < numToSend; i++) {
         world.send(rankA1,
@@ -36,7 +35,7 @@ TEST_CASE_METHOD(MpiTestFixture,
           rankA1, rankA2, BYTES(buffer), MPI_INT, messageData.size(), &status);
     }
 
-    REQUIRE(msg.intexecgraphdetails_size() == 1);
+    REQUIRE(msg.intexecgraphdetails_size() == NUM_MPI_EXEC_GRAPH_DETAILS);
     REQUIRE(msg.execgraphdetails_size() == 0);
     REQUIRE(msg.intexecgraphdetails().count(expectedKey) == 1);
     REQUIRE(msg.intexecgraphdetails().at(expectedKey) == numToSend);
@@ -58,8 +57,7 @@ TEST_CASE_METHOD(MpiTestFixture,
     auto buffer = new int[messageData.size()];
 
     int numToSend = 10;
-    std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rankA2);
+    std::string expectedKey = MPI_MSG_COUNT_PREFIX + std::to_string(rankA2);
 
     for (int i = 0; i < numToSend; i++) {
         world.send(rankA1,
@@ -123,11 +121,76 @@ TEST_CASE_METHOD(MpiBaseTestFixture,
         otherWorldThread.join();
     }
 
-    std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rank);
+    std::string expectedKey = MPI_MSG_COUNT_PREFIX + std::to_string(rank);
     REQUIRE(otherMsg.mpirank() == otherRank);
-    REQUIRE(otherMsg.intexecgraphdetails_size() == 1);
     REQUIRE(otherMsg.intexecgraphdetails().count(expectedKey) == 1);
     REQUIRE(otherMsg.intexecgraphdetails().at(expectedKey) == 1);
 }
+
+TEST_CASE_METHOD(MpiTestFixture,
+                 "Test tracing the number of MPI messages by type",
+                 "[util][exec-graph]")
+{
+    msg.set_recordexecgraph(true);
+
+    // Send one message
+    int rankA1 = 0;
+    int rankA2 = 1;
+    MPI_Status status{};
+
+    std::vector<int> messageData = { 0, 1, 2 };
+    auto buffer = new int[messageData.size()];
+
+    std::string expectedKey;
+    int msgCount;
+
+    SECTION("Normal send")
+    {
+        expectedKey = fmt::format("{}-{}-{}",
+                                  MPI_MSGTYPE_COUNT_PREFIX,
+                                  faabric::MPIMessage::NORMAL,
+                                  std::to_string(rankA2));
+        msgCount = 1;
+
+        world.send(rankA1,
+                   rankA2,
+                   BYTES(messageData.data()),
+                   MPI_INT,
+                   messageData.size());
+        world.recv(
+          rankA1, rankA2, BYTES(buffer), MPI_INT, messageData.size(), &status);
+    }
+
+    SECTION("Reduce")
+    {
+        std::vector<int> data(2, 0);
+
+        expectedKey = fmt::format("{}-{}-{}",
+                                  MPI_MSGTYPE_COUNT_PREFIX,
+                                  faabric::MPIMessage::REDUCE,
+                                  std::to_string(rankA2));
+        msgCount = worldSize - 1;
+
+        // Reduce expects to receive a message from all other ranks
+        for (int r = 0; r < worldSize; r++) {
+            if (r != rankA2) {
+                world.reduce(
+                  r, rankA2, BYTES(&data[0]), nullptr, MPI_INT, 1, MPI_SUM);
+            }
+        }
+
+        world.reduce(rankA2,
+                     rankA2,
+                     BYTES(&data[1]),
+                     BYTES(&data[1]),
+                     MPI_INT,
+                     1,
+                     MPI_SUM);
+    }
+
+    REQUIRE(msg.intexecgraphdetails_size() == NUM_MPI_EXEC_GRAPH_DETAILS);
+    REQUIRE(msg.execgraphdetails_size() == 0);
+    REQUIRE(msg.intexecgraphdetails().count(expectedKey) == 1);
+    REQUIRE(msg.intexecgraphdetails().at(expectedKey) == msgCount);
+}
 }