operator.h 8.6 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 22
#include <boost/variant.hpp>
#include <string>
#include <unordered_map>
#include <vector>

Q
qijun 已提交
23 24
#include "paddle/framework/attr_checker.h"
#include "paddle/framework/op_desc.pb.h"
Y
Yan Chunwei 已提交
25
#include "paddle/framework/op_proto.pb.h"
Q
qijun 已提交
26 27 28 29 30
#include "paddle/framework/scope.h"
#include "paddle/framework/tensor.h"
#include "paddle/platform/device_context.h"
#include "paddle/platform/place.h"
#include "paddle/utils/Error.h"
Q
Qiao Longfei 已提交
31 32 33 34

namespace paddle {
namespace framework {

Q
qijun 已提交
35 36 37 38
template <typename T>
struct EigenDeviceConverter;

template <>
Q
qijun 已提交
39
struct EigenDeviceConverter<platform::CPUPlace> {
Q
qijun 已提交
40 41 42 43 44
  using EigenDeviceType = Eigen::DefaultDevice;
};

#ifndef PADDLE_ONLY_CPU
template <>
Q
qijun 已提交
45
struct EigenDeviceConverter<platform::GPUPlace> {
Q
qijun 已提交
46 47 48 49
  using EigenDeviceType = Eigen::GpuDevice;
};
#endif

Q
Qiao Longfei 已提交
50 51 52 53 54 55 56 57 58
class OperatorBase;
/**
 * 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:
59 60 61 62 63 64 65
  /// If a variable is a empty variable, that name will be used.
  static std::string EMPTY_VAR_NAME() { return "@EMPTY@"; }

  /// 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.
  static std::string TMP_VAR_NAME() { return "@TEMP@"; }

F
fengjiayi 已提交
66 67 68 69 70
  /// 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".
  static std::string GRAD_VAR_SUFFIX() { return "@GRAD"; }

71 72 73
  /// Variables with this suffix are supposed to be filled up with zeros.
  static std::string ZERO_VAR_SUFFIX() { return "@ZERO"; }

Q
Qiao Longfei 已提交
74 75 76 77 78 79 80 81 82
  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));
  }

83
  virtual std::string DebugString() const;
Q
Qiao Longfei 已提交
84

Q
Qiao Longfei 已提交
85 86 87 88
  /// Init will be called after CreateOperator, you can put some initialization
  /// logic here.
  virtual void Init() {}

Q
Qiao Longfei 已提交
89 90
  /// InferShape infer the size of Variables used by this Operator with
  /// information inside scope
Y
Yu Yang 已提交
91
  virtual void InferShape(const std::shared_ptr<Scope>& scope) const = 0;
Q
Qiao Longfei 已提交
92 93

  /// Net will call this function to Run an op.
Y
Yu Yang 已提交
94
  virtual void Run(const std::shared_ptr<Scope>& scope,
Y
Yu Yang 已提交
95 96
                   const platform::DeviceContext& dev_ctx) const = 0;

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

D
dongzhihong 已提交
99 100 101
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

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

Q
Qiao Longfei 已提交
113
 public:
Q
Qiao Longfei 已提交
114
  std::string type_;
D
dongzhihong 已提交
115 116 117 118
  // NOTE: in case of OpGrad, inputs_ contains:
  // I (Inputs)
  // O (Outputs)
  // OG (Output Gradients)
Q
Qiao Longfei 已提交
119
  std::vector<std::string> inputs_;
D
dongzhihong 已提交
120 121
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Q
Qiao Longfei 已提交
122 123
  std::vector<std::string> outputs_;
  AttributeMap attrs_;
Y
Yan Chunwei 已提交
124
  // store the arguments' offset described in op_desc.
Y
Yu Yang 已提交
125
  std::shared_ptr<std::unordered_map<std::string, int>> in_out_idxs_;
Y
Yan Chunwei 已提交
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167
};

class KernelContext {
 public:
  KernelContext(const OperatorBase* op, const std::shared_ptr<Scope>& scope,
                const platform::DeviceContext& device_context)
      : op_(*op), scope_(scope), device_context_(device_context) {}

  const Variable* Input(int index) const {
    return scope_->GetVariable(op_.inputs_[index]);
  }

  Variable* Output(int index) const {
    return scope_->GetVariable(op_.outputs_[index]);
  }

  const Variable* Input(const std::string& name) const {
    return scope_->GetVariable(op_.Input(name));
  }

  const Variable* Output(const std::string& name) const {
    return scope_->GetVariable(op_.Output(name));
  }

  const std::vector<const Variable*> Inputs(const std::string& name) const {
    auto names = op_.Inputs(name);
    std::vector<const Variable*> res;
    std::transform(
        names.begin(), names.end(), res.begin(),
        [this](const std::string& name) { return scope_->GetVariable(name); });
    return res;
  }

  const std::vector<const Variable*> Outputs(const std::string& name) const {
    auto names = op_.Outputs(name);
    std::vector<const Variable*> res;
    std::transform(
        names.begin(), names.end(), res.begin(),
        [this](const std::string& name) { return scope_->GetVariable(name); });
    return res;
  }

Q
qijun 已提交
168 169 170 171 172 173 174
  template <typename PlaceType,
            typename DeviceType =
                typename EigenDeviceConverter<PlaceType>::EigenDeviceType>
  DeviceType* GetEigenDevice() const;

  platform::Place GetPlace() const { return device_context_.GetPlace(); }

Y
Yan Chunwei 已提交
175 176 177
  const OperatorBase& op_;
  const std::shared_ptr<Scope>& scope_;
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
178 179
};

Q
qijun 已提交
180 181
class OpKernel {
 public:
Q
qijun 已提交
182 183 184 185 186 187 188
  /**
   * KernelContext is the only parameter of Kernel Run function.
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
   * KernelContext. User should construct it before run the Operator.
   */

Y
Yu Yang 已提交
189 190 191 192 193
  virtual void Compute(const KernelContext& context) const = 0;

