gather_op_handle_test.cc 6.6 KB
Newer Older
C
chengduoZH 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
//   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/gather_op_handle.h"
#include "gtest/gtest.h"

#include "paddle/fluid/platform/device_context.h"

C
chengduoZH 已提交
20 21 22
namespace paddle {
namespace framework {
namespace details {
C
chengduoZH 已提交
23 24 25 26 27 28
namespace f = paddle::framework;
namespace p = paddle::platform;

// test data amount
const f::DDim kDims = {20, 20};

29 30 31
struct TestGatherOpHandle {
  std::vector<std::unique_ptr<p::DeviceContext>> ctxs_;
  std::vector<Scope*> local_scopes_;
C
chengduoZH 已提交
32
  std::vector<Scope*> param_scopes_;
33 34 35 36 37 38 39 40 41 42 43
  Scope g_scope_;
  std::unique_ptr<OpHandleBase> op_handle_;
  std::vector<std::unique_ptr<VarHandleBase>> vars_;
  std::vector<p::Place> gpu_list_;

  void WaitAll() {
    for (size_t j = 0; j < ctxs_.size(); ++j) {
      ctxs_[j]->Wait();
    }
  }

C
chengduoZH 已提交
44
  void InitCtxOnGpu(bool use_gpu) {
C
chengduoZH 已提交
45 46 47 48
    if (use_gpu) {
#ifdef PADDLE_WITH_CUDA
      int count = p::GetCUDADeviceCount();
      if (count <= 1) {
49
        LOG(WARNING) << "Cannot test multi-gpu Broadcast, because the CUDA "
C
chengduoZH 已提交
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
                        "device count is "
                     << count;
        exit(0);
      }
      for (int i = 0; i < count; ++i) {
        auto p = p::CUDAPlace(i);
        gpu_list_.push_back(p);
        ctxs_.emplace_back(new p::CUDADeviceContext(p));
      }
#else
      PADDLE_THROW("CUDA is not support.");
#endif
    } else {
      int count = 8;
      for (int i = 0; i < count; ++i) {
        auto p = p::CPUPlace();
        gpu_list_.push_back(p);
        ctxs_.emplace_back(new p::CPUDeviceContext(p));
      }
    }
  }

72
  void InitGatherOp(size_t input_scope_idx) {
X
Xin Pan 已提交
73
    std::vector<std::unique_ptr<ir::Node>> nodes;
C
chengduoZH 已提交
74
    for (size_t j = 0; j < gpu_list_.size(); ++j) {
75
      local_scopes_.push_back(&(g_scope_.NewScope()));
C
chengduoZH 已提交
76 77 78 79 80 81
      Scope& local_scope = local_scopes_.back()->NewScope();
      *local_scopes_.back()
           ->Var(details::kLocalExecScopeName)
           ->GetMutable<Scope*>() = &local_scope;
      local_scope.Var("input");
      param_scopes_.emplace_back(&local_scope);
C
chengduoZH 已提交
82
    }
C
chengduoZH 已提交
83
    param_scopes_[input_scope_idx]->Var("out");
C
chengduoZH 已提交
84

X
Xin Pan 已提交
85
    nodes.emplace_back(new ir::Node("node"));
X
Xin Pan 已提交
86 87
    op_handle_.reset(
        new GatherOpHandle(nodes.back().get(), local_scopes_, gpu_list_));
88
    // add input
C
chengduoZH 已提交
89
    for (size_t j = 0; j < gpu_list_.size(); ++j) {
X
Xin Pan 已提交
90
      op_handle_->SetDeviceContext(gpu_list_[j], ctxs_[j].get());
X
Xin Pan 已提交
91
      nodes.emplace_back(new ir::Node("node1"));
X
Xin Pan 已提交
92 93
      auto* in_var_handle =
          new VarHandle(nodes.back().get(), 1, j, "input", gpu_list_[j]);
Y
Yu Yang 已提交
94
      vars_.emplace_back(in_var_handle);
95
      op_handle_->AddInput(in_var_handle);
C
chengduoZH 已提交
96 97
    }

98
    // add dummy var
X
Xin Pan 已提交
99
    nodes.emplace_back(new ir::Node("node2"));
X
Xin Pan 已提交
100
    vars_.emplace_back(new DummyVarHandle(nodes.back().get()));
101 102
    DummyVarHandle* in_dummy_var_handle =
        static_cast<DummyVarHandle*>(vars_.back().get());
X
Xin Pan 已提交
103
    in_dummy_var_handle->ClearGeneratedOp();
104 105 106
    op_handle_->AddInput(in_dummy_var_handle);

    // add output
X
Xin Pan 已提交
107
    nodes.emplace_back(new ir::Node("node3"));
X
Xin Pan 已提交
108 109
    auto* out_var_handle = new VarHandle(nodes.back().get(), 2, input_scope_idx,
                                         "out", gpu_list_[input_scope_idx]);
Y
Yu Yang 已提交
110
    vars_.emplace_back(out_var_handle);
111
    op_handle_->AddOutput(out_var_handle);
C
chengduoZH 已提交
112

113
    // add dummy var
X
Xin Pan 已提交
114
    nodes.emplace_back(new ir::Node("node4"));
X
Xin Pan 已提交
115
    vars_.emplace_back(new DummyVarHandle(nodes.back().get()));
116 117 118 119
    DummyVarHandle* dummy_var_handle =
        static_cast<DummyVarHandle*>(vars_.back().get());
    op_handle_->AddOutput(dummy_var_handle);
  }
C
chengduoZH 已提交
120

121
  void TestGatherSelectedRows(size_t output_scope_idx) {
C
chengduoZH 已提交
122 123 124 125 126 127 128 129 130 131
    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};
    std::vector<float> send_vector(f::product(kDims));
    for (size_t k = 0; k < send_vector.size(); ++k) {
      send_vector[k] = k;
    }

