Overlap comm and compute (#48)

* adding overlap-comm-and-compute pass
* implementing async comm support

Author: fschlimb

imex_version.txt

@@ -1 +1 @@
-e5d49d469ec53f62d5befc7c30f4f404a7870031
+ff9d13d5e71213a99a9a436895072b853236c040

src/Deferred.cpp

@@ -13,8 +13,6 @@
 #include "include/ddptensor/Service.hpp"
 #include "include/ddptensor/Transceiver.hpp"
 
-#include <imex/Dialect/Dist/IR/DistOps.h>
-#include <imex/Dialect/DistRuntime/IR/DistRuntimeOps.h>
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/LLVMIR/LLVMDialect.h>
@@ -95,8 +93,6 @@ void process_promises() {
 
     // Create a MLIR module
     auto module = builder.create<::mlir::ModuleOp>(loc);
-    auto protos = builder.create<::imex::distruntime::RuntimePrototypesOp>(loc);
-    module.push_back(protos);
     // Create the jit func
     // create dummy type, we'll replace it with the actual type later
     auto dummyFuncType = builder.getFunctionType({}, {});

src/MPITransceiver.cpp

@@ -187,13 +187,17 @@ void MPITransceiver::reduce_all(void *inout, DTypeId T, size_t N,
   MPI_Allreduce(MPI_IN_PLACE, inout, N, to_mpi(T), to_mpi(op), _comm);
 }
 
-void MPITransceiver::alltoall(const void *buffer_send, const int *counts_send,
-                              const int *displacements_send, DTypeId datatype,
-                              void *buffer_recv, const int *counts_recv,
-                              const int *displacements_recv) {
-  MPI_Alltoallv(buffer_send, counts_send, displacements_send, to_mpi(datatype),
-                buffer_recv, counts_recv, displacements_recv, to_mpi(datatype),
-                _comm);
+Transceiver::WaitHandle
+MPITransceiver::alltoall(const void *buffer_send, const int *counts_send,
+                         const int *displacements_send, DTypeId datatype,
+                         void *buffer_recv, const int *counts_recv,
+                         const int *displacements_recv) {
+  MPI_Request request;
+  MPI_Ialltoallv(buffer_send, counts_send, displacements_send, to_mpi(datatype),
+                 buffer_recv, counts_recv, displacements_recv, to_mpi(datatype),
+                 _comm, &request);
+  static_assert(sizeof(request == sizeof(WaitHandle)));
+  return static_cast<WaitHandle>(request);
 }
 
 void MPITransceiver::alltoall(const void *buffer_send, const int counts,
@@ -226,3 +230,10 @@ void MPITransceiver::send_recv(void *buffer_send, int count_send,
   MPI_Sendrecv_replace(buffer_send, count_send, to_mpi(datatype_send), dest,
                        SRTAG, source, SRTAG, _comm, MPI_STATUS_IGNORE);
 }
+
+void MPITransceiver::wait(WaitHandle h) {
+  if (h) {
+    auto r = static_cast<MPI_Request>(h);
+    MPI_Wait(&r, MPI_STATUS_IGNORE);
+  }
+}

src/idtr.cpp

@@ -30,7 +30,33 @@ template <typename T> T *mr_to_ptr(void *ptr, intptr_t offset) {
   return reinterpret_cast<T *>(ptr) + offset; // &mr.aligned[mr.offset]
 }
 
+// abstract handle providing an abstract wait method
+struct WaitHandleBase {
+  virtual ~WaitHandleBase(){};
+  virtual void wait() = 0;
+};
+
+// concrete handle to be instantiated with a lambda or alike
+// the lambda will be executed within wait()
+template <typename T> class WaitHandle : public WaitHandleBase {
+  T _fini;
+
+public:
+  WaitHandle(T fini) : _fini(fini) {}
+  virtual void wait() override { _fini(); }
+};
+
+template <typename T> WaitHandle<T> *mkWaitHandle(T fini) {
+  return new WaitHandle<T>(fini);
+};
+
 extern "C" {
+void _idtr_wait(WaitHandleBase *handle) {
+  if (handle) {
+    handle->wait();
+    delete handle;
+  }
+}
 
 #define NO_TRANSCEIVER
 #ifdef NO_TRANSCEIVER
@@ -401,8 +427,9 @@ void _idtr_reshape(DTypeId ddpttype, int64_t lRank, int64_t *gShapePtr,
   Buffer outbuff(totSSz * sizeof_dtype(ddpttype), 2); // FIXME debug value
   bufferizeN(lDataPtr, ddpttype, lShapePtr, lStridesPtr, lsOffs.data(),
              lsEnds.data(), lRank, N, outbuff.data());
-  tc->alltoall(outbuff.data(), sszs.data(), soffs.data(), ddpttype, oDataPtr,
-               rszs.data(), roffs.data());
+  auto hdl = tc->alltoall(outbuff.data(), sszs.data(), soffs.data(), ddpttype,
+                          oDataPtr, rszs.data(), roffs.data());
+  tc->wait(hdl);
 }
 
 /// @brief reshape tensor
@@ -451,21 +478,22 @@ TYPED_RESHAPE(i1, bool);
 /// @brief Update data in halo parts
 /// We assume tensor is partitioned along the first dimension only
 /// (row partitioning) and partitions are ordered by ranks
-void _idtr_update_halo(DTypeId ddpttype, int64_t ndims, int64_t *ownedOff,
-                       int64_t *ownedShape, int64_t *ownedStride,
-                       int64_t *bbOff, int64_t *bbShape, void *ownedData,
-                       int64_t *leftHaloShape, int64_t *leftHaloStride,
-                       void *leftHaloData, int64_t *rightHaloShape,
-                       int64_t *rightHaloStride, void *rightHaloData,
-                       Transceiver *tc) {
+/// @return (MPI) handles
+void *_idtr_update_halo(DTypeId ddpttype, int64_t ndims, int64_t *ownedOff,
+                        int64_t *ownedShape, int64_t *ownedStride,
+                        int64_t *bbOff, int64_t *bbShape, void *ownedData,
+                        int64_t *leftHaloShape, int64_t *leftHaloStride,
+                        void *leftHaloData, int64_t *rightHaloShape,
+                        int64_t *rightHaloStride, void *rightHaloData,
+                        Transceiver *tc) {
 
 #ifdef NO_TRANSCEIVER
   initMPIRuntime();
   tc = getTransceiver();
 #endif
   auto nworkers = tc->nranks();
   if (nworkers <= 1 || getenv("DDPT_SKIP_COMM"))
-    return;
+    return nullptr;
   auto myWorkerIndex = tc->rank();
 
   // Gather table with bounding box offsets and shapes for all workers
@@ -582,7 +610,6 @@ void _idtr_update_halo(DTypeId ddpttype, int64_t ndims, int64_t *ownedOff,
   void *sendData;
   bool bufferizeLRecv = !is_contiguous(leftHaloShape, leftHaloStride, ndims);
   bool bufferizeRRecv = !is_contiguous(rightHaloShape, rightHaloStride, ndims);
-  std::vector<int64_t> recvBufferStart(nworkers * ndims, 0);
   std::vector<int64_t> lRecvBufferSize(nworkers * ndims, 0);
   std::vector<int64_t> rRecvBufferSize(nworkers * ndims, 0);
 
@@ -624,35 +651,50 @@ void _idtr_update_halo(DTypeId ddpttype, int64_t ndims, int64_t *ownedOff,
     bufferize(ownedData, ddpttype, ownedShape, ownedStride, lBufferStart.data(),
               lBufferSize.data(), ndims, nworkers, sendBuff.data());
   }
-  tc->alltoall(sendData, lSendSize.data(), lSendOff.data(), ddpttype, lRecvData,
-               lRecvSize.data(), lRecvOff.data());
-  if (bufferizeLRecv) {
-    unpack(lRecvData, ddpttype, leftHaloShape, leftHaloStride,
-           recvBufferStart.data(), lRecvBufferSize.data(), ndims, nworkers,
-           leftHaloData);
-  }
+  auto lwh = tc->alltoall(sendData, lSendSize.data(), lSendOff.data(), ddpttype,
+                          lRecvData, lRecvSize.data(), lRecvOff.data());
   if (bufferizeSend) {
     bufferize(ownedData, ddpttype, ownedShape, ownedStride, rBufferStart.data(),
               rBufferSize.data(), ndims, nworkers, sendBuff.data());
   }
-  tc->alltoall(sendData, rSendSize.data(), rSendOff.data(), ddpttype, rRecvData,
-               rRecvSize.data(), rRecvOff.data());
-  if (bufferizeRRecv) {
-    unpack(rRecvData, ddpttype, rightHaloShape, rightHaloStride,
-           recvBufferStart.data(), rRecvBufferSize.data(), ndims, nworkers,
-           rightHaloData);
+  auto rwh = tc->alltoall(sendData, rSendSize.data(), rSendOff.data(), ddpttype,
+                          rRecvData, rRecvSize.data(), rRecvOff.data());
+
+  auto wait = [=, _lRecvBufferSize = std::move(lRecvBufferSize),
+               _rRecvBufferSize = std::move(rRecvBufferSize)]() {
+    tc->wait(lwh);
+    std::vector<int64_t> recvBufferStart(nworkers * ndims, 0);
+    if (bufferizeLRecv) {
+      unpack(lRecvData, ddpttype, leftHaloShape, leftHaloStride,
+             recvBufferStart.data(), _lRecvBufferSize.data(), ndims, nworkers,
+             leftHaloData);
+    }
+    tc->wait(rwh);
+    if (bufferizeRRecv) {
+      unpack(rRecvData, ddpttype, rightHaloShape, rightHaloStride,
+             recvBufferStart.data(), _rRecvBufferSize.data(), ndims, nworkers,
+             rightHaloData);
+    }
+  };
+  assert(lRecvBufferSize.empty() && rRecvBufferSize.empty());
+
+  // FIXME (in imex) buffer-dealloc pass deallocs halo strides and sizes
+  if (bufferizeLRecv || bufferizeRRecv || getenv("DDPT_NO_ASYNC")) {
+    wait();
+    return nullptr;
   }
+  return mkWaitHandle(wait);
 }
 
 /// @brief templated wrapper for typed function versions calling
 /// _idtr_update_halo
 template <typename T>
-void _idtr_update_halo(Transceiver *tc, int64_t gShapeRank, void *gShapeDescr,
-                       int64_t oOffRank, void *oOffDescr, int64_t oDataRank,
-                       void *oDataDescr, int64_t bbOffRank, void *bbOffDescr,
-                       int64_t bbShapeRank, void *bbShapeDescr,
-                       int64_t lHaloRank, void *lHaloDescr, int64_t rHaloRank,
-                       void *rHaloDescr) {
+void *_idtr_update_halo(Transceiver *tc, int64_t gShapeRank, void *gShapeDescr,
+                        int64_t oOffRank, void *oOffDescr, int64_t oDataRank,
+                        void *oDataDescr, int64_t bbOffRank, void *bbOffDescr,
+                        int64_t bbShapeRank, void *bbShapeDescr,
+                        int64_t lHaloRank, void *lHaloDescr, int64_t rHaloRank,
+                        void *rHaloDescr) {
 
   auto ddpttype = DTYPE<T>::value;
 
@@ -664,25 +706,25 @@ void _idtr_update_halo(Transceiver *tc, int64_t gShapeRank, void *gShapeDescr,
   UnrankedMemRefType<T> leftHalo(lHaloRank, lHaloDescr);
   UnrankedMemRefType<T> rightHalo(rHaloRank, rHaloDescr);
 
-  _idtr_update_halo(ddpttype, ownedData.rank(), ownedOff.data(),
-                    ownedData.sizes(), ownedData.strides(), bbOff.data(),
-                    bbShape.data(), ownedData.data(), leftHalo.sizes(),
-                    leftHalo.strides(), leftHalo.data(), rightHalo.sizes(),
-                    rightHalo.strides(), rightHalo.data(), tc);
+  return _idtr_update_halo(
+      ddpttype, ownedData.rank(), ownedOff.data(), ownedData.sizes(),
+      ownedData.strides(), bbOff.data(), bbShape.data(), ownedData.data(),
+      leftHalo.sizes(), leftHalo.strides(), leftHalo.data(), rightHalo.sizes(),
+      rightHalo.strides(), rightHalo.data(), tc);
 }
 
 extern "C" {
 #define TYPED_UPDATE_HALO(_sfx, _typ)                                          \
-  void _idtr_update_halo_##_sfx(                                               \
+  void *_idtr_update_halo_##_sfx(                                              \
       Transceiver *tc, int64_t gShapeRank, void *gShapeDescr,                  \
       int64_t oOffRank, void *oOffDescr, int64_t oDataRank, void *oDataDescr,  \
       int64_t bbOffRank, void *bbOffDescr, int64_t bbShapeRank,                \
       void *bbShapeDescr, int64_t lHaloRank, void *lHaloDescr,                 \
       int64_t rHaloRank, void *rHaloDescr) {                                   \
-    _idtr_update_halo<_typ>(tc, gShapeRank, gShapeDescr, oOffRank, oOffDescr,  \
-                            oDataRank, oDataDescr, bbOffRank, bbOffDescr,      \
-                            bbShapeRank, bbShapeDescr, lHaloRank, lHaloDescr,  \
-                            rHaloRank, rHaloDescr);                            \
+    return _idtr_update_halo<_typ>(                                            \
+        tc, gShapeRank, gShapeDescr, oOffRank, oOffDescr, oDataRank,           \
+        oDataDescr, bbOffRank, bbOffDescr, bbShapeRank, bbShapeDescr,          \
+        lHaloRank, lHaloDescr, rHaloRank, rHaloDescr);                         \
   }
 
 TYPED_UPDATE_HALO(f64, double);

src/include/ddptensor/MPITransceiver.hpp

@@ -27,16 +27,18 @@ class MPITransceiver : public Transceiver {
   virtual void barrier();
   virtual void bcast(void *ptr, size_t N, rank_type root);
   virtual void reduce_all(void *inout, DTypeId T, size_t N, RedOpType op);
-  virtual void alltoall(const void *buffer_send, const int *counts_send,
-                        const int *displacements_send, DTypeId datatype_send,
-                        void *buffer_recv, const int *counts_recv,
-                        const int *displacements_recv);
+  virtual WaitHandle alltoall(const void *buffer_send, const int *counts_send,
+                              const int *displacements_send,
+                              DTypeId datatype_send, void *buffer_recv,
+                              const int *counts_recv,
+                              const int *displacements_recv);
   virtual void alltoall(const void *buffer_send, const int counts,
                         DTypeId datatype, void *buffer_recv);
   virtual void gather(void *buffer, const int *counts, const int *displacements,
                       DTypeId datatype, rank_type root);
   virtual void send_recv(void *buffer_send, int count_send,
                          DTypeId datatype_send, int dest, int source);
+  virtual void wait(WaitHandle);
 
 private:
   rank_type _nranks, _rank;

src/include/ddptensor/Transceiver.hpp

@@ -10,6 +10,8 @@
 
 class Transceiver {
 public:
+  using WaitHandle = uint32_t;
+
   virtual ~Transceiver(){};
 
   virtual bool is_cw() = 0;
@@ -38,10 +40,11 @@ class Transceiver {
   virtual void reduce_all(void *inout, DTypeId T, size_t N, RedOpType op) = 0;
 
   // umm, can this be higher-level?
-  virtual void alltoall(const void *buffer_send, const int *counts_send,
-                        const int *displacements_send, DTypeId datatype_send,
-                        void *buffer_recv, const int *counts_recv,
-                        const int *displacements_recv) = 0;
+  virtual WaitHandle alltoall(const void *buffer_send, const int *counts_send,
+                              const int *displacements_send,
+                              DTypeId datatype_send, void *buffer_recv,
+                              const int *counts_recv,
+                              const int *displacements_recv) = 0;
   virtual void alltoall(const void *buffer_send, const int counts,
                         DTypeId datatype, void *buffer_recv) = 0;
 
@@ -50,6 +53,7 @@ class Transceiver {
 
   virtual void send_recv(void *buffer_send, int count_send,
                          DTypeId datatype_send, int dest, int source) = 0;
+  virtual void wait(WaitHandle) = 0;
 };
 
 extern void init_transceiver(Transceiver *);

src/jit/mlir.cpp

@@ -435,6 +435,7 @@ static const char *pass_pipeline =
                             "func.func(dist-infer-elementwise-cores),"
                             "convert-dist-to-standard,"
                             "canonicalize,"
+                            "overlap-comm-and-compute,"
                             "lower-distruntime-to-idtr,"
                             "convert-ptensor-to-linalg,"
                             "canonicalize,"