operator.h 15.0 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 "glog/logging.h"  // For VLOG
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"
Y
Yu Yang 已提交
29
#include "paddle/framework/op_info.h"
Q
qijun 已提交
30
#include "paddle/framework/scope.h"
Q
QI JUN 已提交
31
#include "paddle/framework/selected_rows.h"
Q
qijun 已提交
32 33 34
#include "paddle/framework/tensor.h"
#include "paddle/platform/device_context.h"
#include "paddle/platform/place.h"
Y
Yu Yang 已提交
35
#include "paddle/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 56 57 58 59

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

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

Q
Qiao Longfei 已提交
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

Q
Qiao Longfei 已提交
74 75 76
  virtual ~OperatorBase() {}

  template <typename T>
Y
Yu Yang 已提交
77
  inline const T& Attr(const std::string& name) const {
Q
Qiao Longfei 已提交
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;
Q
Qiao Longfei 已提交
84 85

  /// Net will call this function to Run an op.
Y
Yu Yang 已提交
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

Y
Yu Yang 已提交
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.
Y
Yu Yang 已提交
119 120
  // Use unique_ptr to prevent caller forget to delete this pointer.
  virtual std::unique_ptr<OperatorBase> Clone() const = 0;
Y
Yu Yang 已提交
121

Q
qiaolongfei 已提交
122
 protected:
Q
Qiao Longfei 已提交
123
  std::string type_;
D
dongzhihong 已提交
124
  // NOTE: in case of OpGrad, inputs_ contains:
125
  // I (Inputs)
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_;
Q
Qiao Longfei 已提交
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.
143 144 145
#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 已提交
146
  }
Y
Yu Yang 已提交
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) {}
Y
Yu Yang 已提交
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 ExecutionContext {
Y
Yan Chunwei 已提交
170
 public:
171 172 173
  ExecutionContext(const OperatorBase& op, const Scope& scope,
                   const platform::DeviceContext& device_context)
      : op_(op), scope_(scope), device_context_(device_context) {}
174

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

D
dzhwinter 已提交
290
  //! Get actual name vector for this input.
D
Dong Zhihong 已提交
291 292 293
  const std::vector<std::string>& Inputs(const std::string& name) const {
    return op_.Inputs(name);
  }
D
Dong Zhihong 已提交
294

D
dzhwinter 已提交
295
  //! Get actual name vector for this output.
D
Dong Zhihong 已提交
296 297 298 299
  const std::vector<std::string>& Outputs(const std::string& name) const {
    return op_.Outputs(name);
  }

武毅 已提交
300 301 302 303 304 305 306 307 308
#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

309
 private:
310 311
  const OperatorBase& op_;
  const Scope& scope_;
312
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
313 314
};

315 316 317 318 319 320 321 322 323 324 325 326 327 328
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 已提交
329
class OpKernelBase {
Q
qijun 已提交
330
 public:
Q
qijun 已提交
331
  /**
332
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
333 334
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
335
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
336 337
   */

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

Y
Yu Yang 已提交
340 341 342 343 344 345 346
  virtual ~OpKernelBase() = default;
};

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

Q
Qiao Longfei 已提交
349 350
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
351 352
  struct OpKernelKey {
    platform::Place place_;
Y
Yu Yang 已提交
353
    DataType data_type_;
Q
Qiao Longfei 已提交
354

Y
Yu Yang 已提交
355 356 357 358 359
    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 已提交
360

Q
qijun 已提交
361
    bool operator==(const OpKernelKey& o) const {
Y
Yu Yang 已提交
362 363
      return platform::places_are_same_class(place_, o.place_) &&
             data_type_ == o.data_type_;
Q
qijun 已提交
364
    }
Y
Yu Yang 已提交
365 366 367
  };

  struct OpKernelHash {
Y
Yu Yang 已提交
368
    std::hash<int> hash_;
Y
Yu Yang 已提交
369
    size_t operator()(const OpKernelKey& key) const {
Y
Yu Yang 已提交
370 371
      int place = key.place_.which();
      int data_type = static_cast<int>(key.data_type_);
Y
Yu Yang 已提交
372 373
      int pre_hash = data_type << NUM_PLACE_TYPE_LIMIT_IN_BIT |
                     (place & ((1 << NUM_PLACE_TYPE_LIMIT_IN_BIT) - 1));
Y
Yu Yang 已提交
374
      return hash_(pre_hash);
Y
Yu Yang 已提交
375 376 377 378
    }
  };

  using OpKernelMap =
Y
Yu Yang 已提交
379 380
      std::unordered_map<OpKernelKey, std::unique_ptr<OpKernelBase>,
                         OpKernelHash>;
Q
Qiao Longfei 已提交
381

Y
Yu Yang 已提交
382 383
  OperatorWithKernel(const std::string& type, const VariableNameMap& inputs,
                     const VariableNameMap& outputs, const AttributeMap& attrs)
Y
Yu Yang 已提交
384 385
      : OperatorBase(type, inputs, outputs, attrs) {}

Y
Yu Yang 已提交
386
  void Run(const Scope& scope,
387
           const platform::DeviceContext& dev_ctx) const final;
Q
Qiao Longfei 已提交
388

Y
Yu Yang 已提交
389 390 391 392
  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 已提交
393
  }
Y
Yan Chunwei 已提交
394

395
  bool SupportGPU() const override {
Y
Yu Yang 已提交
396 397 398 399 400
    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_);
                       });
401 402
  }

403 404 405
  virtual void InferShape(InferShapeContext* ctx) const {
    OpInfoMap::Instance().Get(Type()).infer_shape_(ctx);
  }
Y
Yu Yang 已提交
406

Q
qiaolongfei 已提交
407
 protected:
Y
Yu Yang 已提交
408 409 410
  // indicate kernel DataType by input data. Defaultly all input data must be
  // same.
  virtual DataType IndicateDataType(const ExecutionContext& ctx) const {
Q
Qiao Longfei 已提交
411
    VLOG(3) << "Default IndicateDataType " << this->Type();
Y
Yu Yang 已提交
412 413 414 415 416 417 418 419 420 421 422
    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>();
Q
QI JUN 已提交
423 424
          } else if (var->IsType<SelectedRows>()) {
            t = &(var->Get<SelectedRows>().value());
Y
Yu Yang 已提交
425 426 427
          }
          if (t != nullptr) {
            int tmp = static_cast<int>(ToDataType(t->type()));
Y
Yu Yang 已提交
428
            VLOG(3) << "Input " << ipt_name << " with data_type " << tmp;
Y
Yu Yang 已提交
429
            PADDLE_ENFORCE(tmp == data_type || data_type == -1,
Y
Yu Yang 已提交
430
                           "DataType of Paddle Op %s must be same.", Type());
Y
Yu Yang 已提交
431 432 433 434 435 436 437 438
            data_type = tmp;
          }
        }
      }
    }
    PADDLE_ENFORCE(data_type != -1, "DataType should be indicated by input");
    return static_cast<DataType>(data_type);
  }
Q
Qiao Longfei 已提交
439 440
};

441 442 443
std::ostream& operator<<(std::ostream& os,
                         const OperatorWithKernel::OpKernelKey& kernel_key);

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

Q
Qiao Longfei 已提交
446 447
}  // namespace framework
}  // namespace paddle