op_registry.h 11.8 KB
Newer Older
1 2
#pragma once

3
#include <algorithm>
4
#include <atomic>
Y
Yu Yang 已提交
5
#include <type_traits>
6 7
#include <unordered_map>
#include <unordered_set>
Q
Qiao Longfei 已提交
8
#include "paddle/framework/attr_checker.h"
9 10
#include "paddle/framework/op_desc.pb.h"
#include "paddle/framework/op_proto.pb.h"
Q
Qiao Longfei 已提交
11
#include "paddle/framework/operator.h"
12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

namespace paddle {
namespace framework {

// helper class to set attribute type
struct AttrTypeHelper {
  template <typename T>
  static void SetAttrType(AttrProto* attr);

  static Attribute GetAttrValue(const AttrDesc& attr_desc) {
    switch (attr_desc.type()) {
      case paddle::framework::AttrType::INT: {
        return attr_desc.i();
      }
      case paddle::framework::AttrType::FLOAT: {
        return attr_desc.f();
      }
      case paddle::framework::AttrType::STRING: {
        return attr_desc.s();
      }
      case paddle::framework::AttrType::INTS: {
        std::vector<int> val(attr_desc.ints_size());
        for (int i = 0; i < attr_desc.ints_size(); ++i) {
          val[i] = attr_desc.ints(i);
        }
        return val;
      }
      case paddle::framework::AttrType::FLOATS: {
        std::vector<float> val(attr_desc.floats_size());
        for (int i = 0; i < attr_desc.floats_size(); ++i) {
          val[i] = attr_desc.floats(i);
        }
        return val;
      }
      case paddle::framework::AttrType::STRINGS: {
        std::vector<std::string> val(attr_desc.strings_size());
        for (int i = 0; i < attr_desc.strings_size(); ++i) {
          val[i] = attr_desc.strings(i);
        }
        return val;
      }
    }
    PADDLE_ENFORCE(false, "Unknown OpDesc::AttrDesc::type !");
    return boost::blank();
  }
};

// this class not only make proto but also init attribute checkers.
class OpProtoAndCheckerMaker {
 public:
  OpProtoAndCheckerMaker(OpProto* proto, OpAttrChecker* op_checker)
      : proto_(proto), op_checker_(op_checker) {}

65 66
  ~OpProtoAndCheckerMaker() { CheckNoDuplicatedAttrs(); }

67
 protected:
68 69
  void AddInput(const std::string& name, const std::string& comment,
                bool multiple = false) {
70
    auto input = proto_->mutable_inputs()->Add();
71 72
    *input->mutable_name() = name;
    *input->mutable_comment() = comment;
73 74 75 76 77 78 79 80
    input->set_multiple(multiple);
    if (multiple) {
      SetHasMultipleInput();
    }
  }

  void AddInputs(const std::string& name, const std::string& comment) {
    AddInput(name, comment, true);
81 82
  }

83 84
  void AddOutput(const std::string& name, const std::string& comment,
                 bool temporary = false, bool multiple = false) {
85
    auto output = proto_->mutable_outputs()->Add();
86 87
    *output->mutable_name() = name;
    *output->mutable_comment() = comment;
88 89 90 91 92 93 94 95 96 97 98 99 100
    output->set_multiple(multiple);
    if (multiple) {
      SetHasMultipleOutput();
    }
    output->set_temporary(temporary);
    if (temporary) {
      SetHasTemporaryOutput();
    }
  }

  void AddOutputs(const std::string& name, const std::string& comment,
                  bool temporary = false) {
    AddOutput(name, comment, temporary, true);
101 102 103 104
  }

  template <typename T>
  TypedAttrChecker<T>& AddAttr(const std::string& name,
105 106
                               const std::string& comment,
                               bool generated = false) {
107
    auto attr = proto_->mutable_attrs()->Add();
108 109
    *attr->mutable_name() = name;
    *attr->mutable_comment() = comment;
110
    attr->set_generated(generated);
111 112 113 114 115 116 117 118
    AttrTypeHelper::SetAttrType<T>(attr);
    return op_checker_->AddAttrChecker<T>(name);
  }

  void AddComment(const std::string& comment) {
    *(proto_->mutable_comment()) = comment;
  }

119 120 121 122 123 124 125 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 168 169 170 171 172 173 174 175 176 177
 private:
  void SetHasMultiple(const std::string& in_out, bool* flag) {
    if (!*flag) {
      AddAttr<std::vector<int>>(in_out + "_format",
                                "The multiple index of " + in_out +
                                    "\n"
                                    R"DOC(
This attribute is used by Paddle core framework. Paddle's Op support each input
or output could be a list of variable. This attribute is used to show how that
list organized.

e.g.
  input = ["a", "b", "c", "d", "e", "f"]
  input_format = [0, 4, 5, 6]

means
  The number of all input variables this op is six, and they are segmented into
  three inputs.

  The first input is input[0:4], second is input[4:5], third is input[5:6].
)DOC",
                                /*generated*/ true);
      *flag = true;
    }
  }

