operator.h 20.8 KB
Newer Older
Q
Qiao Longfei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

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

D
dongzhihong 已提交
17
#include <algorithm>
18
#include <atomic>
Q
Qiao Longfei 已提交
19 20 21 22
#include <string>
#include <unordered_map>
#include <vector>

Y
Yu Yang 已提交
23
#include "op_info.h"
Y
Yi Wang 已提交
24
#include "paddle/framework/attribute.h"
Q
qiaolongfei 已提交
25
#include "paddle/framework/block_desc.h"
Y
Yu Yang 已提交
26
#include "paddle/framework/data_type.h"
Y
Yu Yang 已提交
27
#include "paddle/framework/framework.pb.h"
28
#include "paddle/framework/lod_tensor.h"
Q
qijun 已提交
29
#include "paddle/framework/scope.h"
Q
Qiao Longfei 已提交
30
#include "paddle/framework/shape_inference.h"
Q
qijun 已提交
31 32 33
#include "paddle/framework/tensor.h"
#include "paddle/platform/device_context.h"
#include "paddle/platform/place.h"
Y
Yu Yang 已提交
34
#include "paddle/platform/variant.h"
Q
qijun 已提交
35
#include "paddle/utils/Error.h"
Q
Qiao Longfei 已提交
36 37 38 39

namespace paddle {
namespace framework {

40
/// If a variable is a empty variable, that name will be used.
41
constexpr char kEmptyVarName[] = "@EMPTY@";
42 43 44

/// If a variable is a temporary variable, that name will be set in Python,
/// but it will be convert to a unique name in scope after OpCreator.
45
constexpr char kTempVarName[] = "@TEMP@";
46 47 48 49

/// If a variable's name has a certain suffix, it means that the
/// variable is the gradient of another varibale.
/// e.g. Variable "x@GRAD" is the gradient of varibale "x".
50
constexpr char kGradVarSuffix[] = "@GRAD";
51 52

/// Variables with this suffix are supposed to be filled up with zeros.
53
constexpr char kZeroVarSuffix[] = "@ZERO";
54 55 56 57 58

inline std::string GradVarName(const std::string& var_name) {
  return var_name + kGradVarSuffix;
}

Q
Qiao Longfei 已提交
59
class OperatorBase;
60
class ExecutionContext;
61

Q
Qiao Longfei 已提交
62 63 64
extern const Tensor* GetTensorFromVar(const Variable* var);
extern Tensor* GetTensorFromVar(Variable* var);

Q
Qiao Longfei 已提交
65 66 67 68 69 70 71 72
/**
 * OperatorBase has the basic element that Net will call to do computation.
 * Only CreateOperator from OpRegistry will new Operator directly. User
 * should always construct a proto message OpDesc and call
 * OpRegistry::CreateOp(op_desc) to get an Operator instance.
 */
class OperatorBase {
 public:
Y
Yu Yang 已提交
73 74
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
75

Q
Qiao Longfei 已提交
76 77 78
  virtual ~OperatorBase() {}

  template <typename T>
Y
Yu Yang 已提交
79
  inline const T& Attr(const std::string& name) const {
Q
Qiao Longfei 已提交
80 81 82 83 84
    PADDLE_ENFORCE(attrs_.count(name) != 0, "%s should be in AttributeMap",
                   name);
    return boost::get<T>(attrs_.at(name));
  }

85
  virtual std::string DebugString() const;
Q
Qiao Longfei 已提交
86 87

