operator.h 13.3 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>
Q
Qiao Longfei 已提交
18 19 20 21
#include <string>
#include <unordered_map>
#include <vector>

Y
Yu Yang 已提交
22
#include "op_info.h"
Y
Yi Wang 已提交
23
#include "paddle/framework/attribute.h"
Y
Yu Yang 已提交
24
#include "paddle/framework/framework.pb.h"
Q
qijun 已提交
25 26 27 28
#include "paddle/framework/scope.h"
#include "paddle/framework/tensor.h"
#include "paddle/platform/device_context.h"
#include "paddle/platform/place.h"
Y
Yu Yang 已提交
29
#include "paddle/platform/variant.h"
Q
qijun 已提交
30
#include "paddle/utils/Error.h"
Q
Qiao Longfei 已提交
31 32 33 34

namespace paddle {
namespace framework {

35
/// If a variable is a empty variable, that name will be used.
36
constexpr char kEmptyVarName[] = "@EMPTY@";
37 38 39

/// 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.
40
constexpr char kTempVarName[] = "@TEMP@";
41 42 43 44

/// 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".
45
constexpr char kGradVarSuffix[] = "@GRAD";
46 47

/// Variables with this suffix are supposed to be filled up with zeros.
48
constexpr char kZeroVarSuffix[] = "@ZERO";
49 50 51 52 53

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

Q
Qiao Longfei 已提交
54
class OperatorBase;
55 56
class InferShapeContext;
class ExecutionContext;
57

Q
Qiao Longfei 已提交
58 59 60 61 62 63 64 65
/**
 * 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 已提交
66 67
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
68

Q
Qiao Longfei 已提交
69 70 71 72 73 74 75 76 77
  virtual ~OperatorBase() {}

  template <typename T>
  inline const T& GetAttr(const std::string& name) const {
    PADDLE_ENFORCE(attrs_.count(name) != 0, "%s should be in AttributeMap",
                   name);
    return boost::get<T>(attrs_.at(name));
  }

78
  virtual std::string DebugString() const;
Q
Qiao Longfei 已提交
79 80 81

  /// InferShape infer the size of Variables used by this Operator with
  /// information inside scope
Y
Yu Yang 已提交
82
  virtual void InferShape(const Scope& scope) const = 0;
Q
Qiao Longfei 已提交
83 84

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

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

90 91
  virtual bool SupportGPU() const { return false; }

D
dongzhihong 已提交
92 93 94
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

Y
Yu Yang 已提交
95 96
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
Y
Yu Yang 已提交
97
  //! Get a input with argument's name described in `op_proto`
Y
Yan Chunwei 已提交
98
  const std::string& Input(const std::string& name) const;
Y
Yu Yang 已提交
99
  //! Get a input which has multiple variables.
Y
Yu Yang 已提交
100
  const std::vector<std::string>& Inputs(const std::string& name) const;
Y
Yi Wang 已提交
101

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

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

Q
qiaolongfei 已提交
110
  const std::string& Type() const { return type_; }
Q
qiaolongfei 已提交
111
  void SetType(const std::string& type) { type_ = type; }
Y
Yi Wang 已提交
112 113
  const AttributeMap& Attrs() const { return attrs_; }

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

Q
qiaolongfei 已提交
118
 protected:
Q
Qiao Longfei 已提交
119
  std::string type_;
D
dongzhihong 已提交
120
  // NOTE: in case of OpGrad, inputs_ contains:
Y
Yu Yang 已提交
121
  // I (Inputs)opear
D
dongzhihong 已提交
122 123
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
124
  VariableNameMap inputs_;
Y
Yu Yang 已提交
125

D
dongzhihong 已提交
126 127
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
128
  VariableNameMap outputs_;
Q
Qiao Longfei 已提交
129
  AttributeMap attrs_;
Y
Yan Chunwei 已提交
130 131
};

Y
Yu Yang 已提交
132 133
// Macro for define a clone method.
// If you are writing an kernel operator, `Clone` will be defined when you
134
// register it. i.e. `Clone` method is not needed to define by yourself.
Y
Yu Yang 已提交
135
#define DEFINE_OP_CLONE_METHOD(cls)                       \
Y
Yu Yang 已提交
136
  std::unique_ptr<OperatorBase> Clone() const final {     \
Y
Yu Yang 已提交
137
    return std::unique_ptr<OperatorBase>(new cls(*this)); \
Y
Yu Yang 已提交
138
  }
Y
Yu Yang 已提交
139

Y
Yu Yang 已提交
140 141 142 143
// 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 已提交
144 145
#define DEFINE_OP_CONSTRUCTOR(cls, parent_cls)             \
  cls(const std::string& type,                             \
Y
Yu Yang 已提交
146 147 148
      const ::paddle::framework::VariableNameMap& inputs,  \
      const ::paddle::framework::VariableNameMap& outputs, \
      const paddle::framework::AttributeMap& attrs)        \
Y
Yu Yang 已提交
149
      : parent_cls(type, inputs, outputs, attrs) {}
Y
Yu Yang 已提交
150

151 152
class NOP : public OperatorBase {
 public:
153
  using OperatorBase::OperatorBase;
154 155 156
  void InferShape(const Scope& scope) const override {}
  void Run(const Scope& scope,
           const platform::DeviceContext& dev_ctx) const override {}
157 158 159
  std::unique_ptr<OperatorBase> Clone() const override {
    return std::unique_ptr<OperatorBase>(new NOP(*this));
  }
160 161
};

162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187
// this class not only make proto but also init attribute checkers.
class OpProtoAndCheckerMaker {
 public:
  OpProtoAndCheckerMaker(OpProto* proto, OpAttrChecker* op_checker)
      : proto_(proto), op_checker_(op_checker) {}

  ~OpProtoAndCheckerMaker() {
    PADDLE_ENFORCE(validated_, "should call Validate after build");
  }

  void Validate();

 protected:
  struct VariableBuilder {
    OpProto::Var* var_;

    VariableBuilder& AsDuplicable() {
      var_->set_duplicable(true);
      return *this;
    }

    VariableBuilder& AsIntermediate() {
      var_->set_intermediate(true);
      return *this;
    }

188 189
    VariableBuilder& NotInGradient() {
      var_->set_not_in_gradient(true);
190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226
      return *this;
    }
  };

  VariableBuilder AddInput(const std::string& name, const std::string& comment);

  VariableBuilder AddOutput(const std::string& name,
                            const std::string& comment);

  template <typename T>
  TypedAttrChecker<T>& AddAttr(const std::string& name,
                               const std::string& comment,
                               bool generated = false) {
    auto* attr = proto_->add_attrs();
    attr->set_name(name);
    attr->set_comment(comment);
    attr->set_generated(generated);
    attr->set_type(AttrTypeID<T>());
    return op_checker_->AddAttrChecker<T>(name);
  }

  void AddComment(const std::string& comment) { proto_->set_comment(comment); }

 private:
  void CheckNoDuplicatedInOutAttrs();

  OpProto* proto_;
  OpAttrChecker* op_checker_;
  bool validated_{false};
};

class NOPMaker : public OpProtoAndCheckerMaker {
 public:
  NOPMaker(framework::OpProto* proto, framework::OpAttrChecker* op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {}
};

227
class InferShapeContext {
Y
Yan Chunwei 已提交
228
 public:
229 230
  InferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}
231

Y
Yu Yang 已提交
232
  size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
233
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
234 235
  }

Y
Yu Yang 已提交
236
  size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
237
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
238 239
  }

240
  const Variable* InputVar(const std::string& name) const {
Y
Yu Yang 已提交
241
    return scope_.FindVar(op_.Input(name));
Y
Yan Chunwei 已提交
242 243
  }

244
  Variable* OutputVar(const std::string& name) const {
Y
Yu Yang 已提交
245
    return scope_.FindVar(op_.Output(name));
Y
Yan Chunwei 已提交
246 247
  }

248 249
  const std::vector<const Variable*> MultiInputVar(
      const std::string& name) const {
Y
Yan Chunwei 已提交
250 251
    auto names = op_.Inputs(name);
    std::vector<const Variable*> res;
252
    res.reserve(names.size());
Y
Yan Chunwei 已提交
253
    std::transform(
254
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
255
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
256 257 258
    return res;
  }

259
  std::vector<const Variable*> MultiOutputVar(const std::string& name) const {
Y
Yan Chunwei 已提交
260 261
    auto names = op_.Outputs(name);
    std::vector<const Variable*> res;
262
    res.reserve(names.size());
Y
Yan Chunwei 已提交
263
    std::transform(
264
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
265
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
266 267 268
    return res;
  }

269 270
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
271
    auto* var = InputVar(name);
Y
Yan Chunwei 已提交
272
    PADDLE_ENFORCE_NOT_NULL(var, "Input(%s) should not be nullptr", name);
273
    return &var->Get<T>();
274 275 276 277
  }

  template <typename T>
  T* Output(const std::string& name) const {
278
    auto var = OutputVar(name);
Y
Yan Chunwei 已提交
279
    PADDLE_ENFORCE_NOT_NULL(var, "Output(%s) should not be nullptr", name);
280
    return var->GetMutable<T>();
281 282 283 284 285 286 287 288
  }

  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),
289
                   [&](const std::string& sub_name) {
290
                     auto var = scope_.FindVar(sub_name);
Y
Yan Chunwei 已提交
291 292 293
                     PADDLE_ENFORCE_NOT_NULL(
                         var, "MultiInput(%s:%s) should not be nullptr", name,
                         sub_name);
294
                     return &var->Get<T>();
295 296 297 298 299 300 301 302 303 304
                   });
    return res;
  }

  template <typename T>
  std::vector<const T*> MultiOutput(const std::string& name) const {
    auto names = op_.Outputs(name);
    std::vector<const T*> res;
    res.reserve(names.size());
    std::transform(names.begin(), names.end(), std::back_inserter(res),
305
                   [&](const std::string& sub_name) {
306
                     auto var = scope_.FindVar(sub_name);
Y
Yan Chunwei 已提交
307
                     PADDLE_ENFORCE_NOT_NULL(
F
fengjiayi 已提交
308
                         var, "MultiOutput(%s:%s) should not be nullptr.", name,
Y
Yan Chunwei 已提交
309
                         sub_name);
310
                     return var->GetMutable<T>();
311 312 313 314 315
                   });
    return res;
  }

  const OperatorBase& op_;
316
  const Scope& scope_;
317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333
};

template <typename T>
struct EigenDeviceConverter;

template <>
struct EigenDeviceConverter<platform::CPUPlace> {
  using EigenDeviceType = Eigen::DefaultDevice;
};

#ifndef PADDLE_ONLY_CPU
template <>
struct EigenDeviceConverter<platform::GPUPlace> {
  using EigenDeviceType = Eigen::GpuDevice;
};
#endif

334
class ExecutionContext : public InferShapeContext {
335
 public:
336
  ExecutionContext(const OperatorBase& op, const Scope& scope,
D
dongzhihong 已提交
337
                   const platform::DeviceContext* device_context)
338
      : InferShapeContext(op, scope), device_context_(device_context) {}
339

Q
qijun 已提交
340 341 342
  template <typename PlaceType,
            typename DeviceType =
                typename EigenDeviceConverter<PlaceType>::EigenDeviceType>
343
  DeviceType& GetEigenDevice() const;
Q
qijun 已提交
344

D
dongzhihong 已提交
345
  platform::Place GetPlace() const { return device_context_->GetPlace(); }
Q
qijun 已提交
346

Q
qijun 已提交
347
  const platform::DeviceContext* device_context() const {
Q
qijun 已提交
348
    return device_context_;
Q
qijun 已提交
349
  }
Q
qijun 已提交
350

D
dongzhihong 已提交
351
  const platform::DeviceContext* device_context_;
Q
Qiao Longfei 已提交
352 353
};

Q
qijun 已提交
354 355
class OpKernel {
 public:
Q
qijun 已提交
356
  /**
357
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
358 359
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
360
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
361 362
   */

363
  virtual void Compute(const ExecutionContext& context) const = 0;
Y
Yu Yang 已提交
364 365 366 367