  void SetHasMultipleInput() { SetHasMultiple("input", &has_multiple_input_); }
  void SetHasMultipleOutput() {
    SetHasMultiple("output", &has_multiple_output_);
  }

  void SetHasTemporaryOutput() {
    if (!has_temporary_output_) {
      AddAttr<std::vector<int>>("temporary_index",
                                R"DOC(The temporary index of output.

Not all output of Paddle Op is used by user. For faster computation, each op
could output some its internal state to other op, other op could take that
output to make compute faster.

Add a mark to which output is temporary is helpful for future optimization.
)DOC",
                                /*generated*/ true)
          .SetDefault(std::vector<int>());
      has_temporary_output_ = true;
    }
  }

  void CheckNoDuplicatedAttrs() {
    std::unordered_set<std::string> names;
    size_t cnt = 0;
    for (auto& attr : proto_->attrs()) {
      names.insert(attr.name());
      ++cnt;
    }
    PADDLE_ENFORCE(names.size() == cnt,
                   "Cannot register two attribute in same name!");
  }

178 179
  OpProto* proto_;
  OpAttrChecker* op_checker_;
180 181 182
  bool has_multiple_input_{false};
  bool has_multiple_output_{false};
  bool has_temporary_output_{false};
183 184 185
};

class OpRegistry {
Q
Qiao Longfei 已提交
186
  using OpCreator = std::function<OperatorBase*()>;
187 188 189 190

 public:
  template <typename OpType, typename ProtoMakerType>
  static void RegisterOp(const std::string& op_type) {
191 192 193
    creators()[op_type] = [] { return new OpType; };
    OpProto& op_proto = protos()[op_type];
    OpAttrChecker& op_checker = op_checkers()[op_type];
194
    ProtoMakerType(&op_proto, &op_checker);
Y
Yu Yang 已提交
195 196 197 198 199
    *op_proto.mutable_type() = op_type;
    PADDLE_ENFORCE(
        op_proto.IsInitialized(),
        "Fail to initialize %s's OpProto, because %s is not initialized",
        op_type, op_proto.InitializationErrorString());
200 201
  }

Q
Qiao Longfei 已提交
202
  static OperatorPtr CreateOp(const OpDesc& op_desc) {
203
    //! Create a OpPtr by type.
204
    std::string op_type = op_desc.type();
Q
Qiao Longfei 已提交
205
    OperatorPtr op(creators().at(op_type)());
206 207
    //! Fill op's data member. Not use constructor because it will be noising
    //! for Op developer.
Q
Qiao Longfei 已提交
208
    op->type_ = op_desc.type();
209 210 211 212 213 214
    op->inputs_.reserve((size_t)op_desc.inputs_size());
    std::copy(op_desc.inputs().begin(), op_desc.inputs().end(),
              std::back_inserter(op->inputs_));
    op->outputs_.reserve((size_t)op_desc.outputs_size());
    std::copy(op_desc.outputs().begin(), op_desc.outputs().end(),
              std::back_inserter(op->outputs_));
215 216

    //! Fill attrs, and validate attrs.
217
    for (auto& attr : op_desc.attrs()) {
Q
Qiao Longfei 已提交
218
      op->attrs_[attr.name()] = AttrTypeHelper::GetAttrValue(attr);
219
    }
Q
Qiao Longfei 已提交
220
    op_checkers().at(op_type).Check(op->attrs_);
221 222

    //! Convert Temporary variable name to an unique variable name.
223
    GenerateTempVariableName(op.get());
224 225 226

    //! Other op's custom Init for a complex Op. For simple Op, the Init
    //! method do nothing.
Q
Qiao Longfei 已提交
227
    op->Init();
228 229 230
    return op;
  }

Y
Yu Yang 已提交
231 232 233 234 235
  static std::unordered_map<std::string, OpProto>& protos() {
    static std::unordered_map<std::string, OpProto> protos_;
    return protos_;
  };

236
 private:
237
  static void GenerateTempVariableName(OperatorBase* op) {
238 239 240
    static std::atomic<size_t> gUniqId(0UL);
    for (auto& outname : op->outputs_) {
      if (outname == OperatorBase::TMP_VAR_NAME()) {
241
        outname += op->type_;
242 243 244 245 246 247
        outname += "@";
        outname += std::to_string(gUniqId.fetch_add(1));
      }
    }
  }

248 249 250 251
  static std::unordered_map<std::string, OpCreator>& creators() {
    static std::unordered_map<std::string, OpCreator> creators_;
    return creators_;
  }
252

253 254 255 256 257
  static std::unordered_map<std::string, OpAttrChecker>& op_checkers() {
    static std::unordered_map<std::string, OpAttrChecker> op_checkers_;
    return op_checkers_;
  };
};
258 259 260 261

template <typename OpType, typename ProtoMakerType>
class OpRegisterHelper {
 public:
Y
Yu Yang 已提交
262
  OpRegisterHelper(const char* op_type) {
263 264 265 266
    OpRegistry::RegisterOp<OpType, ProtoMakerType>(op_type);
  }
};

