operator.h 20.5 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.
Y
Yu Yang 已提交
145
#define DEFINE_OP_CLONE_METHOD(cls)                       \
Y
Yu Yang 已提交
146
  std::unique_ptr<OperatorBase> Clone() const final {     \
Y
Yu Yang 已提交
147
    return std::unique_ptr<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
 private:
293 294
  const OperatorBase& op_;
  const Scope& scope_;
295
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
296 297
};

298 299 300 301 302 303 304 305 306 307 308 309 310 311
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;

Q
tmp  
qiaolongfei 已提交
312 313
class CompileTimeInferShapeContext : public InferShapeContextBase {
 public:
Q
qiaolongfei 已提交
314 315
  CompileTimeInferShapeContext(const OpDescBind& op, const BlockDescBind& block)
      : op_(op), block_(block) {}
Q
tmp  
qiaolongfei 已提交
316

Q
qiaolongfei 已提交
317
  bool HasInput(const std::string& name) const override {
Q
qiaolongfei 已提交
318
    const std::vector<std::string>& input_names = op_.Input(name);
319 320 321 322 323
    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 已提交
324
    return block_.HasVar(input_names[0]);
Q
tmp  
qiaolongfei 已提交
325 326
  }

Q
qiaolongfei 已提交
327
  bool HasOutput(const std::string& name) const override {
Q
qiaolongfei 已提交
328
    const std::vector<std::string>& output_names = op_.Output(name);
329 330 331 332 333
    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 已提交
334
    return block_.HasVar(output_names[0]);
Q
tmp  
qiaolongfei 已提交
335 336
  }

Q
qiaolongfei 已提交
337
  bool HasInputs(const std::string& name) const override {
Q
qiaolongfei 已提交
338
    const std::vector<std::string>& input_names = op_.Input(name);
Q
qiaolongfei 已提交
339
    PADDLE_ENFORCE(!input_names.empty(), "Inputs(%s) length is 0", name);
Q
qiaolongfei 已提交
340 341
    for (auto& input : input_names) {
      if (!block_.HasVar(input)) return false;
Q
tmp  
qiaolongfei 已提交
342 343 344 345
    }
    return true;
  }

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

Q
qiaolongfei 已提交
355
  DDim GetInputDim(const std::string& name) const override {
Q
qiaolongfei 已提交
356
    std::vector<DDim> ddims = GetInputsDim(name);
357 358 359 360 361
    auto length = ddims.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Input(%s) should have 1 value, "
                      "but it has %d now",
                      name, length);
Q
qiaolongfei 已提交
362
    return ddims[0];
Q
tmp  
qiaolongfei 已提交
363 364
  }

Q
qiaolongfei 已提交
365
  void SetInputDim(const std::string& name, const DDim& dim) override {
Q
qiaolongfei 已提交
366
    SetInputsDim(name, {dim});
Q
tmp  
qiaolongfei 已提交
367 368
  }

Q
qiaolongfei 已提交
369
  DDim GetOutputDim(const std::string& name) const override {
Q
qiaolongfei 已提交
370
    std::vector<DDim> ddims = GetOutputsDim(name);
371 372 373 374 375
    auto length = ddims.size();
    PADDLE_ENFORCE_EQ(length, 1UL,
                      "Output(%s) should have 1 value, "
                      "but it has %d now",
                      name, length);
Q
qiaolongfei 已提交
376
    return ddims[0];
Q
tmp  
qiaolongfei 已提交
377 378
  }

Q
qiaolongfei 已提交
379
  void SetOutputDim(const std::string& name, const DDim& dim) override {
Q
qiaolongfei 已提交
380
    SetOutputsDim(name, {dim});
Q
tmp  
qiaolongfei 已提交
381 382
  }

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

Q
qiaolongfei 已提交
385 386
  const std::vector<std::string>& Inputs(
      const std::string& name) const override {
Q
qiaolongfei 已提交
387
    return op_.Input(name);
Q
tmp  
qiaolongfei 已提交
388 389
  }

Q
qiaolongfei 已提交
390 391
  const std::vector<std::string>& Outputs(
      const std::string& name) const override {
Q
qiaolongfei 已提交
392
    return op_.Output(name);
Q
tmp  
qiaolongfei 已提交
393 394 395
  }

 private:
Q
qiaolongfei 已提交
396
  DDim GetDim(const std::string& name) const override {
Q
qiaolongfei 已提交
397
    return framework::make_ddim(block_.Var(name)->Shape());
Q
tmp  
qiaolongfei 已提交
398 399
  }

Q
qiaolongfei 已提交
400
  void SetDim(const std::string& name, const DDim& dim) override {
Q
qiaolongfei 已提交
401
    block_.Var(name)->SetShape(framework::vectorize(dim));
Q
tmp  
qiaolongfei 已提交
402 403
  }

Q
qiaolongfei 已提交
404 405
  const OpDescBind& op_;
  const BlockDescBind& block_;
Q
tmp  
qiaolongfei 已提交
406 407
};

Q
Qiao Longfei 已提交
408 409 410 411 412
class RuntimeInferShapeContext : public InferShapeContextBase {
 public:
  RuntimeInferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}