  virtual ~OpKernel() {}
};

Q
Qiao Longfei 已提交
368 369
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
370 371
  struct OpKernelKey {
    platform::Place place_;
Q
Qiao Longfei 已提交
372

Y
Yu Yang 已提交
373
    OpKernelKey() = default;
L
liaogang 已提交
374
    explicit OpKernelKey(const platform::DeviceContext& dev_ctx) {
Y
Yu Yang 已提交
375 376 377
      place_ = dev_ctx.GetPlace();
    }

Q
qijun 已提交
378 379 380
    bool operator==(const OpKernelKey& o) const {
      return platform::places_are_same_class(place_, o.place_);
    }
Y
Yu Yang 已提交
381 382 383 384 385 386 387 388 389 390 391
  };

  struct OpKernelHash {
    std::hash<bool> hash_;
    size_t operator()(const OpKernelKey& key) const {
      return hash_(platform::is_gpu_place(key.place_));
    }
  };

  using OpKernelMap =
      std::unordered_map<OpKernelKey, std::unique_ptr<OpKernel>, OpKernelHash>;
Q
Qiao Longfei 已提交
392

Y
Yu Yang 已提交
393 394
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
395 396
      : OperatorBase(type, inputs, outputs, attrs) {}

397
  void InferShape(const Scope& scope) const override {
398
    InferShape(InferShapeContext(*this, scope));
399 400
  }

Y
Yu Yang 已提交
401
  void Run(const Scope& scope,
Y
Yu Yang 已提交
402
           const platform::DeviceContext& dev_ctx) const final {
Q
Qiao Longfei 已提交
403
    auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
404
    opKernel->Compute(ExecutionContext(*this, scope, &dev_ctx));
Q
Qiao Longfei 已提交
405 406
  }

Y
Yu Yang 已提交
407 408 409 410
  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 已提交
411
  }
Y
Yan Chunwei 已提交
412

413 414 415 416 417 418
  bool SupportGPU() const override {
    OperatorWithKernel::OpKernelKey key;
    key.place_ = platform::GPUPlace();
    return OperatorWithKernel::AllOpKernels().at(type_).count(key) != 0;
  }

Y
Yu Yang 已提交
419
 protected:
420
  virtual void InferShape(const InferShapeContext& ctx) const = 0;
Q
Qiao Longfei 已提交
421 422 423 424
};

}  // namespace framework
}  // namespace paddle