267 268 269
/**
 * check if MACRO is used in GLOBAL NAMESPACE.
 */
Y
Yu Yang 已提交
270 271 272 273 274 275
#define STATIC_ASSERT_GLOBAL_NAMESPACE(uniq_name, msg)                        \
  struct __test_global_namespace_##uniq_name##__ {};                          \
  static_assert(std::is_same<::__test_global_namespace_##uniq_name##__,       \
                             __test_global_namespace_##uniq_name##__>::value, \
                msg)

276 277 278
/**
 * Macro to Register Operator.
 */
Y
Yu Yang 已提交
279 280 281 282 283 284 285
#define REGISTER_OP(__op_type, __op_class, __op_maker_class)                 \
  STATIC_ASSERT_GLOBAL_NAMESPACE(__reg_op__##__op_type,                      \
                                 "REGISTER_OP must be in global namespace"); \
  static ::paddle::framework::OpRegisterHelper<__op_class, __op_maker_class> \
      __op_register_##__op_type##__(#__op_type);                             \
  int __op_register_##__op_type##_handle__() { return 0; }

286 287 288 289
/**
 * Macro to Register OperatorKernel.
 */
#define REGISTER_OP_KERNEL(type, DEVICE_TYPE, PlaceType, KernelType)      \
Y
Yu Yang 已提交
290
  STATIC_ASSERT_GLOBAL_NAMESPACE(                                         \
291
      __reg_op_kernel_##type##_##DEVICE_TYPE##__,                         \
Y
Yu Yang 已提交
292 293 294 295 296 297 298 299 300 301
      "REGISTER_OP_KERNEL must be in global namespace");                  \
  struct __op_kernel_register__##type##__ {                               \
    __op_kernel_register__##type##__() {                                  \
      ::paddle::framework::OperatorWithKernel::OpKernelKey key;           \
      key.place_ = PlaceType();                                           \
      ::paddle::framework::OperatorWithKernel::AllOpKernels()[#type][key] \
          .reset(new KernelType());                                       \
    }                                                                     \
  };                                                                      \
  static __op_kernel_register__##type##__ __reg_kernel_##type##__;        \
302
  int __op_kernel_register_##type##_handle_##DEVICE_TYPE##__() { return 0; }
Y
Yu Yang 已提交
303 304 305 306 307 308 309

#define REGISTER_OP_GPU_KERNEL(type, KernelType) \
  REGISTER_OP_KERNEL(type, GPU, ::paddle::platform::GPUPlace, KernelType)

#define REGISTER_OP_CPU_KERNEL(type, KernelType) \
  REGISTER_OP_KERNEL(type, CPU, ::paddle::platform::CPUPlace, KernelType)

310 311 312 313
/**
 * Macro to mark what Operator and Kernel we will use and tell the compiler to
 * link them into target.
 */
Y
Yu Yang 已提交
314 315 316 317 318 319 320 321
#define USE_OP_WITHOUT_KERNEL(op_type)                      \
  STATIC_ASSERT_GLOBAL_NAMESPACE(                           \
      __use_op_without_kernel_##op_type,                    \
      "USE_OP_WITHOUT_KERNEL must be in global namespace"); \
  extern int __op_register_##op_type##_handle__();          \
  static int __use_op_ptr_##op_type##_without_kernel__      \
      __attribute__((unused)) = __op_register_##op_type##_handle__()

Y
Yu Yang 已提交
322 323 324 325 326 327 328 329
#define USE_OP_KERNEL(op_type, DEVICE_TYPE)                               \
  STATIC_ASSERT_GLOBAL_NAMESPACE(                                         \
      __use_op_kernel_##op_type##_##DEVICE_TYPE##__,                      \
      "USE_OP_KERNEL must be in global namespace");                       \
  extern int __op_kernel_register_##op_type##_handle_##DEVICE_TYPE##__(); \
  static int __use_op_ptr_##op_type##_##DEVICE_TYPE##_kernel__            \
      __attribute__((unused)) =                                           \
          __op_kernel_register_##op_type##_handle_##DEVICE_TYPE##__()
Y
Yu Yang 已提交
330

331 332
// use Operator with only cpu kernel.
#define USE_OP_CPU(op_type)       \
Y
Yu Yang 已提交
333
  USE_OP_WITHOUT_KERNEL(op_type); \
334
  USE_OP_KERNEL(op_type, CPU)
Y
Yu Yang 已提交
335

336 337
#ifdef PADDLE_ONLY_CPU
#define USE_OP(op_type) USE_OP_CPU(op_type)
Y
Yu Yang 已提交
338
#else
339 340
#define USE_OP(op_type) \
  USE_OP_CPU(op_type);  \
Y
Yu Yang 已提交
341 342
  USE_OP_KERNEL(op_type, GPU)
#endif
343 344 345

}  // namespace framework
}  // namespace paddle