“ea4d08dab6b76a685e9277f28daf6b594912b97f”上不存在“paddle/legacy/gserver/layers/ContextProjection.cpp”
parallel_do_op.cc 13.1 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
Yi Wang 已提交
17 18 19 20
#include "paddle/fluid/framework/executor.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/threadpool.h"
#include "paddle/fluid/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 82 83 84 85 86
    } else {
      Copy(src.Get<LoDTensor>(), dst_place, dst->GetMutable<LoDTensor>());
    }
  } 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 已提交
87 88
    if (src_sr.value().place() == dst_place) {
      dst_sr->mutable_value()->ShareDataWith(src_sr.value());
D
dzhwinter 已提交
89
      dst_sr->set_rows(src_sr.rows());
Y
Yang Yang 已提交
90 91 92
    } else {
      Copy(src_sr.value(), dst_place, dst_sr->mutable_value());
    }
Y
Yang Yang 已提交
93 94 95 96 97
  } else {
    PADDLE_THROW("Expect LoDTensor/SelectedRows, get %s", src.Type().name());
  }
}

Y
Yang Yang 已提交
98 99 100 101 102 103
void WaitOnPlace(const platform::Place place) {
  platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
  auto &dev_ctx = *pool.Get(place);
  dev_ctx.Wait();
}

104 105 106 107 108 109 110 111 112
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 已提交
113
class ParallelDoOp : public framework::OperatorBase {
Y
Yang Yang 已提交
114 115 116 117 118
 public:
  ParallelDoOp(const std::string &type,
               const framework::VariableNameMap &inputs,
               const framework::VariableNameMap &outputs,
               const framework::AttributeMap &attrs)
Y
Yu Yang 已提交
119
      : framework::OperatorBase(type, inputs, outputs, attrs) {}
Y
Yang Yang 已提交
120 121

  void Run(const framework::Scope &scope,
Y
Yang Yang 已提交
122 123 124 125 126 127
           const platform::Place &place) const override {
    // 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 已提交
128
    auto *program = block->Program();
Y
Yang Yang 已提交
129

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

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

135
    // split input
Y
Yu Yang 已提交
136
    SplitTensorAndMoveTensorToScopes(scope, &sub_scopes, places,
Y
Yang Yang 已提交
137
                                     Inputs(kInputs));
Y
Yu Yang 已提交
138

139 140 141 142 143
    // 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 已提交
144
      for (size_t i = 0; i < sub_scopes.size(); ++i) {
145 146 147 148 149 150 151
        auto &place = places[i];
        auto *sub_scope = sub_scopes[i];
        auto *dst = sub_scope->Var(param)->GetMutable<LoDTensor>();
        framework::Copy(src, place, dst);
      }
    }
    WaitOnPlaces(places);
Y
Yang Yang 已提交
152

Y
Yang Yu 已提交
153 154
    std::vector<std::future<void>> workers;
    workers.reserve(places.size());
Y
Yu Yang 已提交
155
    for (size_t place_idx = 0; place_idx < sub_scopes.size(); ++place_idx) {
Y
Yang Yang 已提交
156 157 158
      auto &place = places[place_idx];
      auto *cur_scope = sub_scopes[place_idx];

Y
Yang Yu 已提交
159 160
      workers.emplace_back(framework::Async([program, cur_scope, place, block] {
        framework::Executor executor(place);
Y
Yang Yang 已提交
161 162 163 164 165
        executor.Run(*program, cur_scope, block->ID(),
                     false /*create_local_scope*/);
      }));
    }
    for (auto &worker : workers) {
Y
Yang Yu 已提交
166
      worker.wait();
Y
Yang Yang 已提交
167
    }
168
    WaitOnPlaces(places);
Y
Yang Yang 已提交
169 170 171 172

