operator.h 16.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>
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"
Y
Yu Yang 已提交
25
#include "paddle/framework/data_type.h"
Y
Yu Yang 已提交
26
#include "paddle/framework/framework.pb.h"
27
#include "paddle/framework/lod_tensor.h"
Q
qijun 已提交
28
#include "paddle/framework/scope.h"
Q
Qiao Longfei 已提交
29
#include "paddle/framework/shape_inference.h"
Q
qijun 已提交
30 31 32
#include "paddle/framework/tensor.h"
#include "paddle/platform/device_context.h"
#include "paddle/platform/place.h"
Y
Yu Yang 已提交
33
#include "paddle/platform/variant.h"
Q
qijun 已提交
34
#include "paddle/utils/Error.h"
Q
Qiao Longfei 已提交
35 36 37 38

namespace paddle {
namespace framework {

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

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

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

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

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

Q
Qiao Longfei 已提交
58
class OperatorBase;
59 60
class InferShapeContext;
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 InferShapeContext {
Y
Yan Chunwei 已提交
172
 public:
173 174
  InferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}
175

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

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

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

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

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

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

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

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

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

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

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

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

  template <typename T>
254
  std::vector<T*> MultiOutput(const std::string& name) const {
255
    auto names = op_.Outputs(name);
256
    std::vector<T*> res;
257 258
    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->GetMutable<T>();
262 263 264 265
                   });
    return res;
  }

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

Q
qiaolongfei 已提交
280
 private:
281
  const OperatorBase& op_;
282
  const Scope& scope_;
283 284
};

285 286 287 288 289 290 291
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;

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

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

299
class ExecutionContext : public InferShapeContext {
300
 public:
301
  ExecutionContext(const OperatorBase& op, const Scope& scope,
302
                   const platform::DeviceContext& device_context)
303
      : InferShapeContext(op, scope), device_context_(device_context) {}
304

Q
qijun 已提交
305
  template <typename PlaceType,
306 307
            typename DeviceType = typename platform::EigenDeviceConverter<
                PlaceType>::EigenDeviceType>
308
  DeviceType& GetEigenDevice() const;
Q
qijun 已提交
309

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

312
  const platform::DeviceContext& device_context() const {
Q
qijun 已提交
313
    return device_context_;
Q
qijun 已提交
314
  }
Q
qijun 已提交
315

316 317
 private:
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
318 319
};

Q
Qiao Longfei 已提交
320 321 322 323 324 325 326 327 328 329 330 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
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_;
};

Y
Yu Yang 已提交
392
class OpKernelBase {
Q
qijun 已提交
393
 public:
Q
qijun 已提交
394
  /**
395
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
396 397
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
398
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
399 400
   */

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

Y
Yu Yang 已提交
403 404 405 406 407 408 409
  virtual ~OpKernelBase() = default;
};

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

Q
Qiao Longfei 已提交
412 413
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
414 415
  struct OpKernelKey {
    platform::Place place_;
Y
Yu Yang 已提交
416
    DataType data_type_;
Q
Qiao Longfei 已提交
417

Y
Yu Yang 已提交
418 419 420 421 422
    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 已提交
423

Q
qijun 已提交
424
    bool operator==(const OpKernelKey& o) const {
Y
Yu Yang 已提交
425 426
      return platform::places_are_same_class(place_, o.place_) &&
             data_type_ == o.data_type_;
Q
qijun 已提交
427
    }
Y
Yu Yang 已提交
428 429 430
  };

  struct OpKernelHash {
Y
Yu Yang 已提交
431
    std::hash<int> hash_;
Y
Yu Yang 已提交
432
    size_t operator()(const OpKernelKey& key) const {
Y
Yu Yang 已提交
433 434
      int place = key.place_.which();
      int data_type = static_cast<int>(key.data_type_);
Y
Yu Yang 已提交
435 436
      int pre_hash = data_type << NUM_PLACE_TYPE_LIMIT_IN_BIT |
                     (place & ((1 << NUM_PLACE_TYPE_LIMIT_IN_BIT) - 1));
Y
Yu Yang 已提交
437
      return hash_(pre_hash);
Y
Yu Yang 已提交
438 439 440 441
    }
  };

  using OpKernelMap =
Y
Yu Yang 已提交
442 443
      std::unordered_map<OpKernelKey, std::unique_ptr<OpKernelBase>,
                         OpKernelHash>;
Q
Qiao Longfei 已提交
444

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

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

Y
Yu Yang 已提交
454 455 456 457
    ExecutionContext ctx(*this, scope, dev_ctx);
    auto& opKernel = AllOpKernels().at(type_).at(
        OpKernelKey(IndicateDataType(ctx), dev_ctx));
    opKernel->Compute(ctx);
Q
Qiao Longfei 已提交
458 459
  }

Y
Yu Yang 已提交
460 461 462 463
  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 已提交
464
  }
Y
Yan Chunwei 已提交
465

466
  bool SupportGPU() const override {
Y
Yu Yang 已提交
467 468 469 470 471
    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_);
                       });
472 473
  }

Y
Yu Yang 已提交
474
 protected:
Q
Qiao Longfei 已提交
475
  virtual void InferShape(InferShapeContextBase* ctx) const = 0;
Y
Yu Yang 已提交
476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503

  // 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 已提交
504 505 506 507
};

}  // namespace framework
}  // namespace paddle