send_recv variables

paddle/framework/lod_tensor.cc

@@ -214,9 +214,10 @@ void SerializeToStream(std::ostream &os, const LoDTensor &tensor,
   SerializeToStream(os, static_cast<Tensor>(tensor), dev_ctx);
 }
 
-void DeserializeFromStream(std::istream &is, LoDTensor *tensor) {
+void DeserializeFromStream(std::istream &is, LoDTensor *tensor,
+                           const platform::DeviceContext &dev_ctx) {
   {
-    // the 1st field, unit32_t version for SelectedRows
+    // the 1st field, unit32_t version for LoDTensor
     uint32_t version;
     is.read(reinterpret_cast<char *>(&version), sizeof(version));
     PADDLE_ENFORCE_EQ(version, 0U, "Only version 0 is supported");
@@ -237,7 +238,7 @@ void DeserializeFromStream(std::istream &is, LoDTensor *tensor) {
     }
   }
   // the 3st filed, Tensor
-  DeserializeFromStream(is, static_cast<Tensor *>(tensor));
+  DeserializeFromStream(is, static_cast<Tensor *>(tensor), dev_ctx);
 }
 
 }  // namespace framework

paddle/framework/lod_tensor.h

@@ -208,7 +208,8 @@ void AppendLoD(LoD* lod, const LoD& lod_length);
  */
 void SerializeToStream(std::ostream& os, const LoDTensor& tensor,
                        const platform::DeviceContext& dev_ctx);
-void DeserializeFromStream(std::istream& is, LoDTensor* tensor);
+void DeserializeFromStream(std::istream& is, LoDTensor* tensor,
+                           const platform::DeviceContext& dev_ctx);
 
 }  // namespace framework
 }  // namespace paddle

paddle/framework/lod_tensor_test.cc

@@ -132,7 +132,7 @@ TEST_F(LoDTensorTester, SerializeAndDeserialize) {
   std::ostringstream oss;
   SerializeToStream(oss, lod_tensor_, cpu_ctx);
   std::istringstream iss(oss.str());
-  DeserializeFromStream(iss, &dst_tensor);
+  DeserializeFromStream(iss, &dst_tensor, cpu_ctx);
   float* dst_ptr = dst_tensor.mutable_data<float>(platform::CPUPlace());
   for (int i = 0; i < kLodTensorSize; ++i) {
     EXPECT_EQ(dst_ptr[i], i);

paddle/framework/selected_rows.cc

@@ -37,8 +37,8 @@ void SerializeToStream(std::ostream& os, const SelectedRows& selected_rows,
   SerializeToStream(os, selected_rows.value(), dev_ctx);
 }
 
-void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows) {
-  auto tensor = *selected_rows->mutable_value();
+void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows,
+                           const platform::DeviceContext& dev_ctx) {
   {
     // the 1st field, unit32_t version for SelectedRows
     uint32_t version;
@@ -62,7 +62,7 @@ void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows) {
     selected_rows->set_height(height);
   }
   // the 4st field, tensor which contains the data
-  DeserializeFromStream(is, &tensor);
+  DeserializeFromStream(is, selected_rows->mutable_value(), dev_ctx);
 }
 
 }  // namespace framework

paddle/framework/selected_rows.h

@@ -66,7 +66,8 @@ class SelectedRows {
  */
 void SerializeToStream(std::ostream& os, const SelectedRows& selected_rows,
                        const platform::DeviceContext& dev_ctx);
-void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows);
+void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows,
+                           const platform::DeviceContext& dev_ctx);
 
 }  // namespace framework
 }  // namespace paddle

paddle/framework/selected_rows_test.cc

@@ -51,10 +51,12 @@ TEST_F(SelectedRowsTester, SerializeAndDeseralize) {
   SerializeToStream(oss, *selected_rows_, cpu_ctx);
 
   std::istringstream iss(oss.str());
-  DeserializeFromStream(iss, &dst_tensor);
+  DeserializeFromStream(iss, &dst_tensor, cpu_ctx);
 
   ASSERT_EQ(selected_rows_->rows(), dst_tensor.rows());
   ASSERT_EQ(selected_rows_->height(), dst_tensor.height());
+  ASSERT_EQ(selected_rows_->value().dims(), dst_tensor.value().dims());
+  ASSERT_EQ(selected_rows_->GetCompleteDims(), dst_tensor.GetCompleteDims());
 }
 
 }  // namespace framework

paddle/framework/tensor_util.h

@@ -262,7 +262,23 @@ inline void SerializeToStream(std::ostream& os, const Tensor& tensor,
   }
 }
 
