send mpi and rpc framework

paddle/fluid/operators/detail/mpi_utils.cpp

@@ -42,48 +42,46 @@ void MPIUtils::InitMPI() {
                     MPI_Init(0, 0);
                     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
                     MPI_Comm_size(MPI_COMM_WORLD, &size);
-    MPI_Get_processor_name(host_name, &len)
+                    MPI_Get_processor_name(host_name, &len);
                 }
             };
 
-MPIIsend::MPIIsend(int dst, const char* req) {
-  done1 = 0;
-  done2 = 0;
-  length = strlen(req);
-  req = req;
-}
 
-MPIIsend::Send() {
-  MPI_Isend(&req, length, MPI_CHAR, dst, mpi_tag, MPI_COMM_WORLD,
-            &msg1_);
-  MPI_Test(&msg1_, &done1_, MPI_STATUS_IGNORE)
+            MPISend::MPISend(const Meta &meta) {
+                done1_ = 1;
+                done2_ = 0;
+                this->meta = meta;
             }
 
-  bool MPIIsend::IsFinished() {
-     MPI_Status status;
-     if (!done1_) MPI_Test(&msg1_, &done1_, &status);
-     return done1;
+            MPISend::Send() {
+                MPI_Send(&meta.request, meta.count, meta.datatype, meta.dst, meta.tag,
+                         MPI_COMM_WORLD);
+                done2_ = 1;
             }
 
-MPIIsend::~MPIIsend(){
-  MPI_Wait(&msg1_, MPI_STATUS_IGNORE);
-  MPI_Free_mem(req);
+            bool MPISend::IsReady() {
+                return true;
             }
 
-MPIIrecv::MPIIrecv(){
+            bool MPISend::IsFinished() { return done1_ && done2_; }
 
-}
+            MPISend::~MPISend() { MPI_Free_mem(meta); }
 
-MPIIrecv::Recv(){
 
+            MPIRecv::MPIRecv(const Meta &meta) {
+                this->meta = meta;
             }
 
-MPIIrecv::IsFinished(){
+            MPIRecv::Recv() {}
 
+            bool MPIRecv::IsReady() {
+                return true;
             }
 
-MPIIrecv::~MPIIrecv(){
+            MPIRecv::IsFinished() {}
 
+            MPIRecv::~MPIRecv() {
+                MPI_Free_mem(meta);
             }
 
         }  // namespace detail

paddle/fluid/operators/detail/mpi_utils.h

@@ -10,6 +10,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+
 #include <mpi.h>
 #include <map>
 #include <string>
@@ -21,33 +22,60 @@ namespace detail {
             class MPIUtils {
             public:
                 MPIUtils(const std::string &worker_name);
+
                 const int GetRankID(const std::string &task_id);
 
             private:
                 void InitMPI();
+
                 std::map<std::string, int> name_id_map;
             };
 
-class MPIIsend {
+            class Meta {
+            public:
+                int src;
+                int dst;
+                MPI_Datatype datatype;
+                char *request;
+                int count;
+                int tag;
+                int device;
+            };
+
+            class MPISend {
             public:
-  MPIIsend(int dst, const char* buf);
+                MPISend(const Meta &meta);
+
                 bool IsFinished();
+
+                bool IsReady();
+
                 void Send();
-  ~MPIIsend();
+
+                ~MPISend();
 
             private:
-  int done1;
-  int length;
-  char* req;
-  MPI_Request msg1_;
+                int done1_;
+                int done2_;
+                Meta *meta;
             };
 
-class MPIIrecv {
+            class MPIRecv {
             public:
-MPIIrecv();
+                MPIRecv(const Meta &meta);
+
+                bool IsReady();
+
                 bool IsFinished();
+
                 void Recv();
-  ~MPIIrecv();
+
+                ~MPIRecv();
+
+            private:
+                int done1_;
+                int done2_;
+                Meta *meta;
             };
 
         }  // namespace detail