  /// Net will call this function to Run an op.
Y
Yu Yang 已提交
88
  virtual void Run(const Scope& scope,
Y
Yu Yang 已提交
89 90
                   const platform::DeviceContext& dev_ctx) const = 0;

Y
Yu Yang 已提交
91 92
  virtual bool IsNetOp() const { return false; }

93 94
  virtual bool SupportGPU() const { return false; }

D
dongzhihong 已提交
95 96 97
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

Y
Yu Yang 已提交
98 99
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
100

Y
Yu Yang 已提交
101
  //! Get a input with argument's name described in `op_proto`
102
  std::string Input(const std::string& name) const;
Y
Yu Yang 已提交
103
  //! Get a input which has multiple variables.
Y
Yu Yang 已提交
104
  const std::vector<std::string>& Inputs(const std::string& name) const;
Y
Yi Wang 已提交
105

Q
qijun 已提交
106 107
  std::vector<std::string> InputVars() const;

Y
Yu Yang 已提交
108
  //! Get a output with argument's name described in `op_proto`
109
  std::string Output(const std::string& name) const;
Y
Yu Yang 已提交
110 111
  //! Get an output which has multiple variables.
  //! TODO add a vector_view to prevent memory copy.
Y
Yu Yang 已提交
112
  const std::vector<std::string>& Outputs(const std::string& name) const;
Y
Yan Chunwei 已提交
113

Y
Yu Yang 已提交
114
  virtual std::vector<std::string> OutputVars(bool has_intermediate) const;
115

Q
qiaolongfei 已提交
116
  const std::string& Type() const { return type_; }
Q
qiaolongfei 已提交
117
  void SetType(const std::string& type) { type_ = type; }
Y
Yi Wang 已提交
118 119
  const AttributeMap& Attrs() const { return attrs_; }

Y
Yu Yang 已提交
120
  // Return a new operator instance, which is as same as this.
Y
Yu Yang 已提交
121 122
  // Use unique_ptr to prevent caller forget to delete this pointer.
  virtual std::unique_ptr<OperatorBase> Clone() const = 0;
Y
Yu Yang 已提交
123

Q
qiaolongfei 已提交
124
 protected:
Q
Qiao Longfei 已提交
125
  std::string type_;
D
dongzhihong 已提交
126
  // NOTE: in case of OpGrad, inputs_ contains:
Y
Yu Yang 已提交
127
  // I (Inputs)opear
D
dongzhihong 已提交
128 129
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
130
  VariableNameMap inputs_;
Y
Yu Yang 已提交
131

D
dongzhihong 已提交
132 133
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
134
  VariableNameMap outputs_;
Q
Qiao Longfei 已提交
135
  AttributeMap attrs_;
136 137 138 139

 private:
  void GenerateTemporaryNames();
  void CheckAllInputOutputSet() const;
Y
Yan Chunwei 已提交
140 141
};

Y
Yu Yang 已提交
142 143
// Macro for define a clone method.
// If you are writing an kernel operator, `Clone` will be defined when you
144
// register it. i.e. `Clone` method is not needed to define by yourself.
145 146 147
#define DEFINE_OP_CLONE_METHOD(cls)                                            \
  std::unique_ptr<::paddle::framework::OperatorBase> Clone() const final {     \
    return std::unique_ptr<::paddle::framework::OperatorBase>(new cls(*this)); \
Y
Yu Yang 已提交
148
  }
Y
Yu Yang 已提交
149

Y
Yu Yang 已提交
150 151 152 153
// Macro for define a default constructor for Operator.
// You can also use
//   using PARENT_CLASS::PARENT_CLASS;
// to use parent's constructor.
Y
Yu Yang 已提交
154 155
#define DEFINE_OP_CONSTRUCTOR(cls, parent_cls)             \
  cls(const std::string& type,                             \
Y
Yu Yang 已提交
156 157 158
      const ::paddle::framework::VariableNameMap& inputs,  \
      const ::paddle::framework::VariableNameMap& outputs, \
      const paddle::framework::AttributeMap& attrs)        \
Y
Yu Yang 已提交
159
      : parent_cls(type, inputs, outputs, attrs) {}
Y
Yu Yang 已提交
160

161 162
class NOP : public OperatorBase {
 public:
163
  using OperatorBase::OperatorBase;
164 165
  void Run(const Scope& scope,
           const platform::DeviceContext& dev_ctx) const override {}
166 167 168
  std::unique_ptr<OperatorBase> Clone() const override {
    return std::unique_ptr<OperatorBase>(new NOP(*this));
  }
169 170
};

171
class ExecutionContext {
Y
Yan Chunwei 已提交
172
 public:
173 174 175
  ExecutionContext(const OperatorBase& op, const Scope& scope,
                   const platform::DeviceContext& device_context)
      : op_(op), scope_(scope), device_context_(device_context) {}
176

Q
qiaolongfei 已提交
177 178 179 180
  const OperatorBase& op() const { return op_; }