    for (size_t input_scope_idx = 0; input_scope_idx < gpu_list_.size();
         ++input_scope_idx) {
C
chengduoZH 已提交
132 133
      auto in_var = param_scopes_.at(input_scope_idx)->FindVar("input");
      PADDLE_ENFORCE_NOT_NULL(in_var);
C
chengduoZH 已提交
134 135 136 137 138 139 140 141 142 143 144 145
      auto in_selected_rows = in_var->GetMutable<f::SelectedRows>();
      auto value = in_selected_rows->mutable_value();
      value->mutable_data<float>(kDims, gpu_list_[input_scope_idx]);

      in_selected_rows->set_height(height);
      in_selected_rows->set_rows(rows);

      paddle::framework::TensorFromVector<float>(
          send_vector, *(ctxs_[input_scope_idx]), value);
      value->Resize(kDims);
    }

C
chengduoZH 已提交
146 147
    auto out_var = param_scopes_.at(output_scope_idx)->FindVar("out");
    PADDLE_ENFORCE_NOT_NULL(out_var);
148 149
    auto out_selected_rows = out_var->GetMutable<f::SelectedRows>();

C
chengduoZH 已提交
150
    auto in_var = param_scopes_.at(output_scope_idx)->FindVar("input");
151 152 153 154 155 156
    auto in_selected_rows = in_var->GetMutable<f::SelectedRows>();

    out_selected_rows->mutable_value()->ShareDataWith(
        in_selected_rows->value());

    op_handle_->Run(false);
C
chengduoZH 已提交
157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173

    WaitAll();

    p::CPUPlace cpu_place;

    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 % rows.size()]);
    }

    f::Tensor result_tensor;
    f::TensorCopy(rt, cpu_place, *(ctxs_[output_scope_idx]), &result_tensor);
    float* ct = result_tensor.data<float>();

C
chengduoZH 已提交
174 175
    for (int64_t j = 0;
         j < f::product(kDims) * static_cast<int64_t>(gpu_list_.size()); ++j) {
C
chengduoZH 已提交
176 177 178 179 180
      ASSERT_NEAR(ct[j], send_vector[j % send_vector.size()], 1e-5);
    }
  }
};

181 182 183 184 185 186
TEST(GatherTester, TestCPUGatherTestSelectedRows) {
  TestGatherOpHandle test_op;
  size_t input_scope_idx = 0;
  test_op.InitCtxOnGpu(false);
  test_op.InitGatherOp(input_scope_idx);
  test_op.TestGatherSelectedRows(input_scope_idx);
C
chengduoZH 已提交
187 188 189 190
}

#ifdef PADDLE_WITH_CUDA

191 192 193 194 195 196
TEST(GatherTester, TestGPUGatherTestSelectedRows) {
  TestGatherOpHandle test_op;
  size_t input_scope_idx = 0;
  test_op.InitCtxOnGpu(false);
  test_op.InitGatherOp(input_scope_idx);
  test_op.TestGatherSelectedRows(input_scope_idx);
C
chengduoZH 已提交
197 198
}
#endif
C
chengduoZH 已提交
199 200 201
}  // namespace details
}  // namespace framework
}  // namespace paddle