parallel_do_op.cc 13.5 KB
Newer Older
Y
Yang Yang 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Y
Yang Yang 已提交
3 4 5
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
Y
Yang Yang 已提交
6

Y
Yang Yang 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yang Yang 已提交
8

Y
Yang Yang 已提交
9 10 11 12 13
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. */
Y
Yang Yang 已提交
14 15

#include <vector>
Y
Yang Yang 已提交
16

Y
Yang Yang 已提交
17 18
#include "paddle/framework/executor.h"
#include "paddle/framework/op_registry.h"
Y
Yang Yu 已提交
19
#include "paddle/framework/threadpool.h"
Y
Yang Yu 已提交
20
#include "paddle/operators/detail/safe_ref.h"
Y
Yang Yang 已提交
21 22 23 24

namespace paddle {
namespace operators {

Y
Yang Yu 已提交
25 26 27
static constexpr char kInputs[] = "inputs";
static constexpr char kParameters[] = "parameters";
static constexpr char kPlaces[] = "places";
Y
Yang Yang 已提交
28

Y
Yang Yu 已提交
29 30
static constexpr char kOutputs[] = "outputs";
static constexpr char kParallelScopes[] = "parallel_scopes";
Y
Yang Yang 已提交
31

Y
Yang Yu 已提交
32
static constexpr char kParallelBlock[] = "sub_block";
Y
Yang Yang 已提交
33

Y
Yang Yang 已提交
34
using LoDTensor = framework::LoDTensor;
Y
Yang Yang 已提交
35
using SelectedRows = framework::SelectedRows;
Y
Yang Yang 已提交
36

Y
Yu Yang 已提交
37 38
static void SplitTensorAndMoveTensorToScopes(
    const framework::Scope &scope, std::vector<framework::Scope *> *sub_scopes,
Y
Yang Yang 已提交
39 40
    const std::vector<platform::Place> &places,
    const std::vector<std::string> &names) {
Y
Yu Yang 已提交
41
  size_t num_sub_scopes = 0;
Y
Yang Yang 已提交
42
  for (auto &argu : names) {
Y
Yang Yu 已提交
43 44 45 46
    const auto &tensor =
        detail::Ref(scope.FindVar(argu),
                    "Cannot find variable %s in the parent scope", argu)
            .Get<LoDTensor>();
Y
Yang Yang 已提交
47 48 49
    auto lod_tensors = tensor.SplitLoDTensor(places);

    for (auto &lod : lod_tensors) {
Y
Yang Yang 已提交
50
      VLOG(3) << lod.dims();
Y
Yang Yang 已提交
51
    }
Y
Yu Yang 已提交
52 53 54 55 56 57 58 59 60 61 62 63
    if (num_sub_scopes == 0) {
      num_sub_scopes = lod_tensors.size();
    } else {
      PADDLE_ENFORCE_EQ(num_sub_scopes, lod_tensors.size());
    }
    PADDLE_ENFORCE_NE(num_sub_scopes, 0);
    if (sub_scopes->size() == 0) {
      sub_scopes->reserve(num_sub_scopes);
      for (size_t i = 0; i < num_sub_scopes; ++i) {
        sub_scopes->emplace_back(&scope.NewScope());
      }
    }
Y
Yang Yang 已提交
64

Y
Yu Yang 已提交
65
    for (size_t i = 0; i < lod_tensors.size(); ++i) {
Y
Yang Yu 已提交
66 67 68
      *detail::Ref(sub_scopes->at(i)->Var(argu),
                   "Cannot find variable in the sub-scope", argu)
           .GetMutable<LoDTensor>() = lod_tensors[i];
Y
Yang Yang 已提交
69 70 71 72
    }
  }
}

Y
Yang Yang 已提交
73 74 75
inline void CopyOrShare(const framework::Variable &src,
                        const platform::Place &dst_place,
                        framework::Variable *dst) {
Y
Yang Yang 已提交
76 77 78
  if (src.IsType<LoDTensor>()) {
    if (src.Get<LoDTensor>().place() == dst_place) {
      dst->GetMutable<LoDTensor>()->ShareDataWith(src.Get<LoDTensor>());
D
dzhwinter 已提交
79
      dst->GetMutable<LoDTensor>()->set_lod(src.Get<LoDTensor>().lod());
Y
Yang Yang 已提交
80 81
    } else {
      Copy(src.Get<LoDTensor>(), dst_place, dst->GetMutable<LoDTensor>());
D
dzhwinter 已提交
82
      framework::LoD lod(src.Get<LoDTensor>().lod());
D
dzhwinter 已提交
83 84
      lod.CopyToPeer(dst_place);
      dst->GetMutable<LoDTensor>()->set_lod(lod);
Y
Yang Yang 已提交
85 86 87 88 89
    }
  } else if (src.IsType<SelectedRows>()) {
    auto &src_sr = src.Get<SelectedRows>();
    auto *dst_sr = dst->GetMutable<SelectedRows>();
    dst_sr->set_height(src_sr.height());
Y
Yang Yang 已提交
90 91
    if (src_sr.value().place() == dst_place) {
      dst_sr->mutable_value()->ShareDataWith(src_sr.value());
D
dzhwinter 已提交
92
      dst_sr->set_rows(src_sr.rows());
Y
Yang Yang 已提交
93 94
    } else {
      Copy(src_sr.value(), dst_place, dst_sr->mutable_value());
D
dzhwinter 已提交
95
      framework::Vector<int64_t> lod(src_sr.rows());
D
dzhwinter 已提交
96 97
      lod.CopyToPeer(dst_place);
      dst_sr->set_rows(lod);
Y
Yang Yang 已提交
98
    }
Y
Yang Yang 已提交
99 100 101 102 103
  } else {
    PADDLE_THROW("Expect LoDTensor/SelectedRows, get %s", src.Type().name());
  }
}

Y
Yang Yang 已提交
104 105 106 107 108 109
void WaitOnPlace(const platform::Place place) {
  platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
  auto &dev_ctx = *pool.Get(place);
  dev_ctx.Wait();
}

110 111 112 113 114 115 116 117 118
void WaitOnPlaces(const std::vector<platform::Place> places) {
  platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();

  for (auto &place : places) {
    auto &dev_ctx = *pool.Get(place);
    dev_ctx.Wait();
  }
}

Y
Yang Yang 已提交
119
class ParallelDoOp : public framework::OperatorBase {
Y
Yang Yang 已提交
120 121 122 123 124
 public:
  ParallelDoOp(const std::string &type,
               const framework::VariableNameMap &inputs,
               const framework::VariableNameMap &outputs,
               const framework::AttributeMap &attrs)
Y
Yu Yang 已提交
125
      : framework::OperatorBase(type, inputs, outputs, attrs) {}
Y
Yang Yang 已提交
126

127 128 129
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
Y
Yang Yang 已提交
130 131 132 133 134
    // get device context from pool
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(place);

