operator.h 10.4 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"
Q
qijun 已提交
24
#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 {

35 36 37 38 39 40 41 42 43 44 45 46 47
/// If a variable is a empty variable, that name will be used.
const std::string kEmptyVarName = "@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.
const std::string kTempVarName = "@TEMP@";

/// 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".
const std::string kGradVarSuffix = "@GRAD";

/// Variables with this suffix are supposed to be filled up with zeros.
Y
Yi Wang 已提交
48
const std::string kZeroVarSuffix = "@ZERO";
49 50 51 52 53

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

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

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

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

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

Q
Qiao Longfei 已提交
81 82
  /// InferShape infer the size of Variables used by this Operator with
  /// information inside scope
Y
Yu Yang 已提交
83
  virtual void InferShape(const Scope& scope) const = 0;
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; }

D
dongzhihong 已提交
91 92 93
  /// rename inputs outputs name
  void Rename(const std::string& old_name, const std::string& new_name);

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

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_;
D
dongzhihong 已提交
108 109 110 111
  // NOTE: in case of OpGrad, inputs_ contains:
  // I (Inputs)
  // O (Outputs)
  // OG (Output Gradients)
Q
Qiao Longfei 已提交
112
  std::vector<std::string> inputs_;
D
dongzhihong 已提交
113 114
  // NOTE: in case of OpGrad, outputs_ contains
  // IG (Inputs Gradients)
Q
Qiao Longfei 已提交
115 116
  std::vector<std::string> outputs_;
  AttributeMap attrs_;
Y
Yan Chunwei 已提交
117
  // store the arguments' offset described in op_desc.
Y
Yu Yang 已提交
118
  std::shared_ptr<std::unordered_map<std::string, int>> in_out_idxs_;
Y
Yan Chunwei 已提交
119 120
};

121
class OperatorContext {
Y
Yan Chunwei 已提交
122
 public:
123
  OperatorContext(const OperatorBase* op, const Scope& scope)
124 125 126
      : op_(*op), scope_(scope) {}

  size_t InputSize() const { return op_.inputs_.size(); }
Y
Yan Chunwei 已提交
127

128 129
  size_t OutputSize() const { return op_.outputs_.size(); }

130
  const Variable* InputVar(const size_t index) const {
131
    return scope_.FindVar(op_.inputs_.at(index));
Y
Yan Chunwei 已提交
132 133
  }

134
  Variable* OutputVar(const size_t index) const {
135
    return scope_.FindVar(op_.outputs_.at(index));
Y
Yan Chunwei 已提交
136 137
  }

138
  const Variable* InputVar(const std::string& name) const {
Y
Yu Yang 已提交
139
    return scope_.FindVar(op_.Input(name));
Y
Yan Chunwei 已提交
140 141
  }

142
  Variable* OutputVar(const std::string& name) const {
Y
Yu Yang 已提交
143
    return scope_.FindVar(op_.Output(name));
Y
Yan Chunwei 已提交
144 145
  }

146 147
  const std::vector<const Variable*> MultiInputVar(
      const std::string& name) const {
Y
Yan Chunwei 已提交
148 149
    auto names = op_.Inputs(name);
    std::vector<const Variable*> res;
150
    res.reserve(names.size());
Y
Yan Chunwei 已提交
151
    std::transform(
152
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
153
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
154 155 156
    return res;
  }

157
  std::vector<const Variable*> MultiOutputVar(const std::string& name) const {
Y
Yan Chunwei 已提交
158 159
    auto names = op_.Outputs(name);
    std::vector<const Variable*> res;
160
    res.reserve(names.size());
Y
Yan Chunwei 已提交
161
    std::transform(
162
        names.begin(), names.end(), std::back_inserter(res),
Y
Yu Yang 已提交
163
        [this](const std::string& name) { return scope_.FindVar(name); });
Y
Yan Chunwei 已提交
164 165 166
    return res;
  }

167
  template <typename T>
168 169 170 171
  const T* Input(const size_t index) const {
    auto var = InputVar(index);
    PADDLE_ENFORCE(var != nullptr, "Input(%d) should not be nullptr", index);
    return &var->Get<T>();
172 173 174
  }

  template <typename T>
175 176
  T* Output(const size_t index) const {
    auto var = OutputVar(index);
Y
Yan Chunwei 已提交
177 178 179 180 181
    PADDLE_ENFORCE(
        var != nullptr,
        "Output(%d) not be nullptr, which means variable [%s] does not "
        "exist in scope",
        index, op_.outputs_[index]);
182
    return var->GetMutable<T>();
183 184 185 186
  }

  template <typename T>
  const T* Input(const std::string& name) const {
187 188 189
    auto var = InputVar(name);
    PADDLE_ENFORCE(var != nullptr, "Input(%s) should not be nullptr", name);
    return &var->Get<T>();
190 191 192 193
  }

  template <typename T>
  T* Output(const std::string& name) const {
194 195 196
    auto var = OutputVar(name);
    PADDLE_ENFORCE(var != nullptr, "Output(%s) should not be nullptr", name);
    return var->GetMutable<T>();
197 198 199 200 201 202 203 204
  }

  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),
205
                   [&](const std::string& sub_name) {
206
                     auto var = scope_.FindVar(sub_name);
207 208 209 210
                     PADDLE_ENFORCE(var != nullptr,
                                    "MultiInput(%s:%s) should not be nullptr",
                                    name, sub_name);
                     return &var->Get<T>();
211 212 213 214 215 216 217 218 219 220
                   });
    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),
221
                   [&](const std::string& sub_name) {
222
                     auto var = scope_.FindVar(sub_name);
223 224 225 226
                     PADDLE_ENFORCE(var != nullptr,
                                    "MultiOutput(%s:%s) should not be nullptr",
                                    name, sub_name);
                     return var->GetMutable<T>();
227 228 229 230 231
                   });
    return res;
  }

  const OperatorBase& op_;
232
  const Scope& scope_;
233 234 235 236
};

class InferShapeContext : public OperatorContext {
 public:
237
  InferShapeContext(const OperatorBase* op, const Scope& scope)
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257
      : OperatorContext(op, scope) {}
};

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

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

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

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

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

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

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

  virtual ~OpKernel() {}
};

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

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

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

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

311
  void InferShape(const Scope& scope) const {
312 313 314
    InferShape(InferShapeContext(this, scope));
  }

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

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

Y
Yu Yang 已提交
327
 protected:
328
  virtual void InferShape(const InferShapeContext& ctx) const = 0;
Q
Qiao Longfei 已提交
329 330 331 332
};

}  // namespace framework
}  // namespace paddle