  const Scope& scope() const { return scope_; }

Q
qiaolongfei 已提交
181
  template <typename T>
Y
Yu Yang 已提交
182 183
  inline const T& Attr(const std::string& name) const {
    return op_.Attr<T>(name);
Q
qiaolongfei 已提交
184 185
  }

Y
Yu Yang 已提交
186
  size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
187
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
188 189
  }

Y
Yu Yang 已提交
190
  size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
191
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
192 193
  }

194
  const Variable* InputVar(const std::string& name) const {
195
    auto ipt = op_.Input(name);
Y
Yu Yang 已提交
196
    return ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
Y
Yan Chunwei 已提交
197 198
  }

199
  Variable* OutputVar(const std::string& name) const {
200
    auto opt = op_.Output(name);
Y
Yu Yang 已提交
201
    return opt == kEmptyVarName ? nullptr : scope_.FindVar(opt);
Y
Yan Chunwei 已提交
202 203
  }

204 205
  const std::vector<const Variable*> MultiInputVar(
      const std::string& name) const {
Y
Yan Chunwei 已提交
206 207
    auto names = op_.Inputs(name);
    std::vector<const Variable*> res;
208
    res.reserve(names.size());
209 210
    std::transform(names.begin(), names.end(), std::back_inserter(res),
                   [this](const std::string& name) {
Y
Yu Yang 已提交
211 212
                     return name == kEmptyVarName ? nullptr
                                                  : scope_.FindVar(name);
213
                   });
Y
Yan Chunwei 已提交
214 215 216
    return res;
  }

217
  std::vector<Variable*> MultiOutputVar(const std::string& name) const {
Y
Yan Chunwei 已提交
218
    auto names = op_.Outputs(name);
219
    std::vector<Variable*> res;
220
    res.reserve(names.size());
221 222
    std::transform(names.begin(), names.end(), std::back_inserter(res),
                   [this](const std::string& name) {
Y
Yu Yang 已提交
223 224
                     return name == kEmptyVarName ? nullptr
                                                  : scope_.FindVar(name);
225
                   });
Y
Yan Chunwei 已提交
226 227 228
    return res;
  }

229 230
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
231
    auto* var = InputVar(name);
232
    return var == nullptr ? nullptr : &var->Get<T>();
233 234 235 236
  }

  template <typename T>
  T* Output(const std::string& name) const {
237
    auto var = OutputVar(name);
238
    return var == nullptr ? nullptr : var->GetMutable<T>();
239 240 241 242 243 244 245 246
  }

  template <typename T>
  const std::vector<const T*> MultiInput(const std::string& name) const {
    auto names = op_.Inputs(name);
    std::vector<const T*> res;
    res.reserve(names.size());
    std::transform(names.begin(), names.end(), std::back_inserter(res),
247
                   [&](const std::string& sub_name) {
248
                     auto var = scope_.FindVar(sub_name);
249
                     return var == nullptr ? nullptr : &var->Get<T>();
250 251 252 253 254
                   });
    return res;
  }

  template <typename T>
255
  std::vector<T*> MultiOutput(const std::string& name) const {
256
    auto names = op_.Outputs(name);
257
    std::vector<T*> res;
258 259
    res.reserve(names.size());
    std::transform(names.begin(), names.end(), std::back_inserter(res),
260
                   [&](const std::string& sub_name) {
261
                     auto var = scope_.FindVar(sub_name);
262
                     return var == nullptr ? nullptr : var->GetMutable<T>();
263 264 265 266
                   });
    return res;
  }

267 268 269 270 271 272
  void ShareLoD(const std::string& in, const std::string& out, size_t i = 0,
                size_t j = 0) const {
    PADDLE_ENFORCE_LT(i, InputSize(in));
    PADDLE_ENFORCE_LT(j, OutputSize(out));
    auto* in_var = MultiInputVar(in)[i];
    auto* out_var = MultiOutputVar(out)[j];
273
    if (!in_var->IsType<LoDTensor>()) return;
274
    PADDLE_ENFORCE(out_var->IsType<LoDTensor>(),
275
                   "The %d-th output of Output(%s) must be LoDTensor.", j, out);
276 277 278
    auto in_tensor = in_var->Get<LoDTensor>();
    auto* out_tensor = out_var->GetMutable<LoDTensor>();
    out_tensor->set_lod(in_tensor.lod());
279 280
  }