    auto *block = Attr<framework::BlockDesc *>(kParallelBlock);
Y
Yang Yang 已提交
135
    auto *program = block->Program();
Y
Yang Yang 已提交
136

137
    auto &places = scope.FindVar(Input(kPlaces))->Get<platform::PlaceList>();
Y
Yang Yang 已提交
138

Y
Yang Yang 已提交
139 140 141
    auto &sub_scopes = *scope.FindVar(Output(kParallelScopes))
                            ->GetMutable<std::vector<framework::Scope *>>();

142
    // split input
Y
Yu Yang 已提交
143
    SplitTensorAndMoveTensorToScopes(scope, &sub_scopes, places,
Y
Yang Yang 已提交
144
                                     Inputs(kInputs));
Y
Yu Yang 已提交
145

146 147 148 149 150
    // copy parameter
    for (auto &param : Inputs(kParameters)) {
      PADDLE_ENFORCE(scope.FindVar(param)->IsType<LoDTensor>(),
                     "Only support parameter type as LoDTensor");
      auto &src = scope.FindVar(param)->Get<LoDTensor>();
Y
Yu Yang 已提交
151
      for (size_t i = 0; i < sub_scopes.size(); ++i) {
152 153 154 155
        auto &place = places[i];
        auto *sub_scope = sub_scopes[i];
        auto *dst = sub_scope->Var(param)->GetMutable<LoDTensor>();
        framework::Copy(src, place, dst);
D
dzhwinter 已提交
156 157 158
        framework::LoD lod(src.lod());
        lod.CopyToPeer(place);
        dst->set_lod(lod);
159 160 161
      }
    }
    WaitOnPlaces(places);
Y
Yang Yang 已提交
162

Y
Yang Yu 已提交
163 164
    std::vector<std::future<void>> workers;
    workers.reserve(places.size());
Y
Yu Yang 已提交
165
    for (size_t place_idx = 0; place_idx < sub_scopes.size(); ++place_idx) {
Y
Yang Yang 已提交
166 167 168
      auto &place = places[place_idx];
      auto *cur_scope = sub_scopes[place_idx];

Y
Yang Yu 已提交
169 170
      workers.emplace_back(framework::Async([program, cur_scope, place, block] {
        framework::Executor executor(place);
Y
Yang Yang 已提交
171 172 173 174 175
        executor.Run(*program, cur_scope, block->ID(),
                     false /*create_local_scope*/);
      }));
    }
    for (auto &worker : workers) {
Y
Yang Yu 已提交
176
      worker.wait();
Y
Yang Yang 已提交
177
    }
178
    WaitOnPlaces(places);
Y
Yang Yang 已提交
179 180 181 182

