operator.h 14.8 KB
Newer Older
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 19 20 21
#include <string>
#include <unordered_map>
#include <vector>

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

namespace paddle {
namespace framework {

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

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

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

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

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

56
class OperatorBase;
57 58
class InferShapeContext;
class ExecutionContext;
59

60 61 62
extern const Tensor* GetTensorFromVar(const Variable* var);
extern Tensor* GetTensorFromVar(Variable* var);

63 64 65 66 67 68 69 70
/**
 * 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 已提交
71 72
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
73

74 75 76
  virtual ~OperatorBase() {}

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

83
  virtual std::string DebugString() const;
84 85

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

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

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

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

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

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

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

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

112
  virtual std::vector<std::string> OutputVars(bool has_intermediate) const;
113

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

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

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

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

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

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

Y
Yu Yang 已提交
148 149 150 151
// 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 已提交
152 153
#define DEFINE_OP_CONSTRUCTOR(cls, parent_cls)             \
  cls(const std::string& type,                             \
Y
Yu Yang 已提交
154 155 156
      const ::paddle::framework::VariableNameMap& inputs,  \
      const ::paddle::framework::VariableNameMap& outputs, \
      const paddle::framework::AttributeMap& attrs)        \
Y
Yu Yang 已提交
157
      : parent_cls(type, inputs, outputs, attrs) {}
158

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

169
class InferShapeContext {
Y
Yan Chunwei 已提交
170
 public:
171 172
  InferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}
173

174 175 176 177
  const OperatorBase& op() const { return op_; }

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

178
  template <typename T>
179 180
  inline const T& Attr(const std::string& name) const {
    return op_.Attr<T>(name);
181 182
  }

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

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

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

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

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

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

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

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

  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),
244
                   [&](const std::string& sub_name) {
245
                     auto var = scope_.FindVar(sub_name);
246
                     return var == nullptr ? nullptr : &var->Get<T>();
247 248 249 250 251
                   });
    return res;
  }

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

264 265 266 267 268 269
  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];
270
    if (!in_var->IsType<LoDTensor>()) return;
271
    PADDLE_ENFORCE(out_var->IsType<LoDTensor>(),
272
                   "The %d-th output of Output(%s) must be LoDTensor.", j, out);
273 274 275
    auto in_tensor = in_var->Get<LoDTensor>();
    auto* out_tensor = out_var->GetMutable<LoDTensor>();
    out_tensor->set_lod(in_tensor.lod());
276 277
  }

278
 private:
279
  const OperatorBase& op_;
280
  const Scope& scope_;
281 282
};

283 284 285 286 287 288 289
template <>
const Tensor* InferShapeContext::Input<Tensor>(const std::string& name) const;

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

290 291 292 293 294 295 296
template <>
Tensor* InferShapeContext::Output<Tensor>(const std::string& name) const;

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

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

312
class ExecutionContext : public InferShapeContext {
313
 public:
314
  ExecutionContext(const OperatorBase& op, const Scope& scope,
315
                   const platform::DeviceContext& device_context)
316
      : InferShapeContext(op, scope), device_context_(device_context) {}
317

Q
qijun 已提交
318 319 320
  template <typename PlaceType,
            typename DeviceType =
                typename EigenDeviceConverter<PlaceType>::EigenDeviceType>
321
  DeviceType& GetEigenDevice() const;
Q
qijun 已提交
322

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

325
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
326
    return device_context_;
Q
qijun 已提交
327
  }
Q
qijun 已提交
328

329 330
 private:
  const platform::DeviceContext& device_context_;
331 332
};

333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404
class RuntimeInferShapeContext : public InferShapeContextBase {
 public:
  RuntimeInferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}

  bool HasInput(const std::string& name) const {
    auto ipt = op_.Input(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

  bool HasOutput(const std::string& name) const {
    auto ipt = op_.Output(name);
    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
    return var != nullptr;
  }

  DDim GetInputDim(const std::string& name) const {
    return GetDim(op_.Input(name));
  }

  void SetInputDim(const std::string& name, const DDim& dim) {
    SetDim(op_.Input(name), dim);
  }

  DDim GetOutputDim(const std::string& name) const {
    return GetDim(op_.Output(name));
  }

  void SetOutputDim(const std::string& name, const DDim& dim) {
    SetDim(op_.Output(name), dim);
  }

  AttrReader Attrs() const { return AttrReader(op_.Attrs()); }

  const std::vector<std::string>& Inputs(const std::string& name) const {
    return op_.Inputs(name);
  }

  const std::vector<std::string>& Outputs(const std::string& name) const {
    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;
  }

  DDim GetDim(const std::string& name) const {
    return GetTensor<false>(name)->dims();
  }

  void SetDim(const std::string& name, const DDim& dim) {
    GetTensor<true>(name)->Resize(dim);
  }

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

Q
qijun 已提交
405 406
class OpKernel {
 public:
Q
qijun 已提交
407
  /**
408
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
409 410
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
411
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
412 413
   */

414
  virtual void Compute(const ExecutionContext& context) const = 0;
Y
Yu Yang 已提交
415 416 417 418

  virtual ~OpKernel() {}
};

419 420
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
421 422
  struct OpKernelKey {
    platform::Place place_;
423

Y
Yu Yang 已提交
424
    OpKernelKey() = default;
425
    explicit OpKernelKey(const platform::DeviceContext& dev_ctx) {
Y
Yu Yang 已提交
426 427 428
      place_ = dev_ctx.GetPlace();
    }

429 430 431
    bool operator==(const OpKernelKey& o) const {
      return platform::places_are_same_class(place_, o.place_);
    }
Y
Yu Yang 已提交
432 433 434 435 436 437 438 439 440 441 442
  };

  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>;
443

Y
Yu Yang 已提交
444 445
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
446 447
      : OperatorBase(type, inputs, outputs, attrs) {}

448
  void Run(const Scope& scope,
Y
Yu Yang 已提交
449
           const platform::DeviceContext& dev_ctx) const final {
450 451 452
    RuntimeInferShapeContext infer_shape_ctx(*this, scope);
    this->InferShape(&infer_shape_ctx);

453
    auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
454
    opKernel->Compute(ExecutionContext(*this, scope, dev_ctx));
455 456
  }

Y
Yu Yang 已提交
457 458 459 460
  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 已提交
461
  }
Y
Yan Chunwei 已提交
462

463 464 465 466 467 468
  bool SupportGPU() const override {
    OperatorWithKernel::OpKernelKey key;
    key.place_ = platform::GPUPlace();
    return OperatorWithKernel::AllOpKernels().at(type_).count(key) != 0;
  }

Y
Yu Yang 已提交
469
 protected:
470
  virtual void InferShape(InferShapeContextBase* ctx) const = 0;
471 472 473 474
};

}  // namespace framework
}  // namespace paddle