-inline void DeserializeFromStream(std::istream& is, Tensor* tensor) {
+struct DeserializedDataFunctor {
+  DeserializedDataFunctor(void** buf, Tensor* tensor,
+                          const platform::Place& place)
+      : buf_(buf), tensor_(tensor), place_(place) {}
+
+  template <typename T>
+  void operator()() {
+    *buf_ = tensor_->mutable_data<T>(place_);
+  }
+
+  void** buf_;
+  Tensor* tensor_;
+  platform::Place place_;
+};
+
+inline void DeserializeFromStream(std::istream& is, Tensor* tensor,
+                                  const platform::DeviceContext& dev_ctx) {
   uint32_t version;
   is.read(reinterpret_cast<char*>(&version), sizeof(version));
   PADDLE_ENFORCE_EQ(version, 0U, "Only version 0 is supported");
@@ -281,27 +297,28 @@ inline void DeserializeFromStream(std::istream& is, Tensor* tensor) {
     dims.reserve(static_cast<size_t>(desc.dims().size()));
     std::copy(desc.dims().begin(), desc.dims().end(), std::back_inserter(dims));
     tensor->Resize(framework::make_ddim(dims));
-
     void* buf;
-    platform::Place cpu = platform::CPUPlace();
-    // TODO(Yancey1989): use VisiterDataType instead of DataType switch
-    switch (desc.data_type()) {
-      case proto::FP32:
-        buf = tensor->mutable_data<float>(cpu);
-        break;
-      case proto::FP64:
-        buf = tensor->mutable_data<double>(cpu);
-        break;
-      case proto::INT32:
-        buf = tensor->mutable_data<int>(cpu);
-        break;
-      case proto::INT64:
-        buf = tensor->mutable_data<int64_t>(cpu);
-        break;
-      default:
-        PADDLE_THROW("DataType %d not supported", desc.data_type());
+    auto ctx = platform::CPUDeviceContext();
+    if (platform::is_gpu_place(dev_ctx.GetPlace())) {
+#ifdef PADDLE_WITH_CUDA
+      Tensor cpu_tensor;
+      cpu_tensor.Resize(framework::make_ddim(dims));
+      framework::VisitDataType(
+          desc.data_type(),
+          DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), cpu_tensor.memory_size());
+      auto cpu_place = new platform::CPUPlace();
+      framework::CopyFrom(cpu_tensor, *cpu_place, dev_ctx, tensor);
+      delete cpu_place;
+#else
+      PADDLE_THROW("Unexpected branch");
+#endif
+    } else {
+      framework::VisitDataType(
+          desc.data_type(),
+          DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), tensor->memory_size());
     }
-    is.read(static_cast<char*>(buf), tensor->memory_size());
   }
 }
 

paddle/framework/tensor_util_test.cc

@@ -246,11 +246,12 @@ TEST(Tensor, SerializeAndDeserialize) {
     SerializeToStream(oss, src_tensor, cpu_ctx);
 
     std::istringstream iss(oss.str());
-    DeserializeFromStream(iss, &dst_tensor);
+    DeserializeFromStream(iss, &dst_tensor, cpu_ctx);
     int* dst_ptr = dst_tensor.mutable_data<int>(platform::CPUPlace());
     for (int i = 0; i < 5; ++i) {
       ASSERT_EQ(dst_ptr[i], array[i]);
     }
+    ASSERT_EQ(dst_tensor.dims(), src_tensor.dims());
     delete place;
   }
 #ifdef PADDLE_WITH_CUDA
@@ -268,13 +269,12 @@ TEST(Tensor, SerializeAndDeserialize) {
     SerializeToStream(oss, gpu_tensor, gpu_ctx);
 
     std::istringstream iss(oss.str());
-    DeserializeFromStream(iss, &dst_tensor);
+    DeserializeFromStream(iss, &dst_tensor, gpu_ctx);
 
     int* dst_ptr = dst_tensor.mutable_data<int>(platform::CPUPlace());
     for (int i = 0; i < 6; ++i) {
       ASSERT_EQ(dst_ptr[i], array[i]);
     }
-
     delete gpu_place;
   }
 #endif

paddle/framework/var_type.h

@@ -17,6 +17,8 @@ limitations under the License. */
 #include "paddle/framework/lod_rank_table.h"
 #include "paddle/framework/lod_tensor.h"
 #include "paddle/framework/lod_tensor_array.h"
