update send_op

paddle/fluid/operators/distributed_ops/send_op.cc

@@ -20,6 +20,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/distributed/distributed.h"
+#include "paddle/fluid/operators/distributed/parameter_send.h"
 #include "paddle/fluid/operators/distributed_ops/send_recv_util.h"
 #include "paddle/fluid/platform/profiler.h"
 
@@ -37,30 +38,46 @@ class SendOp : public framework::OperatorBase {
                const platform::Place& place) const override {
     auto ins = Inputs("X");
 
-    std::vector<std::string> epmap = Attr<std::vector<std::string>>("epmap");
+    auto epmap = Attr<std::vector<std::string>>("epmap");
     int sync_send = Attr<int>("sync_mode");
 
-    platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
-    auto& ctx = *pool.Get(place);
-
-    distributed::RPCClient* rpc_client =
-        distributed::RPCClient::GetInstance<RPCCLIENT_T>(
-            Attr<int>("trainer_id"));
-
-    std::vector<distributed::VarHandlePtr> rets;
-    for (size_t i = 0; i < ins.size(); i++) {
-      if (NeedSend(scope, ins[i])) {
-        VLOG(3) << "sending " << ins[i] << " to " << epmap[i];
-        rets.push_back(rpc_client->AsyncSendVar(epmap[i], ctx, scope, ins[i]));
-      } else {
-        VLOG(3) << "don't send no-initialied variable: " << ins[i];
+    auto send_varnames = Attr<std::vector<std::string>>("send_varnames");
+    auto height_sections = Attr<std::vector<int64_t>>("height_sections");
+
+    if (send_varnames.size() > 0) {
+      PADDLE_ENFORCE_EQ(ins.size(), 1, "");
+      framework::RuntimeContext ctx(Inputs(), Outputs(), scope);
+      platform::DeviceContextPool& pool =
+          platform::DeviceContextPool::Instance();
+      auto* dev_ctx = pool.Get(place);
+      auto exe_ctx = framework::ExecutionContext(*this, scope, *dev_ctx, ctx);
+      distributed::send<float>(ins[0], send_varnames, epmap, height_sections,
+                               exe_ctx, scope, static_cast<bool>(sync_send));
+    } else {
+      platform::DeviceContextPool& pool =
+          platform::DeviceContextPool::Instance();
+      auto& ctx = *pool.Get(place);
+
+      distributed::RPCClient* rpc_client =
+          distributed::RPCClient::GetInstance<RPCCLIENT_T>(
+              Attr<int>("trainer_id"));
+
+      std::vector<distributed::VarHandlePtr> rets;
+      for (size_t i = 0; i < ins.size(); i++) {
+        if (NeedSend(scope, ins[i])) {
+          VLOG(3) << "sending " << ins[i] << " to " << epmap[i];
+          rets.push_back(
+              rpc_client->AsyncSendVar(epmap[i], ctx, scope, ins[i]));
+        } else {
+          VLOG(3) << "don't send no-initialied variable: " << ins[i];
+        }
       }
-    }
-    if (sync_send) {
-      for (size_t i = 0; i < rets.size(); i++) {
-        VLOG(7) << "before sync_send " << ins[i] << "from " << epmap[i];
-        PADDLE_ENFORCE(rets[i]->Wait(), "internal error in RPCClient");
-        VLOG(7) << "after sync_send " << ins[i] << "from " << epmap[i];
+      if (sync_send) {
+        for (size_t i = 0; i < rets.size(); i++) {
+          VLOG(7) << "before sync_send " << ins[i] << "from " << epmap[i];
+          PADDLE_ENFORCE(rets[i]->Wait(), "internal error in RPCClient");
+          VLOG(7) << "after sync_send " << ins[i] << "from " << epmap[i];
+        }
       }
     }
   }