text: /* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  
  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
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License. */
  
  #include "paddle/fluid/distributed/service/brpc_utils.h"
  #include <limits>
  #include <memory>
  #include "paddle/fluid/platform/enforce.h"
  #include "paddle/fluid/platform/profiler.h"
  
  namespace paddle {
  namespace framework {
  class Scope;
  class Variable;
  }  // namespace framework
  namespace platform {
  class DeviceContext;
  }  // namespace platform
  }  // namespace paddle
  
  namespace paddle {
  namespace distributed {
  
  framework::proto::VarType::Type VarMessageToVarType(
      VariableMessage::Type type) {
    switch (type) {
      case VariableMessage::FP32:
        return framework::proto::VarType::FP32;  // NOLINT
      case VariableMessage::FP64:
        return framework::proto::VarType::FP64;  // NOLINT
      case VariableMessage::INT32:
        return framework::proto::VarType::INT32;  // NOLINT
      case VariableMessage::INT64:
        return framework::proto::VarType::INT64;  // NOLINT
      case VariableMessage::BOOL:
        return framework::proto::VarType::BOOL;  // NOLINT
      default:
        PADDLE_THROW(platform::errors::InvalidArgument(
            "VarMessageToVarType:Unsupported type %d", type));
    }
  }
  
  void SerializeToMultiVarMsgAndIOBuf(
      const std::string& message_name,
      const std::vector<std::string>& send_var_name_val,
      const std::vector<std::string>& recv_var_name_val,
      const platform::DeviceContext& ctx, const framework::Scope* scope,
      MultiVarMsg* request, butil::IOBuf* iobuf) {
    // 1. message_name
    request->set_message_name(message_name);
  
    // 2. var_names
    for (auto& send_var_name : send_var_name_val) {
      request->add_send_var_names(send_var_name);
    }
    for (auto& recv_var_name : recv_var_name_val) {
      request->add_recv_var_names(recv_var_name);
    }
  
    // 3. VarMessage
    for (auto& send_var_name : send_var_name_val) {
      auto* send_var_msg = request->add_var_messages();
      butil::IOBuf temp_iobuf;
      send_var_msg->set_varname(send_var_name);
  
      framework::Variable* var = scope->FindVar(send_var_name);
  
      if (var->IsType<framework::LoDTensor>()) {
        SerializeLodTensor(var, ctx, send_var_msg, &temp_iobuf);
      } else if (var->IsType<framework::SelectedRows>()) {
        SerializeSelectedRows(var, ctx, send_var_msg, &temp_iobuf);
      }
      iobuf->append(temp_iobuf);
    }
  }
  
  void SerializeLodTensor(framework::Variable* var,
                          const platform::DeviceContext& ctx, VarMsg* var_msg,
                          butil::IOBuf* iobuf) {
    auto* tensor = var->GetMutable<framework::LoDTensor>();
    var_msg->set_type(::paddle::distributed::LOD_TENSOR);
    const framework::LoD lod = tensor->lod();
    if (lod.size() > 0) {
      var_msg->set_lod_level(lod.size());
      for (auto& each : lod) {
        VarMsg::LodData* lod_inner = var_msg->add_lod();
        for (auto& d : each) {
          lod_inner->add_lod_data(d);
        }
      }
    }
    var_msg->set_data_type(static_cast<VarMsg::Type>(tensor->type()));
    for (auto& dim : framework::vectorize(tensor->dims())) {
      var_msg->add_dims(dim);
    }
    // IO Buffer
    if (platform::is_cpu_place(tensor->place())) {
      auto data_len = tensor->numel() * framework::SizeOfType(tensor->type());
      iobuf->append(reinterpret_cast<const char*>(&data_len), 8);
      iobuf->append(reinterpret_cast<const char*>(tensor->data<void>()),
                    data_len);
    } else {
  #ifdef PADDLE_WITH_CUDA
      char* temp_ptr =
          new char[tensor->numel() * framework::SizeOfType(tensor->type())];
      auto stream =
          reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
      memory::Copy(platform::CPUPlace(), temp_ptr,
                   BOOST_GET_CONST(platform::CUDAPlace, tensor->place()),
                   tensor->data<void>(),
                   tensor->numel() * framework::SizeOfType(tensor->type()),
                   stream);
      auto data_len = tensor->numel() * framework::SizeOfType(tensor->type());
      iobuf->append(reinterpret_cast<const char*>(&data_len), 8);
      iobuf->append(reinterpret_cast<const char*>(temp_ptr), data_len);
      delete[] temp_ptr;
  #endif
    }
  }
  
  void SerializeSelectedRows(framework::Variable* var,
                             const platform::DeviceContext& ctx, VarMsg* var_msg,
                             butil::IOBuf* iobuf) {
    framework::SelectedRows* slr = var->GetMutable<framework::SelectedRows>();
    auto* tensor = slr->mutable_value();
    auto* rows = slr->mutable_rows();
  
    var_msg->set_type(::paddle::distributed::SELECTED_ROWS);
    var_msg->set_slr_height(slr->height());
  
    auto* var_data = var_msg->mutable_data();
    var_data->clear();
    var_data->resize(rows->size() * sizeof(int64_t));
    char* data_ptr = const_cast<char*>(var_data->data());
  
    if (platform::is_cpu_place(tensor->place())) {
      memcpy(data_ptr, &(*rows)[0], rows->size() * sizeof(int64_t));
    } else {
  #ifdef PADDLE_WITH_CUDA
      auto stream =
          reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
      memory::Copy(platform::CPUPlace(), data_ptr,
                   BOOST_GET_CONST(platform::CUDAPlace, tensor->place()),
                   &(*rows)[0], rows->size() * sizeof(int64_t), stream);
  #endif
    }
    var_msg->set_data_type(static_cast<VarMsg::Type>(tensor->type()));
    for (auto& dim : framework::vectorize(tensor->dims())) {
      var_msg->add_dims(dim);
    }
  
    // IO Buffer
    if (platform::is_cpu_place(tensor->place())) {
      auto data_len = tensor->numel() * framework::SizeOfType(tensor->type());
      iobuf->append(reinterpret_cast<const char*>(&data_len), 8);
      iobuf->append(reinterpret_cast<const char*>(tensor->data<void>()),
                    data_len);
    } else {
  #ifdef PADDLE_WITH_CUDA
      char* temp_ptr =
          new char[tensor->numel() * framework::SizeOfType(tensor->type())];
      auto stream =
          reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
      memory::Copy(platform::CPUPlace(), temp_ptr,
                   BOOST_GET_CONST(platform::CUDAPlace, tensor->place()),
                   tensor->data<void>(),
                   tensor->numel() * framework::SizeOfType(tensor->type()),
                   stream);
      auto data_len = tensor->numel() * framework::SizeOfType(tensor->type());
      iobuf->append(reinterpret_cast<const char*>(&data_len), 8);
      iobuf->append(reinterpret_cast<const char*>(temp_ptr), data_len);
      delete[] temp_ptr;
  #endif
    }
  }
  
  void DeserializeFromMultiVarMsgAndIOBuf(const MultiVarMsg& multi_msg,
                                          const butil::IOBuf* iobuf,
                                          const platform::DeviceContext& ctx,
                                          framework::Scope* scope) {
    butil::IOBufBytesIterator io_buffer_itr(*iobuf);
    // size_t shard_buffer_remain = res_io_buffer.size();
    for (int recv_var_index = 0; recv_var_index < multi_msg.send_var_names_size();
         ++recv_var_index) {
      const auto& msg = multi_msg.var_messages(recv_var_index);
      auto* var = scope->Var(msg.varname());
      if (msg.type() == ::paddle::distributed::LOD_TENSOR) {
        DeserializeLodTensor(var, msg, io_buffer_itr, ctx);
      } else if (msg.type() == ::paddle::distributed::SELECTED_ROWS) {
        DeserializeSelectedRows(var, msg, io_buffer_itr, ctx);
      }
    }
  }
  
  void DeserializeFromMultiVarMsgAndIOBuf(const MultiVarMsg& multi_msg,
                                          const butil::IOBuf* iobuf,
                                          const platform::DeviceContext& ctx,
                                          const framework::Scope* scope) {
    butil::IOBufBytesIterator io_buffer_itr(*iobuf);
    // size_t shard_buffer_remain = res_io_buffer.size();
    for (int recv_var_index = 0; recv_var_index < multi_msg.send_var_names_size();
         ++recv_var_index) {
      const auto& msg = multi_msg.var_messages(recv_var_index);
      auto* var = scope->FindVar(msg.varname());
      PADDLE_ENFORCE_NE(var, nullptr,
                        platform::errors::InvalidArgument(
                            "Not find variable %s in scope.", msg.varname()));
      if (msg.type() == ::paddle::distributed::LOD_TENSOR) {
        DeserializeLodTensor(var, msg, io_buffer_itr, ctx);
      } else if (msg.type() == ::paddle::distributed::SELECTED_ROWS) {
        DeserializeSelectedRows(var, msg, io_buffer_itr, ctx);
      }
    }
  }
  
  void DeserializeLodTensor(framework::Variable* var, const VarMsg& msg,
                            butil::IOBufBytesIterator& io_buffer_itr,
                            const platform::DeviceContext& ctx) {
    const auto place = ctx.GetPlace();
    framework::LoDTensor* tensor = var->GetMutable<framework::LoDTensor>();
    std::vector<int> vec_dim;
    for (auto& x : msg.dims()) {
      vec_dim.push_back(x);
    }
    tensor->Resize(framework::make_ddim(vec_dim));
  
    framework::LoD lod;
    for (int i = 0; i < msg.lod_level(); ++i) {
      framework::Vector<size_t> v;
      for (int j = 0; j < msg.lod(i).lod_data_size(); ++j) {
        v.push_back(msg.lod(i).lod_data(j));
      }
      lod.push_back(v);
    }
    tensor->set_lod(lod);
  
    void* tensor_data =
        tensor->mutable_data(place, VarMessageToVarType(msg.data_type()));
  
    // IO Buffer
    if (platform::is_cpu_place(place)) {
      unsigned long data_len;
      io_buffer_itr.copy_and_forward((void*)(&data_len), 8);
      io_buffer_itr.copy_and_forward(tensor_data, data_len);
    } else if (platform::is_gpu_place(place)) {
  #ifdef PADDLE_WITH_CUDA
      unsigned long data_len;
      char* temp_ptr =
          new char[tensor->numel() * framework::SizeOfType(tensor->type())];
      io_buffer_itr.copy_and_forward((void*)(&data_len), 8);
      io_buffer_itr.copy_and_forward((void*)temp_ptr, data_len);
      auto stream =
          reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
      memory::Copy(BOOST_GET_CONST(platform::CUDAPlace, place), tensor_data,
                   platform::CPUPlace(), (void*)temp_ptr,
                   tensor->numel() * framework::SizeOfType(tensor->type()),
                   stream);
      delete[] temp_ptr;
  #endif
    }
  }
  
  void DeserializeSelectedRows(framework::Variable* var, const VarMsg& msg,
                               butil::IOBufBytesIterator& io_buffer_itr,
                               const platform::DeviceContext& ctx) {
    const auto place = ctx.GetPlace();
    auto* slr = var->GetMutable<framework::SelectedRows>();
    framework::Tensor* tensor = slr->mutable_value();
    slr->set_height(msg.slr_height());
    std::vector<int64_t> tmp_rows(msg.slr_height());
    memcpy(&tmp_rows[0], msg.data().data(), msg.slr_height() * sizeof(int64_t));
    slr->set_rows(tmp_rows);
    std::vector<int> vec_dim;
    for (auto& x : msg.dims()) {
      vec_dim.push_back(x);
    }
    tensor->Resize(framework::make_ddim(vec_dim));
    void* tensor_data =
        tensor->mutable_data(place, VarMessageToVarType(msg.data_type()));
    // IO Buffer
    if (platform::is_cpu_place(place)) {
      unsigned long data_len;
      io_buffer_itr.copy_and_forward((void*)(&data_len), 8);
      io_buffer_itr.copy_and_forward(tensor_data, data_len);
    } else if (platform::is_gpu_place(place)) {
  #ifdef PADDLE_WITH_CUDA
      char* temp_ptr =
          new char[tensor->numel() * framework::SizeOfType(tensor->type())];
      unsigned long data_len;
      io_buffer_itr.copy_and_forward((void*)(&data_len), 8);
      io_buffer_itr.copy_and_forward(temp_ptr, data_len);
      auto stream =
          reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
      memory::Copy(BOOST_GET_CONST(platform::CUDAPlace, place), tensor_data,
                   platform::CPUPlace(), temp_ptr,
                   tensor->numel() * framework::SizeOfType(tensor->type()),
                   stream);
      delete[] temp_ptr;
  #endif
    }
  }
  
  }  // namespace distributed
  }  // namespace paddle
  