    // merge output
    for (auto &o_name : Outputs(kOutputs)) {
      std::vector<const framework::LoDTensor *> lod_tensors;
Y
Yang Yu 已提交
183
      lod_tensors.reserve(sub_scopes.size());
Y
Yang Yang 已提交
184
      for (auto *sub_scope : sub_scopes) {
Y
Yang Yu 已提交
185
        lod_tensors.emplace_back(&sub_scope->FindVar(o_name)->Get<LoDTensor>());
Y
Yang Yang 已提交
186 187 188 189 190 191
      }

      auto *lod_tensor_to_be_merged =
          scope.FindVar(o_name)->GetMutable<LoDTensor>();
      lod_tensor_to_be_merged->MergeLoDTensor(lod_tensors, dev_ctx.GetPlace());
    }
192
    WaitOnPlaces(places);
Y
Yang Yang 已提交
193
  }
Y
Yang Yang 已提交
194 195 196 197
};

class ParallelDoOpProtoMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yang Yang 已提交
198
  ParallelDoOpProtoMaker(OpProto *proto, framework::OpAttrChecker *op_checker)
Y
Yang Yang 已提交
199 200 201 202 203 204
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput(kInputs, "").AsDuplicable();
    AddInput(kParameters, "").AsDuplicable();
    AddInput(kPlaces, "");
    AddOutput(kOutputs, "").AsDuplicable();
    AddOutput(kParallelScopes, "");
Y
Yang Yang 已提交
205
    AddAttr<framework::BlockDesc *>(kParallelBlock, "");
Y
Yang Yang 已提交
206 207 208 209 210 211
    AddComment(R"DOC(
ParallelDo Operator.
)DOC");
  }
};

Y
Yu Yang 已提交
212
class ParallelDoGradOp : public framework::OperatorBase {
Y
Yang Yang 已提交
213 214 215 216 217
 public:
  ParallelDoGradOp(const std::string &type,
                   const framework::VariableNameMap &inputs,
                   const framework::VariableNameMap &outputs,
                   const framework::AttributeMap &attrs)
Y
Yu Yang 已提交
218
      : framework::OperatorBase(type, inputs, outputs, attrs) {}
Y
Yang Yang 已提交
219

220 221 222
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
Y
Yang Yang 已提交
223
    auto *block = Attr<framework::BlockDesc *>(kParallelBlock);
Y
Yang Yang 已提交
224 225 226 227 228
    auto *program = block->Program();

    auto &sub_scopes = scope.FindVar(Input(kParallelScopes))
                           ->Get<std::vector<framework::Scope *>>();

229
    auto &places = scope.FindVar(Input(kPlaces))->Get<platform::PlaceList>();
Y
Yang Yang 已提交
230 231

    // feed output@grad
Y
Yu Yang 已提交
232 233 234
    SplitTensorAndMoveTensorToScopes(
        scope, const_cast<std::vector<framework::Scope *> *>(&sub_scopes),
        places, Inputs(framework::GradVarName(kOutputs)));
235
    WaitOnPlaces(places);
Y
Yang Yang 已提交
236 237