Q
qijun 已提交
281
  template <typename PlaceType,
282 283
            typename DeviceType = typename platform::EigenDeviceConverter<
                PlaceType>::EigenDeviceType>
284
  DeviceType& GetEigenDevice() const;
Q
qijun 已提交
285

286
  platform::Place GetPlace() const { return device_context_.GetPlace(); }
Q
qijun 已提交
287

288
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
289
    return device_context_;
Q
qijun 已提交
290
  }
Q
qijun 已提交
291

武毅 已提交
292 293 294 295 296 297 298 299 300
#ifdef PADDLE_WITH_CUDA
  const platform::CUDADeviceContext& cuda_device_context() const {
    PADDLE_ENFORCE(platform::is_gpu_place(device_context_.GetPlace()));
    auto cuda_ctx =
        reinterpret_cast<const platform::CUDADeviceContext*>(&device_context_);
    return *cuda_ctx;
  }
#endif

301
 private:
302 303
  const OperatorBase& op_;
  const Scope& scope_;
304
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
305 306
};

307 308 309 310 311 312 313 314 315 316 317 318 319 320
template <>
const Tensor* ExecutionContext::Input<Tensor>(const std::string& name) const;

template <>
const std::vector<const Tensor*> ExecutionContext::MultiInput<Tensor>(
    const std::string& name) const;

template <>
Tensor* ExecutionContext::Output<Tensor>(const std::string& name) const;

template <>
std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
    const std::string& name) const;

321
class CompileTimeInferShapeContext : public InferShapeContext {
Q
tmp  
qiaolongfei 已提交
322
 public:
Q
qiaolongfei 已提交
323 324
  CompileTimeInferShapeContext(const OpDescBind& op, const BlockDescBind& block)
      : op_(op), block_(block) {}
Q
tmp  
qiaolongfei 已提交
325

Q
qiaolongfei 已提交
326
  bool HasInput(const std::string& name) const override {
Q
qiaolongfei 已提交
327
    const std::vector<std::string>& input_names = op_.Input(name);
328 329 330 331 332
    auto length = input_names.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Input(%s) should have only one value, "
                      "but it have %d now",
                      name, length);
Q
qiaolongfei 已提交
333
    return block_.HasVar(input_names[0]);
Q
tmp  
qiaolongfei 已提交
334 335
  }

Q
qiaolongfei 已提交
336
  bool HasOutput(const std::string& name) const override {
Q
qiaolongfei 已提交
337
    const std::vector<std::string>& output_names = op_.Output(name);
338 339 340 341 342
    auto length = output_names.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Output(%s) should have only one value, "
                      "but it have %d now",
                      name, length);
Q
qiaolongfei 已提交
343
    return block_.HasVar(output_names[0]);
Q
tmp  
qiaolongfei 已提交
344 345
  }

Q
qiaolongfei 已提交
346
  bool HasInputs(const std::string& name) const override {
Q
qiaolongfei 已提交
347
    const std::vector<std::string>& input_names = op_.Input(name);
Q
qiaolongfei 已提交
348
    PADDLE_ENFORCE(!input_names.empty(), "Inputs(%s) length is 0", name);
Q
qiaolongfei 已提交
349 350
    for (auto& input : input_names) {
      if (!block_.HasVar(input)) return false;
Q
tmp  
qiaolongfei 已提交
351 352 353 354
    }
    return true;
  }

Q
qiaolongfei 已提交
355
  bool HasOutputs(const std::string& name) const override {
Q
qiaolongfei 已提交
356
    const std::vector<std::string>& output_names = op_.Output(name);
Q
qiaolongfei 已提交
357
    PADDLE_ENFORCE(!output_names.empty(), "Inputs(%s) length is 0", name);
Q
qiaolongfei 已提交
358
    for (auto& output : output_names) {
Q
qiaolongfei 已提交
359
      if (!block_.HasVar(output)) return false;
Q
tmp  
qiaolongfei 已提交
360 361 362 363
    }
    return true;
  }

Q
qiaolongfei 已提交
364
  DDim GetInputDim(const std::string& name) const override {
Q
qiaolongfei 已提交
365
    std::vector<DDim> ddims = GetInputsDim(name);
366 367 368 369 370
    auto length = ddims.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Input(%s) should have 1 value, "
                      "but it has %d now",
                      name, length);
