BP5: Rerouting improvements

source/adios2/helper/adiosComm.cpp

@@ -180,6 +180,20 @@ Comm::Status Comm::Req::Wait(const std::string &hint)
     return status;
 }
 
+bool Comm::Req::TestComplete(const std::string &hint)
+{
+    bool result = false;
+    if (m_Impl)
+    {
+        result = m_Impl->TestComplete(hint);
+        if (result)
+        {
+            m_Impl.reset();
+        }
+    }
+    return result;
+}
+
 Comm::Win::Win() = default;
 
 Comm::Win::Win(std::unique_ptr<CommWinImpl> impl) : m_Impl(std::move(impl)) {}

source/adios2/helper/adiosComm.h

@@ -331,6 +331,11 @@ class Comm::Req
      */
     Comm::Status Wait(const std::string &hint = std::string());
 
+    /**
+     * @brief Test for completion of a request.
+     */
+    bool TestComplete(const std::string &hint = std::string());
+
 private:
     friend class CommImpl;
 
@@ -523,6 +528,7 @@ class CommReqImpl
 public:
     virtual ~CommReqImpl() = 0;
     virtual Comm::Status Wait(const std::string &hint) = 0;
+    virtual bool TestComplete(const std::string &hint) = 0;
 };
 
 class CommWinImpl

source/adios2/helper/adiosCommDummy.cpp

@@ -35,6 +35,7 @@ class CommReqImplDummy : public CommReqImpl
     ~CommReqImplDummy() override;
 
     Comm::Status Wait(const std::string &hint) override;
+    bool TestComplete(const std::string &hint) override;
 };
 
 CommReqImplDummy::~CommReqImplDummy() = default;
@@ -334,6 +335,8 @@ Comm::Status CommReqImplDummy::Wait(const std::string &hint)
     return status;
 }
 
+bool CommReqImplDummy::TestComplete(const std::string &hint) { return true; }
+
 int CommWinImplDummy::Free(const std::string &hint) { return 0; }
 
 Comm CommDummy()

source/adios2/helper/adiosCommMPI.cpp

@@ -112,6 +112,7 @@ class CommReqImplMPI : public CommReqImpl
     ~CommReqImplMPI() override;
 
     Comm::Status Wait(const std::string &hint) override;
+    bool TestComplete(const std::string &hint) override;
 
     /** Encapsulated MPI datatype of the requested operation.  */
     MPI_Datatype m_MPIDatatype = MPI_DATATYPE_NULL;
@@ -658,6 +659,56 @@ Comm::Status CommReqImplMPI::Wait(const std::string &hint)
     return status;
 }
 
+bool CommReqImplMPI::TestComplete(const std::string &hint)
+{
+    if (m_MPIReqs.empty())
+    {
+        return true;
+    }
+
+    int flag = 0;
+    int mpiReturn = MPI_SUCCESS;
+    std::vector<MPI_Status> mpiStatuses(m_MPIReqs.size());
+
+    if (m_MPIReqs.size() > 1)
+    {
+        mpiReturn = MPI_Testall(static_cast<int>(m_MPIReqs.size()), m_MPIReqs.data(), &flag,
+                                mpiStatuses.data());
+
+        // Only check for errors in statuses if the operation actually completed
+        if (flag != 0 && mpiReturn == MPI_ERR_IN_STATUS)
+        {
+            for (auto &mpiStatus : mpiStatuses)
+            {
+                if (mpiStatus.MPI_ERROR != MPI_SUCCESS)
+                {
+                    mpiReturn = mpiStatus.MPI_ERROR;
+                    break;
+                }
+            }
+        }
+    }
+    else
+    {
+        mpiReturn = MPI_Test(m_MPIReqs.data(), &flag, mpiStatuses.data());
+    }
+
+    // non-zero means the operations are complete
+    if (flag != 0)
+    {
+        // Throw or handle any errors we collected
+        CheckMPIReturn(mpiReturn, hint);
+
+        // Since the requests are done, clear the internal container so the object
+        // registers as complete/empty for future calls.
+        m_MPIReqs.clear();
+        return true;
+    }
+
+    // Request(s) not complete, the caller can try again later
+    return false;
+}
+
 int CommWinImplMPI::Free(const std::string &hint) { return MPI_Win_free(&m_Win); }
 
 Comm CommWithMPI(MPI_Comm mpiComm)

source/adios2/helper/adiosRerouting.cpp

@@ -23,7 +23,8 @@ namespace adios2
 namespace helper
 {
 
-void RerouteMessage::NonBlockingSendTo(helper::Comm &comm, int destRank, std::vector<char> &buffer)
+helper::Comm::Req RerouteMessage::NonBlockingSendTo(helper::Comm &comm, int destRank,
+                                                    std::vector<char> &buffer)
 {
     size_t pos = 0;
     buffer.resize(REROUTE_MESSAGE_SIZE);
@@ -45,7 +46,7 @@ void RerouteMessage::NonBlockingSendTo(helper::Comm &comm, int destRank, std::ve
 
     // std::cout << ss.str();
 
-    comm.Isend(buffer.data(), buffer.size(), destRank, 0);
+    return comm.Isend(buffer.data(), buffer.size(), destRank, 0);
 }
 
 void RerouteMessage::BlockingSendTo(helper::Comm &comm, int destRank, std::vector<char> &buffer)

source/adios2/helper/adiosRerouting.h

@@ -63,7 +63,8 @@ class RerouteMessage
     }
 
     // Send the contents of this message to another rank
-    void NonBlockingSendTo(helper::Comm &comm, int destRank, std::vector<char> &buffer);
+    helper::Comm::Req NonBlockingSendTo(helper::Comm &comm, int destRank,
+                                        std::vector<char> &buffer);
     void BlockingSendTo(helper::Comm &comm, int destRank, std::vector<char> &buffer);
 
     // Receive a message from another rank to populate this message

testing/adios2/helper/TestRerouteMessage.cpp

@@ -35,7 +35,7 @@ void SendAndReceiveMessage(helper::Comm &comm, int destRank, int srcRank)
     origMsg.m_DestRank = destRank;
     origMsg.m_Offset = 2138;
     origMsg.m_Size = 1213;
-    origMsg.NonBlockingSendTo(comm, destRank, sendBuffer);
+    adios2::helper::Comm::Req sendReq = origMsg.NonBlockingSendTo(comm, destRank, sendBuffer);
 
     int ready = 0;
     helper::Comm::Status status;
@@ -65,6 +65,9 @@ void SendAndReceiveMessage(helper::Comm &comm, int destRank, int srcRank)
     ASSERT_EQ(receivedMsg.m_DestRank, worldRank);
     ASSERT_EQ(receivedMsg.m_Offset, origMsg.m_Offset);
     ASSERT_EQ(receivedMsg.m_Size, origMsg.m_Size);
+
+    // Make sure the request is complete
+    ASSERT_EQ(sendReq.TestComplete(), true);
 }
 }