    // merge output
    for (auto &o_name : Outputs(kOutputs)) {
      std::vector<const framework::LoDTensor *> lod_tensors;
Y
Yang Yu 已提交
173
      lod_tensors.reserve(sub_scopes.size());
Y
Yang Yang 已提交
174
      for (auto *sub_scope : sub_scopes) {
Y
Yang Yu 已提交
175
        lod_tensors.emplace_back(&sub_scope->FindVar(o_name)->Get<LoDTensor>());
Y
Yang Yang 已提交
176 177 178 179 180 181
      }

      auto *lod_tensor_to_be_merged =
          scope.FindVar(o_name)->GetMutable<LoDTensor>();
      lod_tensor_to_be_merged->MergeLoDTensor(lod_tensors, dev_ctx.GetPlace());
    }
182
    WaitOnPlaces(places);
Y
Yang Yang 已提交
183
  }
Y
Yang Yang 已提交
184 185 186 187
};

class ParallelDoOpProtoMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yang Yang 已提交
188
  ParallelDoOpProtoMaker(OpProto *proto, framework::OpAttrChecker *op_checker)
Y
Yang Yang 已提交
189 190 191 192 193 194
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput(kInputs, "").AsDuplicable();
    AddInput(kParameters, "").AsDuplicable();
    AddInput(kPlaces, "");
    AddOutput(kOutputs, "").AsDuplicable();
    AddOutput(kParallelScopes, "");
Y
Yang Yang 已提交
195
    AddAttr<framework::BlockDesc *>(kParallelBlock, "");
Y
Yang Yang 已提交
196 197 198 199 200 201
    AddComment(R"DOC(
ParallelDo Operator.
)DOC");
  }
};

Y
Yu Yang 已提交
202
class ParallelDoGradOp : public framework::OperatorBase {
Y
Yang Yang 已提交
203 204 205 206 207
 public:
  ParallelDoGradOp(const std::string &type,
                   const framework::VariableNameMap &inputs,
                   const framework::VariableNameMap &outputs,
                   const framework::AttributeMap &attrs)
Y
Yu Yang 已提交
208
      : framework::OperatorBase(type, inputs, outputs, attrs) {}
Y
Yang Yang 已提交
209 210

  void Run(const framework::Scope &scope,
Y
Yang Yang 已提交
211 212
           const platform::Place &place) const override {
    auto *block = Attr<framework::BlockDesc *>(kParallelBlock);
Y
Yang Yang 已提交
213 214 215 216 217
    auto *program = block->Program();

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

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

    // feed output@grad
Y
Yu Yang 已提交
221 222 223
    SplitTensorAndMoveTensorToScopes(
        scope, const_cast<std::vector<framework::Scope *> *>(&sub_scopes),
        places, Inputs(framework::GradVarName(kOutputs)));
224
    WaitOnPlaces(places);
Y
Yang Yang 已提交
225 226

    // exe run
Y
Yang Yu 已提交
227
    std::vector<std::future<void>> workers;
Y
Yu Yang 已提交
228 229 230
    for (size_t i = 0; i < sub_scopes.size(); ++i) {
      auto &place = places[i];
      auto *cur_scope = sub_scopes[i];
Y
Yang Yang 已提交
231 232

      // execute
Y
Yang Yu 已提交
233 234
      workers.emplace_back(framework::Async([program, cur_scope, place, block] {
        framework::Executor executor(place);
Y
Yang Yang 已提交
235 236 237 238 239
        executor.Run(*program, cur_scope, block->ID(),
                     false /*create_local_scope*/);
      }));
    }
    for (auto &worker : workers) {
Y
Yang Yu 已提交
240
      worker.wait();
Y
Yang Yang 已提交
241
    }
242
    WaitOnPlaces(places);
Y
Yang Yang 已提交
243

Y
Yang Yang 已提交
244 245 246 247 248 249 250
    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 已提交
251
    for (auto &s : Outputs(framework::GradVarName(kParameters))) {
Y
Yang Yang(Tony) 已提交
252 253
      VLOG(3) << "Accumulating " << s;
      if (s == framework::kEmptyVarName) continue;
Y
Yang Yang 已提交
254
      std::string tmp_name;
Y
Yang Yang 已提交
255
      auto *tmp = sub_scopes[0]->Var(&tmp_name);
Y
Yu Yang 已提交
256 257

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

Y
Yang Yang 已提交
261
        auto sum_op = framework::OpRegistry::CreateOp(
Y
Yu Yang 已提交
262
            "sum", {{"X", {s, tmp_name}}}, {{"Out", {s}}},
Y
Yang Yang 已提交
263
            framework::AttributeMap{});
Y
Yang Yang(Tony) 已提交
264
        VLOG(10) << sum_op->DebugStringEx(sub_scopes[0]);
265
        sum_op->Run(*sub_scopes[0], places[0]);
Y
Yang Yang 已提交
266
        WaitOnPlace(places[0]);
Y
Yang Yang 已提交
267 268
      }

Y
Yang Yang 已提交
269
      CopyOrShare(*sub_scopes[0]->FindVar(s), place, scope.FindVar(s));
Y
Yang Yang 已提交
270
    }
Y
Yang Yang 已提交
271
    WaitOnPlaces(places);
Y
Yang Yang 已提交
272
  }