Q
qiaolongfei 已提交
371
    return ddims[0];
Q
tmp  
qiaolongfei 已提交
372 373
  }

Q
qiaolongfei 已提交
374
  void SetInputDim(const std::string& name, const DDim& dim) override {
Q
qiaolongfei 已提交
375
    SetInputsDim(name, {dim});
Q
tmp  
qiaolongfei 已提交
376 377
  }

Q
qiaolongfei 已提交
378
  DDim GetOutputDim(const std::string& name) const override {
Q
qiaolongfei 已提交
379
    std::vector<DDim> ddims = GetOutputsDim(name);
380 381 382 383 384
    auto length = ddims.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Output(%s) should have 1 value, "
                      "but it has %d now",
                      name, length);
Q
qiaolongfei 已提交
385
    return ddims[0];
Q
tmp  
qiaolongfei 已提交
386 387
  }

Q
qiaolongfei 已提交
388
  void SetOutputDim(const std::string& name, const DDim& dim) override {
Q
qiaolongfei 已提交
389
    SetOutputsDim(name, {dim});
Q
tmp  
qiaolongfei 已提交
390 391
  }

Q
qiaolongfei 已提交
392
  AttrReader Attrs() const override { return AttrReader(op_.GetAttrMap()); }
Q
tmp  
qiaolongfei 已提交
393

Q
qiaolongfei 已提交
394 395
  const std::vector<std::string>& Inputs(
      const std::string& name) const override {
Q
qiaolongfei 已提交
396
    return op_.Input(name);
Q
tmp  
qiaolongfei 已提交
397 398
  }

Q
qiaolongfei 已提交
399 400
  const std::vector<std::string>& Outputs(
      const std::string& name) const override {
Q
qiaolongfei 已提交
401
    return op_.Output(name);
Q
tmp  
qiaolongfei 已提交
402 403 404
  }

 private:
Q
qiaolongfei 已提交
405
  DDim GetDim(const std::string& name) const override {
D
Dong Zhihong 已提交
406
    return framework::make_ddim(block_.FindVar(name)->Shape());
Q
tmp  
qiaolongfei 已提交
407 408
  }

Q
qiaolongfei 已提交
409
  void SetDim(const std::string& name, const DDim& dim) override {
D
Dong Zhihong 已提交
410
    block_.FindVar(name)->SetShape(framework::vectorize(dim));
Q
tmp  
qiaolongfei 已提交
411 412
  }

Q
qiaolongfei 已提交
413 414
  const OpDescBind& op_;
  const BlockDescBind& block_;
Q
tmp  
qiaolongfei 已提交
415 416
};

417
class RuntimeInferShapeContext : public InferShapeContext {
Q
Qiao Longfei 已提交
418 419 420 421
 public:
  RuntimeInferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}

