op_info.h 3.0 KB
Newer Older
Y
Yu Yang 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

L
Luo Tao 已提交
3 4 5
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
Y
Yu Yang 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yu Yang 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
Y
Yu Yang 已提交
14 15 16 17 18 19

#pragma once
#include <functional>
#include <map>
#include <string>
#include <unordered_map>
20

Y
Yi Wang 已提交
21 22 23
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/type_defs.h"
#include "paddle/fluid/platform/macros.h"
Y
Yu Yang 已提交
24 25 26 27

namespace paddle {
namespace framework {

28 29 30 31 32 33
class InferShapeBase {
 public:
  virtual ~InferShapeBase() = default;
  virtual void operator()(InferShapeContext*) const = 0;
};

Y
Yu Yang 已提交
34 35
struct OpInfo {
  OpCreator creator_;
Y
Yu Yang 已提交
36
  GradOpMakerFN grad_op_maker_;
37
  proto::OpProto* proto_{nullptr};
38
  OpAttrChecker* checker_{nullptr};
Y
Yu Yang 已提交
39
  InferVarTypeFN infer_var_type_;
40
  InferShapeFN infer_shape_;
Y
Yu Yang 已提交
41 42 43 44 45

  bool HasOpProtoAndChecker() const {
    return proto_ != nullptr && checker_ != nullptr;
  }

46
  const proto::OpProto& Proto() const {
Y
Yu Yang 已提交
47 48 49 50 51 52 53 54 55 56 57
    PADDLE_ENFORCE_NOT_NULL(proto_, "Operator Proto has not been registered");
    PADDLE_ENFORCE(proto_->IsInitialized(),
                   "Operator Proto must be initialized in op info");
    return *proto_;
  }

  const OpCreator& Creator() const {
    PADDLE_ENFORCE_NOT_NULL(creator_,
                            "Operator Creator has not been registered");
    return creator_;
  }
Y
Yu Yang 已提交
58 59 60 61 62 63

  const GradOpMakerFN& GradOpMaker() const {
    PADDLE_ENFORCE_NOT_NULL(grad_op_maker_,
                            "Operator GradOpMaker has not been registered.");
    return grad_op_maker_;
  }
F
fengjiayi 已提交
64 65

  const OpAttrChecker* Checker() const { return checker_; }
Y
Yu Yang 已提交
66 67
};

Y
Yu Yang 已提交
68 69 70 71 72 73 74 75 76 77 78 79 80 81
class OpInfoMap {
 public:
  static OpInfoMap& Instance();

  bool Has(const std::string& op_type) const {
    return map_.find(op_type) != map_.end();
  }

  void Insert(const std::string& type, const OpInfo& info) {
    PADDLE_ENFORCE(!Has(type), "Operator %s has been registered", type);
    map_.insert({type, info});
  }

  const OpInfo& Get(const std::string& type) const {
82 83 84 85 86 87 88
    auto op_info_ptr = GetNullable(type);
    PADDLE_ENFORCE_NOT_NULL(op_info_ptr, "Operator %s has not been registered",
                            type);
    return *op_info_ptr;
  }

  const OpInfo* GetNullable(const std::string& type) const {
Y
Yu Yang 已提交
89
    auto it = map_.find(type);
90 91 92 93 94
    if (it == map_.end()) {
      return nullptr;
    } else {
      return &it->second;
    }
Y
Yu Yang 已提交
95 96
  }

97 98 99
  const std::unordered_map<std::string, OpInfo>& map() const { return map_; }

  std::unordered_map<std::string, OpInfo>* mutable_map() { return &map_; }
Y
Yu Yang 已提交
100 101 102

 private:
  OpInfoMap() = default;
103
  std::unordered_map<std::string, OpInfo> map_;
D
format  
dongzhihong 已提交
104 105

  DISABLE_COPY_AND_ASSIGN(OpInfoMap);
Y
Yu Yang 已提交
106
};
Y
Yu Yang 已提交
107 108 109

}  // namespace framework
}  // namespace paddle