  virtual ~OpKernel() {}
};

Y
Yu Yang 已提交
194 195 196 197 198 199 200 201 202 203 204 205 206
template <typename T>
struct VarToTensor {};

template <>
struct VarToTensor<Tensor*> {
  Tensor* operator()(Variable* var) { return var->GetMutable<Tensor>(); }
};

template <>
struct VarToTensor<const Tensor*> {
  const Tensor* operator()(Variable* var) { return &var->Get<Tensor>(); }
};

Q
Qiao Longfei 已提交
207 208
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
209 210
  struct OpKernelKey {
    platform::Place place_;
Q
Qiao Longfei 已提交
211

Y
Yu Yang 已提交
212 213 214 215 216
    OpKernelKey() = default;
    OpKernelKey(const platform::DeviceContext& dev_ctx) {
      place_ = dev_ctx.GetPlace();
    }

Q
qijun 已提交
217 218 219
    bool operator==(const OpKernelKey& o) const {
      return platform::places_are_same_class(place_, o.place_);
    }
Y
Yu Yang 已提交
220 221 222 223 224 225 226 227 228 229 230
  };

  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 已提交
231

Y
Yu Yang 已提交
232
  void Run(const std::shared_ptr<Scope>& scope,
Y
Yu Yang 已提交
233
           const platform::DeviceContext& dev_ctx) const final {
Q
Qiao Longfei 已提交
234
    auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
Y
Yan Chunwei 已提交
235
    opKernel->Compute(KernelContext(this, scope, dev_ctx));
Q
Qiao Longfei 已提交
236 237
  }

Y
Yu Yang 已提交
238 239 240 241
  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 已提交
242
  }
Y
Yan Chunwei 已提交
243

Y
Yu Yang 已提交
244 245 246 247 248 249
  void InferShape(const std::shared_ptr<Scope>& scope) const final {
    std::vector<const Tensor*> ins;
    VarNamesToTensors(scope, inputs_, &ins);
    std::vector<Tensor*> outs;
    VarNamesToTensors(scope, outputs_, &outs);
    InferShape(ins, outs);
Y
Yu Yang 已提交
250
  };
Y
Yu Yang 已提交
251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271

 private:
  template <typename T>
  void VarNamesToTensors(const std::shared_ptr<Scope>& scope,
                         const std::vector<std::string>& var_names,
                         std::vector<T>* container) const {
    container->reserve(var_names.size());
    VarToTensor<T> convert;
    for (auto& name : var_names) {
      auto var = scope->GetVariable(name);
      if (var != nullptr) {
        container->push_back(convert(var));
      } else {
        container->push_back(nullptr);
      }
    }
  }

 protected:
  virtual void InferShape(const std::vector<const Tensor*>& inputs,
                          const std::vector<Tensor*>& outputs) const = 0;
Q
Qiao Longfei 已提交
272 273 274 275
};

}  // namespace framework
}  // namespace paddle