+#include "paddle/framework/selected_rows.h"
+#include "paddle/framework/variable.h"
 
 namespace paddle {
 namespace framework {
@@ -35,7 +37,7 @@ inline proto::VarDesc::VarType ToVarType(std::type_index type) {
 }
 
 template <typename Visitor>
-inline void VisitVarType(const Variable& var, Visitor visitor) {
+inline void VisitVarType(const framework::Variable& var, Visitor visitor) {
   switch (ToVarType(var.Type())) {
     case proto::VarDesc_VarType_LOD_TENSOR:
       visitor(var.Get<framework::LoDTensor>());

paddle/operators/detail/recv_impl.cc

@@ -21,14 +21,9 @@ namespace detail {
 Status SendRecvServerImpl::SendVariable(ServerContext *context,
                                         const VariableMessage *in_var,
                                         VoidMessage *out_var) {
-  // TODO(typhoonzero): support different variable types.
-  std::istringstream iss(in_var->serialized());
-  framework::LoDTensor t;
-  framework::DeserializeFromStream(iss, &t);
-  TensorWithName tensor_with_name =
-      std::make_pair(in_var->varname(), std::move(t));
-
-  var_recv_queue_.Push(std::move(tensor_with_name));
+  MessageWithName msg_with_name =
+      std::make_pair(in_var->varname(), std::move(*in_var));
+  var_recv_queue_.Push(std::move(msg_with_name));
   return Status::OK;
 }
 
@@ -37,14 +32,8 @@ Status SendRecvServerImpl::GetVariable(ServerContext *context,
                                        VariableMessage *out_var) {
   std::string get_var_name = in_var->varname();
   auto *var = scope_->FindVar(get_var_name);
-  auto tensor = var->Get<framework::LoDTensor>();
-  std::ostringstream oss;
-  framework::SerializeToStream(oss, tensor, platform::CPUDeviceContext());
 
-  std::string *varname = out_var->mutable_varname();
-  *varname = get_var_name;
-  std::string *serialized = out_var->mutable_serialized();
-  *serialized = oss.str();
+  SerializeToMessage(get_var_name, var, platform::CPUDeviceContext(), out_var);
   return Status::OK;
 }
 

paddle/operators/detail/send_impl.cc

@@ -27,14 +27,8 @@ bool RPCClient::SendVariable(const framework::Scope& scope,
   auto ctx = platform::CPUDeviceContext();
   auto* var = scope.FindVar(inname);
   PADDLE_ENFORCE(var);
-  // TODO(typhoonzero): support SelectedRows
-  PADDLE_ENFORCE(var->IsType<framework::LoDTensor>(),
-                 "Only support LoDTensor, %s has wrong type", inname);
-  const framework::LoDTensor& tensor = var->Get<framework::LoDTensor>();
-  std::ostringstream oss;
-  framework::SerializeToStream(oss, tensor, ctx);
-  msg.set_varname(inname);
-  msg.set_serialized(oss.str());
+  SerializeToMessage(inname, var, ctx, &msg);
+
   Status status = stub_->SendVariable(&context, msg, &out_msg);
   if (!status.ok()) {
     LOG(ERROR) << "gRPC error: " << status.error_message();
@@ -50,19 +44,15 @@ bool RPCClient::GetVariable(const framework::Scope& scope,
   call_msg.set_varname(outname);
   auto ctx = platform::CPUDeviceContext();
   Status status = stub_->GetVariable(&context, call_msg, &ret_msg);
+  auto* outvar = scope.FindVar(outname);
   if (!status.ok()) {
     LOG(ERROR) << "gRPC error: " << status.error_message();
     return false;
   }
 
   std::istringstream iss(ret_msg.serialized());
+  DeserializeFromMessage(ret_msg, ctx, outvar);
 
-  framework::LoDTensor ret_tensor;
-  framework::DeserializeFromStream(iss, &ret_tensor);
-  auto* outvar = scope.FindVar(outname);
-  framework::LoDTensor* out_tensor = outvar->GetMutable<framework::LoDTensor>();
-  // FIXME(typhoonzero): do not copy.
-  framework::CopyFrom(ret_tensor, ctx.GetPlace(), ctx, out_tensor);
   return true;
 }
 

paddle/operators/detail/send_recv.proto

@@ -1,7 +1,6 @@
-/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve. Licensed under
+the Apache License, Version 2.0 (the "License"); you may not use this file
+except in compliance with the License.
 You may obtain a copy of the License at
 
     http://www.apache.org/licenses/LICENSE-2.0
@@ -13,7 +12,6 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 syntax = "proto3";
-
 package sendrecv;
 
 service SendRecvService {
@@ -29,12 +27,18 @@ service SendRecvService {
 
 // VariableMessage is serialized paddle variable message.
 // It can be:
-// Tensor
 // LoDTensor
 // SelectedRows
+enum VarType {
+  LOD_TENSOR = 0;
+  SELECTED_ROWS = 1;
+}
+
 message VariableMessage {
   string varname = 1;
-  bytes serialized = 2;
+  // TODO(Yancey1989): reference framework::proto::VarDesc::VarType
+  VarType type = 2;
+  bytes serialized = 3;
 }
 
 message VoidMessage {}