Q
qiaolongfei 已提交
413
  bool HasInput(const std::string& name) const override {
Q
Qiao Longfei 已提交
414 415 416 417 418
    auto ipt = op_.Input(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

Q
qiaolongfei 已提交
419
  bool HasOutput(const std::string& name) const override {
Q
Qiao Longfei 已提交
420 421 422 423 424
    auto ipt = op_.Output(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

Q
qiaolongfei 已提交
425
  bool HasInputs(const std::string& name) const override {
426
    auto inputs = op_.Inputs(name);
Q
qiaolongfei 已提交
427
    if (inputs.empty()) {
428 429 430 431 432 433 434 435 436 437
      return false;
    }
    for (auto& input : inputs) {
      if (scope_.FindVar(input) == nullptr) {
        return false;
      }
    }
    return true;
  }

Q
qiaolongfei 已提交
438
  bool HasOutputs(const std::string& name) const override {
439
    auto outputs = op_.Outputs(name);
Q
qiaolongfei 已提交
440
    if (outputs.empty()) {
441 442 443 444 445 446 447 448 449 450
      return false;
    }
    for (auto& output : outputs) {
      if (scope_.FindVar(output) == nullptr) {
        return false;
      }
    }
    return true;
  }

Q
qiaolongfei 已提交
451
  DDim GetInputDim(const std::string& name) const override {
Q
Qiao Longfei 已提交
452 453 454
    return GetDim(op_.Input(name));
  }

Q
qiaolongfei 已提交
455
  void SetInputDim(const std::string& name, const DDim& dim) override {
Q
Qiao Longfei 已提交
456 457 458
    SetDim(op_.Input(name), dim);
  }

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

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

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

Q
qiaolongfei 已提交
469 470
  const std::vector<std::string>& Inputs(
      const std::string& name) const override {
Q
Qiao Longfei 已提交
471 472 473
    return op_.Inputs(name);
  }

Q
qiaolongfei 已提交
474 475
  const std::vector<std::string>& Outputs(
      const std::string& name) const override {
Q
Qiao Longfei 已提交
476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
    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 已提交
496
  DDim GetDim(const std::string& name) const override {
Q
Qiao Longfei 已提交
497 498 499
    return GetTensor<false>(name)->dims();
  }

Q
qiaolongfei 已提交
500
  void SetDim(const std::string& name, const DDim& dim) override {
Q
Qiao Longfei 已提交
501 502 503 504 505 506 507
    GetTensor<true>(name)->Resize(dim);
  }

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

Y
Yu Yang 已提交
508
class OpKernelBase {
Q
qijun 已提交
509
 public:
Q
qijun 已提交
510
  /**
511
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
512 513
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
514
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
515 516
   */

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

Y
Yu Yang 已提交
519 520 521 522 523 524 525
  virtual ~OpKernelBase() = default;
};

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

Q
Qiao Longfei 已提交
528 529
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
530 531
  struct OpKernelKey {
    platform::Place place_;
Y
Yu Yang 已提交
532
    DataType data_type_;
Q
Qiao Longfei 已提交
533

Y
Yu Yang 已提交
534 535 536 537 538
    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 已提交
539

Q
qijun 已提交
540
    bool operator==(const OpKernelKey& o) const {
Y
Yu Yang 已提交
541 542
      return platform::places_are_same_class(place_, o.place_) &&
             data_type_ == o.data_type_;
Q
qijun 已提交
543
    }
Y
Yu Yang 已提交
544 545 546
  };

  struct OpKernelHash {
Y
Yu Yang 已提交
547
    std::hash<int> hash_;
Y
Yu Yang 已提交
548
    size_t operator()(const OpKernelKey& key) const {
Y
Yu Yang 已提交
549 550
      int place = key.place_.which();
      int data_type = static_cast<int>(key.data_type_);
Y
Yu Yang 已提交
551 552
      int pre_hash = data_type << NUM_PLACE_TYPE_LIMIT_IN_BIT |
                     (place & ((1 << NUM_PLACE_TYPE_LIMIT_IN_BIT) - 1));
Y
Yu Yang 已提交
553
      return hash_(pre_hash);
Y
Yu Yang 已提交
554 555 556 557
    }
  };

  using OpKernelMap =
Y
Yu Yang 已提交
558 559
      std::unordered_map<OpKernelKey, std::unique_ptr<OpKernelBase>,
                         OpKernelHash>;
Q
Qiao Longfei 已提交
560

Y
Yu Yang 已提交
561 562
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
563 564
      : OperatorBase(type, inputs, outputs, attrs) {}

Y
Yu Yang 已提交
565
  void Run(const Scope& scope,
Y
Yu Yang 已提交
566
           const platform::DeviceContext& dev_ctx) const final {
Y
Yu Yang 已提交
567 568 569
    RuntimeInferShapeContext infer_shape_ctx(*this, scope);
    this->InferShape(&infer_shape_ctx);

Y
Yu Yang 已提交
570
    ExecutionContext ctx(*this, scope, dev_ctx);
571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589

    // 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 已提交
590 591
  }

Y
Yu Yang 已提交
592 593 594 595
  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 已提交
596
  }
Y
Yan Chunwei 已提交
597

598
  bool SupportGPU() const override {
Y
Yu Yang 已提交
599 600 601 602 603
    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_);
                       });
604 605
  }

Q
Qiao Longfei 已提交
606
  virtual void InferShape(InferShapeContextBase* ctx) const = 0;
Y
Yu Yang 已提交
607

Q
qiaolongfei 已提交
608
 protected:
Y
Yu Yang 已提交
609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635
  // 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 已提交
636 637
};

638 639 640
std::ostream& operator<<(std::ostream& os,
                         const OperatorWithKernel::OpKernelKey& kernel_key);

Q
Qiao Longfei 已提交
641 642
}  // namespace framework
}  // namespace paddle