send mpi and rpc framework

paddle/fluid/operators/detail/mpi_utils.cpp

@@ -12,9 +12,9 @@
 #define mpi_tag = 2008
 
 namespace paddle {
-namespace operators {
-namespace detail {
-MPIUtils::MPIUtils(const std::string& worker_name) {
+    namespace operators {
+        namespace detail {
+            MPIUtils::MPIUtils(const std::string &worker_name) {
                 InitMPI();
 
                 int rank = 0, size = 1;
@@ -29,9 +29,9 @@ MPIUtils::MPIUtils(const std::string& worker_name) {
                 for (int i = 0; i < number_of_procs; i++) {
                     name_to_id_[std::string(&worker_names[i * 128])] = i;
                 }
-}
+            }
 
-void MPIUtils::InitMPI() {
+            void MPIUtils::InitMPI() {
                 int flag = 0;
                 MPI_CHECK(MPI_Initialized(&flag));
 
@@ -42,51 +42,49 @@ 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)
-}
-
-  bool MPIIsend::IsFinished() {
-     MPI_Status status;
-     if (!done1_) MPI_Test(&msg1_, &done1_, &status);
-     return done1;
+            };
+
+
+            MPISend::MPISend(const Meta &meta) {
+                done1_ = 1;
+                done2_ = 0;
+                this->meta = meta;
             }
 
-MPIIsend::~MPIIsend(){
-  MPI_Wait(&msg1_, MPI_STATUS_IGNORE);
-  MPI_Free_mem(req);
-}
+            MPISend::Send() {
+                MPI_Send(&meta.request, meta.count, meta.datatype, meta.dst, meta.tag,
+                         MPI_COMM_WORLD);
+                done2_ = 1;
+            }
 
-MPIIrecv::MPIIrecv(){
+            bool MPISend::IsReady() {
+                return true;
+            }
 
-}
+            bool MPISend::IsFinished() { return done1_ && done2_; }
 
-MPIIrecv::Recv(){
+            MPISend::~MPISend() { MPI_Free_mem(meta); }
 
-}
 
-MPIIrecv::IsFinished(){
+            MPIRecv::MPIRecv(const Meta &meta) {
+                this->meta = meta;
+            }
+
+            MPIRecv::Recv() {}
 
-}
+            bool MPIRecv::IsReady() {
+                return true;
+            }
 
-MPIIrecv::~MPIIrecv(){
+            MPIRecv::IsFinished() {}
 
-}
+            MPIRecv::~MPIRecv() {
+                MPI_Free_mem(meta);
+            }
 
-}  // namespace detail
+        }  // namespace detail
 
-}  // namespace operators
+    }  // namespace operators
 }  // namespace paddle
\ No newline at end of file

paddle/fluid/operators/detail/mpi_utils.h

@@ -10,46 +10,74 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+
 #include <mpi.h>
 #include <map>
 #include <string>
 #include <vector>
 
 namespace paddle {
-namespace operators {
-namespace detail {
-class MPIUtils {
+    namespace operators {
+        namespace detail {
+            class MPIUtils {
             public:
-  MPIUtils(const std::string& worker_name);
-  const int GetRankID(const std::string& task_id);
+                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 Meta {
+            public:
+                int src;
+                int dst;
+                MPI_Datatype datatype;
+                char *request;
+                int count;
+                int tag;
+                int device;
+            };
 
-class MPIIsend {
+            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();
-bool IsFinished();
+                MPIRecv(const Meta &meta);
+
+                bool IsReady();
+
+                bool IsFinished();
+
                 void Recv();
-  ~MPIIrecv();
-};
 
-}  // namespace detail
-}  // namespace operators
+                ~MPIRecv();
+
+            private:
+                int done1_;
+                int done2_;
+                Meta *meta;
+            };
+
+        }  // namespace detail
+    }  // namespace operators
 }  // namespace paddle