Y
Yang Yang 已提交
273 274
};

Y
Yu Yang 已提交
275 276 277 278 279 280 281
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 已提交
282 283 284 285 286
class ParallelDoGradOpDescMaker : public framework::SingleGradOpDescMaker {
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
Y
Yang Yang 已提交
287 288
  virtual std::unique_ptr<framework::OpDesc> Apply() const {
    auto *grad = new framework::OpDesc();
Y
Yang Yang 已提交
289
    grad->SetType("parallel_do_grad");
Y
Yang Yang 已提交
290
    for (auto &input_param : this->InputNames()) {
Y
Yang Yang 已提交
291
      VLOG(3) << input_param;
Y
Yang Yang 已提交
292
      grad->SetInput(input_param, this->Input(input_param));
293 294 295 296
      if (input_param != kPlaces) {
        grad->SetOutput(framework::GradVarName(input_param),
                        this->InputGrad(input_param, false));
      }
Y
Yang Yang 已提交
297
    }
Y
Yang Yu 已提交
298 299 300 301 302 303 304 305 306 307 308
    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 已提交
309 310 311 312 313 314 315 316

    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 已提交
317 318 319
        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 已提交
320 321
            // there are some gradient operators who need the OG. So make this
            // OG as an input of parallel.do
Y
Yang Yu 已提交
322 323
            og_names.push_back(og_name);
          }
Y
Yang Yu 已提交
324 325
          // else, there is no operator who need the OG. Do not use this OG as
          // an input
Y
Yang Yu 已提交
326 327
        }
        grad->SetInput(framework::GradVarName(output_param), og_names);
Y
Yang Yang 已提交
328 329 330 331 332
      }
    }
    grad->SetAttrMap(this->Attrs());
    grad->SetBlockAttr(kParallelBlock, *grad_block_[0]);

Y
Yang Yang 已提交
333
    return std::unique_ptr<framework::OpDesc>(grad);
Y
Yang Yang 已提交
334 335 336 337 338 339
  }
};

class ParallelDoGradOpShapeInference : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *ctx) const override {
Y
Yu Yang 已提交
340
    PADDLE_ENFORCE(ctx->HasInputs(kParameters));
Y
Yang Yang 已提交
341
    PADDLE_ENFORCE(ctx->HasInputs(kInputs));
Y
Yang Yang(Tony) 已提交
342
    PADDLE_ENFORCE(ctx->HasInputs(kOutputs));
Y
Yu Yang 已提交
343 344 345 346 347 348 349 350 351 352 353 354 355 356

    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 已提交
357
    }
Y
Yu Yang 已提交
358

Y
Yang Yang(Tony) 已提交
359 360 361 362 363 364 365 366
    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 已提交
367
    }
Y
Yang Yang 已提交
368 369 370 371 372 373 374 375 376 377 378
  }
};

}  // 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);