operator.h 10.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>

Y
Yi Wang 已提交
23
#include "paddle/framework/attribute.h"
Y
Yu Yang 已提交
24
#include "paddle/framework/framework.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 {

34
/// If a variable is a empty variable, that name will be used.
35
constexpr char kEmptyVarName[] = "@EMPTY@";
36 37 38

/// 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.
39
constexpr char kTempVarName[] = "@TEMP@";
40 41 42 43

/// 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".
44
constexpr char kGradVarSuffix[] = "@GRAD";
45 46

/// Variables with this suffix are supposed to be filled up with zeros.
47
constexpr char kZeroVarSuffix[] = "@ZERO";
48 49 50 51 52

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

53 54
extern std::unordered_map<std::string, OpProto>& OpProtos();

Q
Qiao Longfei 已提交
55
class OperatorBase;
56 57
class InferShapeContext;
class ExecutionContext;
58

Q
Qiao Longfei 已提交
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
/**
 * 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:
  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));
  }

76
  virtual std::string DebugString() const;
Q
Qiao Longfei 已提交
77

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

Q
Qiao Longfei 已提交
82 83
  /// InferShape infer the size of Variables used by this Operator with
  /// information inside scope
Y
Yu Yang 已提交
84
  virtual void InferShape(const Scope& scope) const = 0;
Q
Qiao Longfei 已提交
85 86

  /// Net will call this function to Run an op.
Y
Yu Yang 已提交
87
  virtual void Run(const Scope& scope,
Y
Yu Yang 已提交
88 89
                   const platform::DeviceContext& dev_ctx) const = 0;

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

92 93
  virtual bool SupportGPU() const { return false; }

D
dongzhihong 已提交
94 95 96
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

Y
Yu Yang 已提交
97
  //! Get a input with argument's name described in `op_proto`
Y
Yan Chunwei 已提交
98
  const std::string& Input(const std::string& name) const;
Y
Yu Yang 已提交
99

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

108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136
  virtual std::vector<std::string> OutputVars(bool has_intermediate) const {
    std::vector<std::string> ret_val;
    if (has_intermediate) {
      // push all outputs into ret_val
      for (auto& o : outputs_) {
        ret_val.reserve(ret_val.size() + o.second.size());
        ret_val.insert(ret_val.end(), o.second.begin(), o.second.end());
      }
      return ret_val;
    }
    auto it = OpProtos().find(type_);
    PADDLE_ENFORCE(
        it != OpProtos().end(),
        "Operator %s not registered, cannot figure out intermediate outputs",
        type_);

    // get all OpProto::Var for outputs
    for (auto& o : it->second.outputs()) {
      // ignore all intermediate output
      if (o.intermediate()) continue;
      auto out = outputs_.find(o.name());
      if (out != outputs_.end()) {
        ret_val.reserve(ret_val.size() + out->second.size());
        ret_val.insert(ret_val.end(), out->second.begin(), out->second.end());
      }
    }
    return ret_val;
  }

Q
Qiao Longfei 已提交
137
 public:
Q
Qiao Longfei 已提交
138
  std::string type_;
D
dongzhihong 已提交
139 140 141 142
  // NOTE: in case of OpGrad, inputs_ contains:
  // I (Inputs)
  // O (Outputs)
  // OG (Output Gradients)
Y
Yu Yang 已提交
143 144
  std::unordered_map<std::string, std::vector<std::string>> inputs_;

D
dongzhihong 已提交
145 146
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Y
Yu Yang 已提交
147
  std::unordered_map<std::string, std::vector<std::string>> outputs_;
Q
Qiao Longfei 已提交
148
  AttributeMap attrs_;
Y
Yan Chunwei 已提交
149 150
};

151
class InferShapeContext {
Y
Yan Chunwei 已提交
152
 public:
153 154
  InferShapeContext(const OperatorBase& op, const Scope& scope)
      : op_(op), scope_(scope) {}
155

Y
Yu Yang 已提交
156 157
  size_t InputSize(const std::string& name) const {
    return op_.inputs_.at(name).size();
Y
Yan Chunwei 已提交
158 159
  }

Y
Yu Yang 已提交
160 161
  size_t OutputSize(const std::string& name) const {
    return op_.outputs_.at(name).size();
Y
Yan Chunwei 已提交
162 163
  }

164
  const Variable* InputVar(const std::string& name) const {
Y
Yu Yang 已提交
165
    return scope_.FindVar(op_.Input(name));
Y
Yan Chunwei 已提交
166 167
  }

168
  Variable* OutputVar(const std::string& name) const {
Y
Yu Yang 已提交
169
    return scope_.FindVar(op_.Output(name));
Y
Yan Chunwei 已提交
170 171
  }

172 173
  const std::vector<const Variable*> MultiInputVar(
      const std::string& name) const {
Y
Yan Chunwei 已提交
174 175
    auto names = op_.Inputs(name);
    std::vector<const Variable*> res;
176
    res.reserve(names.size());
Y
Yan Chunwei 已提交
177
    std::transform(
178
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
179
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
180 181 182
    return res;
  }

183
  std::vector<const Variable*> MultiOutputVar(const std::string& name) const {
Y
Yan Chunwei 已提交
184 185
    auto names = op_.Outputs(name);
    std::vector<const Variable*> res;
186
    res.reserve(names.size());
Y
Yan Chunwei 已提交
187
    std::transform(
188
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
189
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
190 191 192
    return res;
  }

193 194
  template <typename T>
  const T* Input(const std::string& name) const {
Y
Yu Yang 已提交
195
    auto* var = InputVar(name);
Y
Yan Chunwei 已提交
196
    PADDLE_ENFORCE_NOT_NULL(var, "Input(%s) should not be nullptr", name);
197
    return &var->Get<T>();
198 199 200 201
  }

  template <typename T>
  T* Output(const std::string& name) const {
202
    auto var = OutputVar(name);
Y
Yan Chunwei 已提交
203
    PADDLE_ENFORCE_NOT_NULL(var, "Output(%s) should not be nullptr", name);
204
    return var->GetMutable<T>();
205 206 207 208 209 210 211 212
  }

  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),
213
                   [&](const std::string& sub_name) {
214
                     auto var = scope_.FindVar(sub_name);
Y
Yan Chunwei 已提交
215 216 217
                     PADDLE_ENFORCE_NOT_NULL(
                         var, "MultiInput(%s:%s) should not be nullptr", name,
                         sub_name);
218
                     return &var->Get<T>();
219 220 221 222 223 224 225 226 227 228
                   });
    return res;
  }

  template <typename T>
  std::vector<const T*> MultiOutput(const std::string& name) const {
    auto names = op_.Outputs(name);
    std::vector<const T*> res;
    res.reserve(names.size());
    std::transform(names.begin(), names.end(), std::back_inserter(res),
229
                   [&](const std::string& sub_name) {
230
                     auto var = scope_.FindVar(sub_name);
Y
Yan Chunwei 已提交
231 232 233
                     PADDLE_ENFORCE_NOT_NULL(
                         var, "MultiOutput(%s:%s) should not be nullptr", name,
                         sub_name);
234
                     return var->GetMutable<T>();
235 236 237 238 239
                   });
    return res;
  }

  const OperatorBase& op_;
240
  const Scope& scope_;
241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257
};

template <typename T>
struct EigenDeviceConverter;

template <>
struct EigenDeviceConverter<platform::CPUPlace> {
  using EigenDeviceType = Eigen::DefaultDevice;
};

#ifndef PADDLE_ONLY_CPU
template <>
struct EigenDeviceConverter<platform::GPUPlace> {
  using EigenDeviceType = Eigen::GpuDevice;
};
#endif

258
class ExecutionContext : public InferShapeContext {
259
 public:
260
  ExecutionContext(const OperatorBase& op, const Scope& scope,
D
dongzhihong 已提交
261
                   const platform::DeviceContext* device_context)
262
      : InferShapeContext(op, scope), device_context_(device_context) {}
263

Q
qijun 已提交
264 265 266
  template <typename PlaceType,
            typename DeviceType =
                typename EigenDeviceConverter<PlaceType>::EigenDeviceType>
267
  DeviceType& GetEigenDevice() const;
Q
qijun 已提交
268

D
dongzhihong 已提交
269
  platform::Place GetPlace() const { return device_context_->GetPlace(); }
Q
qijun 已提交
270

D
dongzhihong 已提交
271
  const platform::DeviceContext* device_context_;
Q
Qiao Longfei 已提交
272 273
};

Q
qijun 已提交
274 275
class OpKernel {
 public:
Q
qijun 已提交
276
  /**
277
   * ExecutionContext is the only parameter of Kernel Run function.
Q
qijun 已提交
278 279
   * Run will get input/output variables, state such as momentum and
   * device resource such as CUDA stream, cublas handle, etc. from
280
   * ExecutionContext. User should construct it before run the Operator.
Q
qijun 已提交
281 282
   */

