add BCast and Gather

paddle/fluid/framework/details/CMakeLists.txt

@@ -2,15 +2,21 @@ cc_library(var_handle SRCS var_handle.cc DEPS place)
 cc_library(op_handle_base SRCS op_handle_base.cc DEPS var_handle device_context)
 cc_library(scale_loss_grad_op_handle SRCS scale_loss_grad_op_handle.cc DEPS op_handle_base scope lod_tensor ddim memory)
 cc_library(fetch_op_handle SRCS fetch_op_handle.cc DEPS op_handle_base scope lod_tensor ddim memory)
-nv_library(nccl_all_reduce_op_handle SRCS nccl_all_reduce_op_handle.cc DEPS op_handle_base scope lod_tensor ddim memory
+if(WITH_GPU)
+    nv_library(nccl_all_reduce_op_handle SRCS nccl_all_reduce_op_handle.cc DEPS op_handle_base scope lod_tensor ddim memory
         dynload_cuda)
+    nv_library(broad_cast_op_handle SRCS broad_cast_op_handle.cc DEPS op_handle_base scope lod_tensor ddim memory)
+endif()
+
 cc_library(computation_op_handle SRCS computation_op_handle.cc DEPS framework_proto scope place operator op_registry)
 
 cc_library(ssa_graph SRCS ssa_graph.cc DEPS var_handle op_handle_base)
 cc_library(ssa_graph_builder SRCS ssa_graph_builder.cc DEPS ssa_graph)
 
 if(WITH_GPU)
     set(multi_devices_graph_builder_deps nccl_all_reduce_op_handle)
+    nv_test(broad_cast_op_test SRCS broad_cast_op_handle_test.cc DEPS var_handle op_handle_base scope lod_tensor ddim memory
+            device_context broad_cast_op_handle)
 else()
     set(multi_devices_graph_builder_deps)
 endif()

paddle/fluid/framework/details/broad_cast_op_handle.cc

@@ -0,0 +1,103 @@
+//   Copyright (c) 2018 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/framework/details/broad_cast_op_handle.h"
+
+namespace paddle {
+namespace framework {
+namespace details {
+
+Tensor *GetTensorFromVar(Variable *in_var) {
+  if (in_var->IsType<LoDTensor>()) {
+    return in_var->GetMutable<LoDTensor>();
+  } else if (in_var->IsType<SelectedRows>()) {
+    return in_var->GetMutable<SelectedRows>()->mutable_value();
+  } else {
+    PADDLE_THROW("Var should be LoDTensor or SelectedRows");
+  }
+  return nullptr;
+}
+BCastOpHandle::BCastOpHandle(const std::vector<Scope *> &local_scopes,
+                             const std::vector<platform::Place> &places,
+                             const platform::ContextMap &ctxs)
+    : local_scopes_(local_scopes), places_(places), ctxs_(ctxs) {
+  for (auto &p : places_) {
+    this->dev_ctxes_[p] = ctxs_.DevCtx(p);
+  }
+}
+
+void BCastOpHandle::RunImpl() {
+  PADDLE_ENFORCE_EQ(this->inputs_.size(), 1);
+  PADDLE_ENFORCE_EQ(this->outputs_.size(), places_.size());
+
+  // Wait input done, this Wait is asynchronous operation
+  auto in_var_handle = static_cast<VarHandle *>(this->inputs_[0]);
+  auto &in_place = in_var_handle->place_;
+  if (inputs_[0]->generated_op_)
+    inputs_[0]->generated_op_->Wait(dev_ctxes_[in_place]);
+
+  auto iter = std::find(places_.begin(), places_.end(), in_place);
+  if (iter == places_.end()) {
+    PADDLE_THROW("The input of BCast is not in the places_.");
+  }
+
+  int offset = iter - places_.begin();
+  auto in_var = local_scopes_[offset]->FindVar(in_var_handle->name_);
+
+  Tensor *in_tensor = GetTensorFromVar(in_var);
+  for (auto *out : outputs_) {
+    auto out_handle = static_cast<VarHandle *>(out);
+    auto &out_p = out_handle->place_;
+
+    auto iter = std::find(places_.begin(), places_.end(), out_p);
+    if (iter == places_.end()) {
+      PADDLE_THROW("The output of BCast is not in the places_.");
+    }
+    int offset = iter - places_.begin();
+
+    auto *s = local_scopes_[offset];
+    auto out_var = s->FindVar(out_handle->name_);
+
+    PADDLE_ENFORCE_EQ(out_var->Type(), in_var->Type(), "");
+
+    if (in_var->IsType<framework::SelectedRows>()) {
+      auto in_sr = in_var->GetMutable<framework::SelectedRows>();
+      auto out = out_var->GetMutable<framework::SelectedRows>();
+      if (in_sr == out) continue;
+      out->set_height(in_sr->height());
+      out->set_rows(in_sr->rows());
+      out->mutable_value()->Resize(in_sr->value().dims());
+      out->mutable_value()->mutable_data(out_p, in_sr->value().type());
+    } else if (in_var->IsType<framework::LoDTensor>()) {
+      auto in_lod = in_var->GetMutable<framework::LoDTensor>();
+      auto out = out_var->GetMutable<framework::LoDTensor>();
+      if (in_lod == out) continue;
+      out->set_lod(in_lod->lod());
+      out->Resize(in_lod->dims());
+      out->mutable_data(out_p, in_lod->type());
+    } else {
+      PADDLE_THROW("Var should be LoDTensor or SelectedRows");
+    }
+
+    Tensor *out_tensor = GetTensorFromVar(out_var);
+
+    paddle::framework::TensorCopy(*in_tensor, out_p, *(dev_ctxes_[in_place]),
+                                  out_tensor);
+  }
+}
+
+std::string BCastOpHandle::Name() const { return "broadcast"; }
+}  // namespace details
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/details/broad_cast_op_handle.h

@@ -0,0 +1,54 @@
+//   Copyright (c) 2018 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.
+
+#pragma once
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/details/op_handle_base.h"
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/framework/selected_rows.h"
+#include "paddle/fluid/platform/device_context.h"
+
+namespace paddle {
+namespace framework {
+namespace details {
+
+/*
+ * BroadCast the input to all scope.
+ *
+ */
+struct BCastOpHandle : public OpHandleBase {
+  const std::vector<Scope *> &local_scopes_;
+  const std::vector<platform::Place> &places_;
+  const platform::ContextMap &ctxs_;
+
+  BCastOpHandle(const std::vector<Scope *> &local_scopes,
+                const std::vector<platform::Place> &places,
+                const platform::ContextMap &ctxs);
+
+  std::string Name() const override;
+
+  bool IsMultiDeviceTransfer() override { return false; };
+
+ protected:
+  void RunImpl() override;
+};
+
+}  // namespace details
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/details/broad_cast_op_handle_test.cc

@@ -0,0 +1,174 @@
+//   Copyright (c) 2018 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/framework/details/broad_cast_op_handle.h"
+#include "gtest/gtest.h"
+
+#include "paddle/fluid/platform/device_context.h"
+
+namespace f = paddle::framework;
+namespace p = paddle::platform;
+
+// test data amount
+const f::DDim kDims = {20, 20};
+
+class BroadCastTester : public ::testing::Test {
+ public:
+  void SetUp() override {
+    int count = p::GetCUDADeviceCount();
+    if (count <= 1) {
+      LOG(WARNING) << "Cannot test multi-gpu BroadCast, because the CUDA "
+                      "device count is "
+                   << count;
+      exit(0);
+    }
+    for (int i = 0; i < count; ++i) {
+      gpu_list_.emplace_back(p::CUDAPlace(i));
+    }
+    ctxs_ = new p::ContextMap(gpu_list_);
+  }
+
+  template <class T>
+  void BroadCastInitOp(int gpu_id = 0) {
+    for (size_t j = 0; j < gpu_list_.size(); ++j) {
+      local_scope_.push_back(&g_scope_.NewScope());
+      auto* out_var = local_scope_[j]->Var("out");
+      out_var->GetMutable<T>();
+    }
+    auto* in_var = local_scope_[gpu_id]->Var("input");
+    in_var->GetMutable<T>();
+
+    bc_op_handle_ =
+        new f::details::BCastOpHandle(local_scope_, gpu_list_, *ctxs_);
+
+    f::details::VarHandle* in_var_handle = new f::details::VarHandle();
+    in_var_handle->place_ = gpu_list_[gpu_id];
+    in_var_handle->name_ = "input";
+    in_var_handle->version_ = 1;
+    in_var_handle->generated_op_ = nullptr;
+    bc_op_handle_->AddInput(in_var_handle);
+
+    for (size_t j = 0; j < gpu_list_.size(); ++j) {
+      f::details::VarHandle* out_var_handle = new f::details::VarHandle();
+      out_var_handle->place_ = gpu_list_[j];
+      out_var_handle->name_ = "out";
+      out_var_handle->version_ = 2;
+      out_var_handle->generated_op_ = bc_op_handle_;
+      bc_op_handle_->AddOutput(out_var_handle);
+    }
+  }
+  void BroadCastDestroy() {
+    delete ctxs_;
+    for (auto in : bc_op_handle_->inputs_) {
+      delete in;
+    }
+    for (auto out : bc_op_handle_->outputs_) {
+      delete out;
+    }
+    delete bc_op_handle_;
+  }
+
+ public:
+  f::Scope g_scope_;
+  p::ContextMap* ctxs_;
+  std::vector<f::Scope*> local_scope_;
+  std::vector<p::Place> gpu_list_;
+  f::details::BCastOpHandle* bc_op_handle_;
+};
+
+TEST_F(BroadCastTester, BroadCastTestLodTensor) {
+  int gpu_id = 0;
+  BroadCastInitOp<f::LoDTensor>(gpu_id);
+
+  auto in_var = local_scope_[gpu_id]->Var("input");
+  auto in_lod_tensor = in_var->GetMutable<f::LoDTensor>();
+  in_lod_tensor->mutable_data<float>(kDims, gpu_list_[gpu_id]);
+
+  std::vector<float> send_vector(f::product(kDims), gpu_id + 12);
+  for (size_t k = 0; k < send_vector.size(); ++k) {
+    send_vector[k] = k;
+  }
+  f::LoD lod{{0, 10, 20}};
+  paddle::framework::TensorFromVector<float>(
+      send_vector, *(ctxs_->DevCtx(gpu_list_[gpu_id])), in_lod_tensor);
+  in_lod_tensor->set_lod(lod);
+  bc_op_handle_->Run(false);
+
+  ctxs_->WaitAll();
+
+  p::CPUPlace cpu_place;
+  for (size_t j = 0; j < gpu_list_.size(); ++j) {
+    auto out_var = local_scope_[j]->Var("out");
+    auto out_tensor = out_var->Get<f::LoDTensor>();
+    PADDLE_ENFORCE_EQ(out_tensor.lod(), lod, "lod is not equal.");
+
+    f::Tensor result_tensor;
+    f::TensorCopy(out_tensor, cpu_place, *(ctxs_->DevCtx(j)), &result_tensor);
+    float* ct = result_tensor.mutable_data<float>(cpu_place);
+
+    for (int64_t j = 0; j < f::product(kDims); ++j) {
+      ASSERT_NEAR(ct[j], send_vector[j], 1e-5);
+    }
+  }
+
+  BroadCastDestroy();
+}
+
+TEST_F(BroadCastTester, BroadCastTestSelectedRows) {
+  int gpu_id = 0;
+  BroadCastInitOp<f::SelectedRows>(gpu_id);
+
+  auto in_var = local_scope_[gpu_id]->Var("input");
+  auto in_selected_rows = in_var->GetMutable<f::SelectedRows>();
+  auto value = in_selected_rows->mutable_value();
+  value->mutable_data<float>(kDims, gpu_list_[gpu_id]);
+  int height = kDims[0] * 2;
+  std::vector<int64_t> rows{0, 1, 2, 3, 3, 0, 14, 7, 3, 1,
+                            2, 4, 6, 3, 1, 1, 1,  1, 3, 7};
+  in_selected_rows->set_height(height);
+  in_selected_rows->set_rows(rows);
+
+  std::vector<float> send_vector(f::product(kDims));
+  for (size_t k = 0; k < send_vector.size(); ++k) {
+    send_vector[k] = k;
+  }
+  paddle::framework::TensorFromVector<float>(
+      send_vector, *(ctxs_->DevCtx(gpu_list_[gpu_id])), value);
+
+  bc_op_handle_->Run(false);
+
+  ctxs_->WaitAll();
+
+  p::CPUPlace cpu_place;
+  for (size_t j = 0; j < gpu_list_.size(); ++j) {
+    auto out_var = local_scope_[j]->Var("out");
+    auto& out_select_rows = out_var->Get<f::SelectedRows>();
+    auto rt = out_select_rows.value();
+
+    PADDLE_ENFORCE_EQ(out_select_rows.height(), height, "height is not equal.");
+    for (size_t k = 0; k < out_select_rows.rows().size(); ++k) {
+      PADDLE_ENFORCE_EQ(out_select_rows.rows()[k], rows[k]);
+    }
+
+    f::Tensor result_tensor;
+    f::TensorCopy(rt, cpu_place, *(ctxs_->DevCtx(j)), &result_tensor);
+    float* ct = result_tensor.data<float>();
+
+    for (int64_t j = 0; j < f::product(kDims); ++j) {
+      ASSERT_NEAR(ct[j], send_vector[j], 1e-5);
+    }
+  }
+
+  BroadCastDestroy();
+}