Q
qiaolongfei 已提交
422
  bool HasInput(const std::string& name) const override {
Q
Qiao Longfei 已提交
423 424 425 426 427
    auto ipt = op_.Input(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

Q
qiaolongfei 已提交
428
  bool HasOutput(const std::string& name) const override {
Q
Qiao Longfei 已提交
429 430 431 432 433
    auto ipt = op_.Output(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

Q
qiaolongfei 已提交
434
  bool HasInputs(const std::string& name) const override {
435
    auto inputs = op_.Inputs(name);
Q
qiaolongfei 已提交
436
    if (inputs.empty()) {
437 438 439 440 441 442 443 444 445 446
      return false;
    }
    for (auto& input : inputs) {
      if (scope_.FindVar(input) == nullptr) {
        return false;
      }
    }
    return true;
  }

Q
qiaolongfei 已提交
447
  bool HasOutputs(const std::string& name) const override {
448
    auto outputs = op_.Outputs(name);
Q
qiaolongfei 已提交
449
    if (outputs.empty()) {
450 451 452 453 454 455 456 457 458 459
      return false;
    }
    for (auto& output : outputs) {
      if (scope_.FindVar(output) == nullptr) {
        return false;
      }
    }
    return true;
  }

Q
qiaolongfei 已提交
460
  DDim GetInputDim(const std::string& name) const override {
Q
Qiao Longfei 已提交
461 462 463
    return GetDim(op_.Input(name));
  }

Q
qiaolongfei 已提交
464
  void SetInputDim(const std::string& name, const DDim& dim) override {
Q
Qiao Longfei 已提交
465 466 467
    SetDim(op_.Input(name), dim);
  }

Q
qiaolongfei 已提交
468
  DDim GetOutputDim(const std::string& name) const override {
Q
Qiao Longfei 已提交
469 470 471
    return GetDim(op_.Output(name));
  }

Q
qiaolongfei 已提交
472
  void SetOutputDim(const std::string& name, const DDim& dim) override {
Q
Qiao Longfei 已提交
473 474 475
    SetDim(op_.Output(name), dim);
  }

Q
qiaolongfei 已提交
476
  AttrReader Attrs() const override { return AttrReader(op_.Attrs()); }
Q
Qiao Longfei 已提交
477

Q
qiaolongfei 已提交
478 479
  const std::vector<std::string>& Inputs(
      const std::string& name) const override {
Q
Qiao Longfei 已提交
480 481 482
    return op_.Inputs(name);
  }

Q
qiaolongfei 已提交
483 484
  const std::vector<std::string>& Outputs(
      const std::string& name) const override {
Q
Qiao Longfei 已提交
485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504
    return op_.Outputs(name);
  }

 private:
  template <bool Allocate>
  Tensor* GetTensor(const std::string& name) const {
    Tensor* t = nullptr;
    auto* var = scope_.FindVar(name);
    if (!var->IsType<LoDTensor>() && !var->IsType<Tensor>()) {
      if (Allocate) {
        t = var->GetMutable<LoDTensor>();
      } else {
        PADDLE_THROW("Variable(%s) should be tensor", name);
      }
    } else {
      t = GetTensorFromVar(scope_.FindVar(name));
    }
    return t;
  }

Q
qiaolongfei 已提交
505
  DDim GetDim(const std::string& name) const override {
Q
Qiao Longfei 已提交
506 507 508
    return GetTensor<false>(name)->dims();
  }

Q
qiaolongfei 已提交
509
  void SetDim(const std::string& name, const DDim& dim) override {
Q
Qiao Longfei 已提交
510 511 512 513 514 515 516
    GetTensor<true>(name)->Resize(dim);
  }

  const OperatorBase& op_;
  const Scope& scope_;
};

Y
Yu Yang 已提交
517
class OpKernelBase {
Q
qijun 已提交
518
 public:
Q
qijun 已提交
519
  /**
520
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
521 522
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
523
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
524 525
   */

526
  virtual void Compute(const ExecutionContext& context) const = 0;
Y
Yu Yang 已提交
527

Y
Yu Yang 已提交
528 529 530 531 532 533 534
  virtual ~OpKernelBase() = default;
};

template <typename T>
class OpKernel : public OpKernelBase {
 public:
  using ELEMENT_TYPE = T;
Y
Yu Yang 已提交
535 536
};

Q
Qiao Longfei 已提交
537 538
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
539 540
  struct OpKernelKey {
    platform::Place place_;
Y
Yu Yang 已提交
541
    DataType data_type_;
Q
Qiao Longfei 已提交
542

Y
Yu Yang 已提交
543 544 545 546 547
    OpKernelKey(DataType data_type, platform::Place place)
        : place_(place), data_type_(data_type) {}

    OpKernelKey(DataType data_type, const platform::DeviceContext& dev_ctx)
        : place_(dev_ctx.GetPlace()), data_type_(data_type) {}
Y
Yu Yang 已提交
548

Q
qijun 已提交
549
    bool operator==(const OpKernelKey& o) const {
Y
Yu Yang 已提交
550 551
      return platform::places_are_same_class(place_, o.place_) &&
             data_type_ == o.data_type_;
Q
qijun 已提交
552
    }
Y
Yu Yang 已提交
553 554 555
  };

  struct OpKernelHash {
Y
Yu Yang 已提交
556
    std::hash<int> hash_;
Y
Yu Yang 已提交
557
    size_t operator()(const OpKernelKey& key) const {
Y
Yu Yang 已提交
558 559
      int place = key.place_.which();
      int data_type = static_cast<int>(key.data_type_);
Y
Yu Yang 已提交
560 561
      int pre_hash = data_type << NUM_PLACE_TYPE_LIMIT_IN_BIT |
                     (place & ((1 << NUM_PLACE_TYPE_LIMIT_IN_BIT) - 1));
Y
Yu Yang 已提交
562
      return hash_(pre_hash);
Y
Yu Yang 已提交
563 564 565 566
    }
  };

  using OpKernelMap =
Y
Yu Yang 已提交
567 568
      std::unordered_map<OpKernelKey, std::unique_ptr<OpKernelBase>,
                         OpKernelHash>;
Q
Qiao Longfei 已提交
569

Y
Yu Yang 已提交
570 571
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
572 573
      : OperatorBase(type, inputs, outputs, attrs) {}

Y
Yu Yang 已提交
574
  void Run(const Scope& scope,
Y
Yu Yang 已提交
575
           const platform::DeviceContext& dev_ctx) const final {
Y
Yu Yang 已提交
576 577 578
    RuntimeInferShapeContext infer_shape_ctx(*this, scope);
    this->InferShape(&infer_shape_ctx);

Y
Yu Yang 已提交
579
    ExecutionContext ctx(*this, scope, dev_ctx);
580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598

    // check if op[type] has kernel registered.
    auto& all_op_kernels = AllOpKernels();
    auto kernels_iter = all_op_kernels.find(type_);
    if (kernels_iter == all_op_kernels.end()) {
      PADDLE_THROW("op[%s] has no kernel", type_);
    }

    // check if op[type] have kernel for kernel_key
    OpKernelMap& kernels = kernels_iter->second;
    auto kernel_key = OpKernelKey(IndicateDataType(ctx), dev_ctx);
    auto kernel_iter = kernels.find(kernel_key);

    if (kernel_iter == kernels.end()) {
      PADDLE_THROW("op[%s] has no kernel with kernel_key[%s]", type_,
                   kernel_key);
    }

    kernel_iter->second->Compute(ctx);
Q
Qiao Longfei 已提交
599 600
  }

Y
Yu Yang 已提交
601 602 603 604
  static std::unordered_map<std::string /* op_type */, OpKernelMap>&
  AllOpKernels() {
    static std::unordered_map<std::string, OpKernelMap> g_all_op_kernels;
    return g_all_op_kernels;
Y
Yu Yang 已提交
605
  }
Y
Yan Chunwei 已提交
606

607
  bool SupportGPU() const override {
Y
Yu Yang 已提交
608 609 610 611 612
    auto& op_kernels = OperatorWithKernel::AllOpKernels().at(type_);
    return std::any_of(op_kernels.begin(), op_kernels.end(),
                       [](OpKernelMap::const_reference kern_pair) {
                         return platform::is_gpu_place(kern_pair.first.place_);
                       });
613 614
  }

615
  virtual void InferShape(InferShapeContext* ctx) const = 0;
Y
Yu Yang 已提交
616

Q
qiaolongfei 已提交
617
 protected:
Y
Yu Yang 已提交
618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644
  // indicate kernel DataType by input data. Defaultly all input data must be
  // same.
  virtual DataType IndicateDataType(const ExecutionContext& ctx) const {
    auto& scope = ctx.scope();
    int data_type = -1;
    for (auto& input : this->inputs_) {
      for (auto& ipt_name : input.second) {
        auto* var = scope.FindVar(ipt_name);
        if (var != nullptr) {
          const Tensor* t = nullptr;
          if (var->IsType<Tensor>()) {
            t = &var->Get<Tensor>();
          } else if (var->IsType<LoDTensor>()) {
            t = &var->Get<LoDTensor>();
          }
          if (t != nullptr) {
            int tmp = static_cast<int>(ToDataType(t->type()));
            PADDLE_ENFORCE(tmp == data_type || data_type == -1,
                           "DataType of Paddle Op must be same.");
            data_type = tmp;
          }
        }
      }
    }
    PADDLE_ENFORCE(data_type != -1, "DataType should be indicated by input");
    return static_cast<DataType>(data_type);
  }
Q
Qiao Longfei 已提交
645 646
};

647 648 649
std::ostream& operator<<(std::ostream& os,
                         const OperatorWithKernel::OpKernelKey& kernel_key);

Q
Qiao Longfei 已提交
650 651
}  // namespace framework
}  // namespace paddle