    // exe run
Y
Yang Yu 已提交
238
    std::vector<std::future<void>> workers;
Y
Yu Yang 已提交
239 240 241
    for (size_t i = 0; i < sub_scopes.size(); ++i) {
      auto &place = places[i];
      auto *cur_scope = sub_scopes[i];
Y
Yang Yang 已提交
242 243

      // execute
Y
Yang Yu 已提交
244 245
      workers.emplace_back(framework::Async([program, cur_scope, place, block] {
        framework::Executor executor(place);
Y
Yang Yang 已提交
246 247 248 249 250
        executor.Run(*program, cur_scope, block->ID(),
                     false /*create_local_scope*/);
      }));
    }
    for (auto &worker : workers) {
Y
Yang Yu 已提交
251
      worker.wait();
Y
Yang Yang 已提交
252
    }
253
    WaitOnPlaces(places);
Y
Yang Yang 已提交
254

Y
Yang Yang 已提交
255 256 257 258 259 260 261
    AccumulateGrad(scope, place, sub_scopes, places);
  }

  void AccumulateGrad(const framework::Scope &scope,
                      const platform::Place &place,
                      const std::vector<framework::Scope *> &sub_scopes,
                      const platform::PlaceList &places) const {
Y
Yang Yang 已提交
262
    for (auto &s : Outputs(framework::GradVarName(kParameters))) {
Y
Yang Yang(Tony) 已提交
263 264
      VLOG(3) << "Accumulating " << s;
      if (s == framework::kEmptyVarName) continue;
Y
Yang Yang 已提交
265
      std::string tmp_name;
Y
Yang Yang 已提交
266
      auto *tmp = sub_scopes[0]->Var(&tmp_name);
Y
Yu Yang 已提交
267 268

      for (size_t i = 1; i < sub_scopes.size(); ++i) {
Y
Yang Yang 已提交
269
        CopyOrShare(*sub_scopes[i]->FindVar(s), places[0], tmp);
Y
Yang Yang(Tony) 已提交
270
        WaitOnPlaces(places);
Y
Yang Yang 已提交
271

Y
Yang Yang 已提交
272
        auto sum_op = framework::OpRegistry::CreateOp(
Y
Yu Yang 已提交
273
            "sum", {{"X", {s, tmp_name}}}, {{"Out", {s}}},
Y
Yang Yang 已提交
274
            framework::AttributeMap{});
Y
Yang Yang(Tony) 已提交
275
        VLOG(10) << sum_op->DebugStringEx(sub_scopes[0]);
276
        sum_op->Run(*sub_scopes[0], places[0]);
Y
Yang Yang 已提交
277
        WaitOnPlace(places[0]);
Y
Yang Yang 已提交
278 279
      }

Y
Yang Yang 已提交
280
      CopyOrShare(*sub_scopes[0]->FindVar(s), place, scope.FindVar(s));
Y
Yang Yang 已提交
281
    }
Y
Yang Yang 已提交
282
    WaitOnPlaces(places);
Y
Yang Yang 已提交
283
  }
Y
Yang Yang 已提交
284 285
};

Y
Yu Yang 已提交
286 287 288 289 290 291 292
std::ostream &operator<<(std::ostream &sout,
                         const std::vector<std::string> &strs) {
  std::copy(strs.begin(), strs.end(),
            std::ostream_iterator<std::string>(sout, ","));
  return sout;
}

