operator.h 8.2 KB
Newer Older
Q
Qiao Longfei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
/* 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

#include <boost/variant.hpp>
#include <string>
#include <unordered_map>
#include <vector>

Q
qijun 已提交
22 23
#include "paddle/framework/attr_checker.h"
#include "paddle/framework/op_desc.pb.h"
Y
Yan Chunwei 已提交
24
#include "paddle/framework/op_proto.pb.h"
Q
qijun 已提交
25 26 27 28 29
#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 已提交
30 31 32 33

namespace paddle {
namespace framework {

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

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

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

Q
Qiao Longfei 已提交
49 50 51 52 53 54 55 56 57
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:
58 59 60 61 62 63 64
  /// 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 已提交
65 66 67 68 69
  /// 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"; }

Q
Qiao Longfei 已提交
70 71 72 73 74 75 76 77 78
  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));
  }

79
  virtual std::string DebugString() const;
Q
Qiao Longfei 已提交
80

Q
Qiao Longfei 已提交
81 82 83 84
  /// Init will be called after CreateOperator, you can put some initialization
  /// logic here.
  virtual void Init() {}

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

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

Y
Yu Yang 已提交
93 94 95
  virtual bool IsNetOp() const { return false; }

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

Q
Qiao Longfei 已提交
106
 public:
Q
Qiao Longfei 已提交
107
  std::string type_;
Q
Qiao Longfei 已提交
108 109 110
  std::vector<std::string> inputs_;
  std::vector<std::string> outputs_;
  AttributeMap attrs_;
Y
Yan Chunwei 已提交
111
  // store the arguments' offset described in op_desc.
Y
Yu Yang 已提交
112
  std::shared_ptr<std::unordered_map<std::string, int>> in_out_idxs_;
Y
Yan Chunwei 已提交
113 114 115 116 117 118 119 120 121
};

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 {
Y
Yi Wang 已提交
122
    return scope_->FindVar(op_.inputs_[index]);
Y
Yan Chunwei 已提交
123 124 125
  }

  Variable* Output(int index) const {
Y
Yi Wang 已提交
126
    return scope_->FindVar(op_.outputs_[index]);
Y
Yan Chunwei 已提交
127 128 129
  }

  const Variable* Input(const std::string& name) const {
Y
Yi Wang 已提交
130
    return scope_->FindVar(op_.Input(name));
Y
Yan Chunwei 已提交
131 132 133
  }

  const Variable* Output(const std::string& name) const {
Y
Yi Wang 已提交
134
    return scope_->FindVar(op_.Output(name));
Y
Yan Chunwei 已提交
135 136 137 138 139 140 141
  }

  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(),
Y
Yi Wang 已提交
142
        [this](const std::string& name) { return scope_->FindVar(name); });
Y
Yan Chunwei 已提交
143 144 145 146 147 148 149 150
    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(),
Y
Yi Wang 已提交
151
        [this](const std::string& name) { return scope_->FindVar(name); });
Y
Yan Chunwei 已提交
152 153 154
    return res;
  }

Q
qijun 已提交
155 156 157 158 159 160 161
  template <typename PlaceType,
            typename DeviceType =
                typename EigenDeviceConverter<PlaceType>::EigenDeviceType>
  DeviceType* GetEigenDevice() const;

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

Y
Yan Chunwei 已提交
162 163 164
  const OperatorBase& op_;
  const std::shared_ptr<Scope>& scope_;
  const platform::DeviceContext& device_context_;
Q
Qiao Longfei 已提交
165 166
};

Q
qijun 已提交
167 168
class OpKernel {
 public:
Q
qijun 已提交
169 170 171 172 173 174 175
  /**
   * 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 已提交
176 177 178 179 180
  virtual void Compute(const KernelContext& context) const = 0;

  virtual ~OpKernel() {}
};

Y
Yu Yang 已提交
181 182 183 184 185 186 187 188 189 190 191 192 193
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 已提交
194 195
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
196 197
  struct OpKernelKey {
    platform::Place place_;
Q
Qiao Longfei 已提交
198

Y
Yu Yang 已提交
199 200 201 202 203
    OpKernelKey() = default;
    OpKernelKey(const platform::DeviceContext& dev_ctx) {
      place_ = dev_ctx.GetPlace();
    }

Q
qijun 已提交
204 205 206
    bool operator==(const OpKernelKey& o) const {
      return platform::places_are_same_class(place_, o.place_);
    }
Y
Yu Yang 已提交
207 208 209 210 211 212 213 214 215 216 217
  };

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

Y
Yu Yang 已提交
219
  void Run(const std::shared_ptr<Scope>& scope,
Y
Yu Yang 已提交
220
           const platform::DeviceContext& dev_ctx) const final {
Q
Qiao Longfei 已提交
221
    auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
Y
Yan Chunwei 已提交
222
    opKernel->Compute(KernelContext(this, scope, dev_ctx));
Q
Qiao Longfei 已提交
223 224
  }

Y
Yu Yang 已提交
225 226 227 228
  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 已提交
229
  }
Y
Yan Chunwei 已提交
230

Y
Yu Yang 已提交
231 232 233 234 235 236
  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 已提交
237
  };
Y
Yu Yang 已提交
238 239 240 241 242 243 244 245 246

 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) {
Y
Yi Wang 已提交
247
      auto var = scope->FindVar(name);
Y
Yu Yang 已提交
248 249 250 251 252 253 254 255 256 257 258
      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 已提交
259 260 261 262
};

}  // namespace framework
}  // namespace paddle