operator.h 13.7 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
#include <string>
D
dzhwinter 已提交
20
#include <tuple>
Q
Qiao Longfei 已提交
21 22 23
#include <unordered_map>
#include <vector>

Y
Yu Yang 已提交
24
#include "glog/logging.h"  // For VLOG
Y
Yi Wang 已提交
25 26 27 28 29 30 31 32 33 34 35
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/block_desc.h"
#include "paddle/fluid/framework/framework.pb.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_info.h"
#include "paddle/fluid/framework/op_kernel_type.h"
#include "paddle/fluid/framework/scope.h"
#include "paddle/fluid/framework/selected_rows.h"
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/variant.h"
Q
qijun 已提交
36
#include "paddle/utils/Error.h"
Q
Qiao Longfei 已提交
37 38 39 40

namespace paddle {
namespace framework {

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

/// 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.
46
constexpr char kTempVarName[] = "@TEMP@";
47 48 49 50

/// 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".
51
constexpr char kGradVarSuffix[] = "@GRAD";
52 53

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

D
dzhwinter 已提交
56
// define some kernel priority
57
/* Define multiple kernel type fallback order*/
D
dzhwinter 已提交
58 59
extern std::vector<std::tuple<platform::Place, LibraryType>> kKernelPriority;

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

Q
Qiao Longfei 已提交
64
class OperatorBase;
65
class ExecutionContext;
66

Q
Qiao Longfei 已提交
67 68 69 70 71 72 73 74
/**
 * 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 已提交
75 76
  OperatorBase(const std::string& type, const VariableNameMap& inputs,
               const VariableNameMap& outputs, const AttributeMap& attrs);
77

Q
Qiao Longfei 已提交
78 79 80
  virtual ~OperatorBase() {}

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

87 88 89 90
  /// if scope is not null, also show dimensions of arguments
  virtual std::string DebugStringEx(const Scope* scope) const;

  std::string DebugString() const { return DebugStringEx(nullptr); }
Q
Qiao Longfei 已提交
91

92 93 94
  /// Net will call this interface function to Run an op.
  //  The implementation should be written at RunImpl
  void Run(const Scope& scope, const platform::Place& place);
Y
Yu Yang 已提交
95

T
typhoonzero 已提交
96 97 98
  // FIXME(typhoonzero): this is only used for recv_op to stop event_loop.
  virtual void Stop() {}

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

101 102
  virtual bool SupportGPU() const { return false; }

D
dongzhihong 已提交
103 104 105
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

Y
Yu Yang 已提交
106 107
  const VariableNameMap& Inputs() const { return inputs_; }
  const VariableNameMap& Outputs() const { return outputs_; }
108

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

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

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

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

Q
qiaolongfei 已提交
124
  const std::string& Type() const { return type_; }
Q
qiaolongfei 已提交
125
  void SetType(const std::string& type) { type_ = type; }
Y
Yi Wang 已提交
126 127
  const AttributeMap& Attrs() const { return attrs_; }

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

Q
qiaolongfei 已提交
132
 protected:
Q
Qiao Longfei 已提交
133
  std::string type_;
D
dongzhihong 已提交
134
  // NOTE: in case of OpGrad, inputs_ contains:
135
  // I (Inputs)
D
dongzhihong 已提交
136 137
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
138
  VariableNameMap inputs_;
Y
Yu Yang 已提交
139

D
dongzhihong 已提交
140 141
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
142
  VariableNameMap outputs_;
Q
Qiao Longfei 已提交
143
  AttributeMap attrs_;
144 145 146 147

 private:
  void GenerateTemporaryNames();
  void CheckAllInputOutputSet() const;
148 149
  virtual void RunImpl(const Scope& scope,
                       const platform::Place& place) const = 0;
Y
Yan Chunwei 已提交
150 151
};

Y
Yu Yang 已提交
152 153
// Macro for define a clone method.
// If you are writing an kernel operator, `Clone` will be defined when you
154
// register it. i.e. `Clone` method is not needed to define by yourself.
155 156 157
#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 已提交
158
  }
Y
Yu Yang 已提交
159

Y
Yu Yang 已提交
160 161 162 163
// 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 已提交
164 165
#define DEFINE_OP_CONSTRUCTOR(cls, parent_cls)             \
  cls(const std::string& type,                             \
Y
Yu Yang 已提交
166 167 168
      const ::paddle::framework::VariableNameMap& inputs,  \
      const ::paddle::framework::VariableNameMap& outputs, \
      const paddle::framework::AttributeMap& attrs)        \
Y
Yu Yang 已提交
169
      : parent_cls(type, inputs, outputs, attrs) {}
Y
Yu Yang 已提交
170

171 172
class NOP : public OperatorBase {
 public:
173
  using OperatorBase::OperatorBase;
174 175 176
  std::unique_ptr<OperatorBase> Clone() const override {
    return std::unique_ptr<OperatorBase>(new NOP(*this));
  }
177 178 179 180

 private:
  void RunImpl(const Scope& scope,
               const platform::Place& place) const override {}
181 182
};

183
class ExecutionContext {
Y
Yan Chunwei 已提交
184
 public:
185 186 187
  ExecutionContext(const OperatorBase& op, const Scope& scope,
                   const platform::DeviceContext& device_context)
      : op_(op), scope_(scope), device_context_(device_context) {}
188

Q
qiaolongfei 已提交
189 190 191 192
  const OperatorBase& op() const { return op_; }

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

Q
qiaolongfei 已提交
193
  template <typename T>
Y
Yu Yang 已提交
194 195
  inline const T& Attr(const std::string& name) const {
    return op_.Attr<T>(name);
Q
qiaolongfei 已提交
196 197
  }

Y
Yu Yang 已提交
198
  size_t InputSize(const std::string& name) const {
Y
Yu Yang 已提交
199
    return op_.Inputs(name).size();
Y
Yan Chunwei 已提交
200 201
  }

Y
Yu Yang 已提交
202
  size_t OutputSize(const std::string& name) const {
Y
Yu Yang 已提交
203
    return op_.Outputs(name).size();
Y
Yan Chunwei 已提交
204 205
  }

206
  const Variable* InputVar(const std::string& name) const {
207
    auto ipt = op_.Input(name);
Y
Yu Yang 已提交
208
    return ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
Y
Yan Chunwei 已提交
209 210
  }

211
  Variable* OutputVar(const std::string& name) const {
212
    auto opt = op_.Output(name);
Y
Yu Yang 已提交
213
    return opt == kEmptyVarName ? nullptr : scope_.FindVar(opt);
Y
Yan Chunwei 已提交
214 215
  }

216 217
  const std::vector<const Variable*> MultiInputVar(
      const std::string& name) const {
Y
Yan Chunwei 已提交
218 219
    auto names = op_.Inputs(name);
    std::vector<const 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
  std::vector<Variable*> MultiOutputVar(const std::string& name) const {
Y
Yan Chunwei 已提交
230
    auto names = op_.Outputs(name);
231
    std::vector<Variable*> res;
232
    res.reserve(names.size());
233 234
    std::transform(names.begin(), names.end(), std::back_inserter(res),
                   [this](const std::string& name) {
Y
Yu Yang 已提交
235 236
                     return name == kEmptyVarName ? nullptr
                                                  : scope_.FindVar(name);
237
                   });
Y
Yan Chunwei 已提交
238 239 240
    return res;
  }

241 242
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
243
    auto* var = InputVar(name);
244
    return var == nullptr ? nullptr : &var->Get<T>();
245 246 247 248
  }

  template <typename T>
  T* Output(const std::string& name) const {
249
    auto var = OutputVar(name);
250
    return var == nullptr ? nullptr : var->GetMutable<T>();
251 252 253 254 255 256 257 258
  }

  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),
259
                   [&](const std::string& sub_name) {
260
                     auto var = scope_.FindVar(sub_name);
261
                     return var == nullptr ? nullptr : &var->Get<T>();
262 263 264 265 266
                   });
    return res;
  }

  template <typename T>
267
  std::vector<T*> MultiOutput(const std::string& name) const {
268
    auto names = op_.Outputs(name);
269
    std::vector<T*> res;
270 271
    res.reserve(names.size());
    std::transform(names.begin(), names.end(), std::back_inserter(res),
272
                   [&](const std::string& sub_name) {
273
                     auto var = scope_.FindVar(sub_name);
274
                     return var == nullptr ? nullptr : var->GetMutable<T>();
275 276 277 278
                   });
    return res;
  }

279 280 281 282 283 284
  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];
285
    if (!in_var->IsType<LoDTensor>()) return;
286
    PADDLE_ENFORCE(out_var->IsType<LoDTensor>(),
287
                   "The %d-th output of Output(%s) must be LoDTensor.", j, out);
288 289 290
    auto in_tensor = in_var->Get<LoDTensor>();
    auto* out_tensor = out_var->GetMutable<LoDTensor>();
    out_tensor->set_lod(in_tensor.lod());
291 292
  }

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

Q
QI JUN 已提交
295 296 297 298 299
  template <typename DeviceContextType>
  const DeviceContextType& device_context() const {
    return *reinterpret_cast<const DeviceContextType*>(&device_context_);
  }

300
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
301
    return device_context_;
Q
qijun 已提交
302
  }
Q
qijun 已提交
303

Q
QI JUN 已提交
304 305 306 307 308 309 310 311
#ifdef PADDLE_WITH_CUDA
  const inline platform::CUDADeviceContext& cuda_device_context() const {
    PADDLE_ENFORCE(platform::is_gpu_place(device_context_.GetPlace()));
    return *reinterpret_cast<const platform::CUDADeviceContext*>(
        &device_context_);
  }
#endif

D
dzhwinter 已提交
312
  //! Get actual name vector for this input.
D
Dong Zhihong 已提交
313 314 315
  const std::vector<std::string>& Inputs(const std::string& name) const {
    return op_.Inputs(name);
  }
D
Dong Zhihong 已提交
316

D
dzhwinter 已提交
317
  //! Get actual name vector for this output.
D
Dong Zhihong 已提交
318 319 320 321
  const std::vector<std::string>& Outputs(const std::string& name) const {
    return op_.Outputs(name);
  }

322
 private:
323 324
  const OperatorBase& op_;
  const Scope& scope_;
325
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
326 327
};

328 329 330 331 332 333 334 335 336 337 338 339 340 341
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;

Y
Yu Yang 已提交
342
class OpKernelBase {
Q
qijun 已提交
343
 public:
Q
qijun 已提交
344
  /**
345
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
346 347
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
348
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
349 350
   */

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

Y
Yu Yang 已提交
353 354 355 356 357 358 359
  virtual ~OpKernelBase() = default;
};

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

Y
Yu Yang 已提交
362 363
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
364
  using OpKernelMap =
Y
Yu Yang 已提交
365 366
      std::unordered_map<OpKernelType, std::unique_ptr<OpKernelBase>,
                         OpKernelType::Hash>;
Q
Qiao Longfei 已提交
367

Y
Yu Yang 已提交
368 369
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
370 371
      : OperatorBase(type, inputs, outputs, attrs) {}

Y
Yu Yang 已提交
372 373 374 375
  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 已提交
376
  }
Y
Yan Chunwei 已提交
377

378
  bool SupportGPU() const override {
Y
Yu Yang 已提交
379 380 381 382 383
    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_);
                       });
384 385
  }

386 387 388
  virtual void InferShape(InferShapeContext* ctx) const {
    OpInfoMap::Instance().Get(Type()).infer_shape_(ctx);
  }
Y
Yu Yang 已提交
389

Q
qiaolongfei 已提交
390
 protected:
391 392 393 394
  virtual OpKernelType GetExpectedKernelType(const ExecutionContext& ctx) const;
  virtual OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const Tensor& tensor,
      const OpKernelType& expected_kernel_type) const;
Y
Yu Yang 已提交
395 396

 private:
Y
Yu Yang 已提交
397 398
  // indicate kernel DataType by input data. Defaultly all input data must be
  // same.
399
  proto::DataType IndicateDataType(const ExecutionContext& ctx) const;
400
  void RunImpl(const Scope& scope, const platform::Place& place) const final;
Q
Qiao Longfei 已提交
401 402
};

Y
Yu Yang 已提交
403 404
extern bool OpSupportGPU(const std::string& op_type);

Q
Qiao Longfei 已提交
405 406
}  // namespace framework
}  // namespace paddle