Y
Yang Yang 已提交
293 294 295 296 297
class ParallelDoGradOpDescMaker : public framework::SingleGradOpDescMaker {
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
Y
Yang Yang 已提交
298 299
  virtual std::unique_ptr<framework::OpDesc> Apply() const {
    auto *grad = new framework::OpDesc();
Y
Yang Yang 已提交
300
    grad->SetType("parallel_do_grad");
Y
Yang Yang 已提交
301
    for (auto &input_param : this->InputNames()) {
Y
Yang Yang 已提交
302
      VLOG(3) << input_param;
Y
Yang Yang 已提交
303
      grad->SetInput(input_param, this->Input(input_param));
304 305 306 307
      if (input_param != kPlaces) {
        grad->SetOutput(framework::GradVarName(input_param),
                        this->InputGrad(input_param, false));
      }
Y
Yang Yang 已提交
308
    }
Y
Yang Yu 已提交
309 310 311 312 313 314 315 316 317 318 319
    auto *g_block = this->grad_block_[0];

    // All variable name that needed by gradient operators
    std::unordered_set<std::string> all_inputs_in_grad_blocks;

    for (size_t i = 0; i < g_block->OpSize(); ++i) {
      auto *op = g_block->Op(i);
      for (auto &var_name : op->InputArgumentNames()) {
        all_inputs_in_grad_blocks.insert(var_name);
      }
    }
Y
Yang Yang 已提交
320 321 322 323 324 325 326 327

    for (auto &output_param : this->OutputNames()) {
      if (output_param == kParallelScopes) {
        grad->SetInput(output_param, this->Output(output_param));
        grad->SetInput(framework::GradVarName(output_param),
                       this->Output(output_param));
      } else {
        grad->SetInput(output_param, this->Output(output_param));
Y
Yang Yu 已提交
328 329 330
        std::vector<std::string> og_names;
        for (auto &og_name : this->OutputGrad(output_param)) {
          if (all_inputs_in_grad_blocks.count(og_name) != 0) {
Y
Yang Yu 已提交
331 332
            // there are some gradient operators who need the OG. So make this
            // OG as an input of parallel.do
Y
Yang Yu 已提交
333 334
            og_names.push_back(og_name);
          }
Y
Yang Yu 已提交
335 336
          // else, there is no operator who need the OG. Do not use this OG as
          // an input
Y
Yang Yu 已提交
337 338
        }
        grad->SetInput(framework::GradVarName(output_param), og_names);
Y
Yang Yang 已提交
339 340 341 342 343
      }
    }
    grad->SetAttrMap(this->Attrs());
    grad->SetBlockAttr(kParallelBlock, *grad_block_[0]);

Y
Yang Yang 已提交
344
    return std::unique_ptr<framework::OpDesc>(grad);
Y
Yang Yang 已提交
345 346 347 348 349 350
  }
};

class ParallelDoGradOpShapeInference : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *ctx) const override {
Y
Yu Yang 已提交
351
    PADDLE_ENFORCE(ctx->HasInputs(kParameters));
Y
Yang Yang 已提交
352
    PADDLE_ENFORCE(ctx->HasInputs(kInputs));
Y
Yang Yang(Tony) 已提交
353
    PADDLE_ENFORCE(ctx->HasInputs(kOutputs));
Y
Yu Yang 已提交
354 355 356 357 358 359 360 361 362 363 364 365 366 367

    ctx->SetOutputsDim(framework::GradVarName(kParameters),
                       ctx->GetInputsDim(kParameters));

    auto i_dims = ctx->GetInputsDim(kInputs);
    auto ig_names = ctx->Outputs(framework::GradVarName(kInputs));

    for (size_t i = 0; i < ig_names.size(); ++i) {
      auto &ig_name = ig_names[i];
      if (ig_name == framework::kEmptyVarName) {
        continue;
      }

      ctx->SetDims({ig_name}, {i_dims[i]});
Y
Yang Yang 已提交
368
    }
Y
Yu Yang 已提交
369

Y
Yang Yang(Tony) 已提交
370 371 372 373 374 375 376 377
    auto p_dims = ctx->GetInputsDim(kParameters);
    auto pg_names = ctx->Outputs(framework::GradVarName(kParameters));
    for (size_t i = 0; i < pg_names.size(); ++i) {
      auto &pg_name = pg_names[i];
      if (pg_name == framework::kEmptyVarName) {
        continue;
      }
      ctx->SetDims({pg_name}, {p_dims[i]});
Y
Yang Yang 已提交
378
    }
Y
Yang Yang 已提交
379 380 381 382 383 384 385 386 387 388 389
  }
};

}  // namespace operators
}  // namespace paddle

REGISTER_OPERATOR(parallel_do, paddle::operators::ParallelDoOp,
                  paddle::operators::ParallelDoOpProtoMaker,
                  paddle::operators::ParallelDoGradOpDescMaker);
REGISTER_OPERATOR(parallel_do_grad, paddle::operators::ParallelDoGradOp,
                  paddle::operators::ParallelDoGradOpShapeInference);