283
  virtual void Compute(const ExecutionContext& context) const = 0;
Y
Yu Yang 已提交
284 285 286 287

  virtual ~OpKernel() {}
};

Q
Qiao Longfei 已提交
288 289
class OperatorWithKernel : public OperatorBase {
 public:
Y
Yu Yang 已提交
290 291
  struct OpKernelKey {
    platform::Place place_;
Q
Qiao Longfei 已提交
292

Y
Yu Yang 已提交
293
    OpKernelKey() = default;
L
liaogang 已提交
294
    explicit OpKernelKey(const platform::DeviceContext& dev_ctx) {
Y
Yu Yang 已提交
295 296 297
      place_ = dev_ctx.GetPlace();
    }

Q
qijun 已提交
298 299 300
    bool operator==(const OpKernelKey& o) const {
      return platform::places_are_same_class(place_, o.place_);
    }
Y
Yu Yang 已提交
301 302 303 304 305 306 307 308 309 310 311
  };

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

313
  void InferShape(const Scope& scope) const override {
314
    InferShape(InferShapeContext(*this, scope));
315 316
  }

Y
Yu Yang 已提交
317
  void Run(const Scope& scope,
Y
Yu Yang 已提交
318
           const platform::DeviceContext& dev_ctx) const final {
Q
Qiao Longfei 已提交
319
    auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
320
    opKernel->Compute(ExecutionContext(*this, scope, &dev_ctx));
Q
Qiao Longfei 已提交
321 322
  }

Y
Yu Yang 已提交
323 324 325 326
  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 已提交
327
  }
Y
Yan Chunwei 已提交
328

329 330 331 332 333 334
  bool SupportGPU() const override {
    OperatorWithKernel::OpKernelKey key;
    key.place_ = platform::GPUPlace();
    return OperatorWithKernel::AllOpKernels().at(type_).count(key) != 0;
  }

Y
Yu Yang 已提交
335
 protected:
336
  virtual void InferShape(const InferShapeContext& ctx) const = 0;
Q
Qiao Longfei 已提交
337 338 339 340
};

}  // namespace framework
}  // namespace paddle