op_info.h 2.5 KB
Newer Older
Z
zhaojiaying01 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

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. */
朔-望's avatar
朔-望 已提交
14 15 16

#pragma once

Z
zhaojiaying01 已提交
17 18
#include <string>
#include "common/log.h"
朔-望's avatar
朔-望 已提交
19
#include "common/type_define.h"
Z
zhaojiaying01 已提交
20
#include "framework/framework.pb.h"
朔-望's avatar
朔-望 已提交
21 22

namespace paddle_mobile {
朔-望's avatar
朔-望 已提交
23 24
namespace framework {

朔-望's avatar
朔-望 已提交
25 26
template <typename Dtype>
struct OpInfo {
27 28 29 30 31 32 33
  OpCreator<Dtype> creator_;
  const OpCreator<Dtype> &Creator() const {
    //    PADDLE_ENFORCE_NOT_NULL(creator_,
    //                            "Operator Creator has not been
    //                            registered");
    return creator_;
  }
朔-望's avatar
朔-望 已提交
34 35
};

朔-望's avatar
朔-望 已提交
36 37
template <typename Dtype>
class OpInfoMap;
朔-望's avatar
朔-望 已提交
38

朔-望's avatar
朔-望 已提交
39 40
template <typename Dtype>
static OpInfoMap<Dtype> *g_op_info_map = nullptr;
朔-望's avatar
朔-望 已提交
41

朔-望's avatar
朔-望 已提交
42 43 44
template <typename Dtype>
class OpInfoMap {
 public:
45
  static OpInfoMap &Instance() {
Z
zhaojiaying01 已提交
46
    LOG(paddle_mobile::kLOG_DEBUG1) << " TODO: fix bug";
47 48
    if (g_op_info_map<Dtype> == nullptr) {
      g_op_info_map<Dtype> = new OpInfoMap();
朔-望's avatar
朔-望 已提交
49
    }
50
    return *g_op_info_map<Dtype>;
Z
zhaojiaying01 已提交
51
  }
52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77

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

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

  const OpInfo<Dtype> &Get(const std::string &type) const {
    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<Dtype> *GetNullable(const std::string &type) const {
    auto it = map_.find(type);
    if (it == map_.end()) {
      return nullptr;
    } else {
      return &it->second;
朔-望's avatar
朔-望 已提交
78
    }
79
  }
朔-望's avatar
朔-望 已提交
80

81 82 83
  const std::unordered_map<std::string, OpInfo<Dtype>> &map() const {
    return map_;
  }
朔-望's avatar
朔-望 已提交
84

85 86 87
  std::unordered_map<std::string, OpInfo<Dtype>> *mutable_map() {
    return &map_;
  }
朔-望's avatar
朔-望 已提交
88

朔-望's avatar
朔-望 已提交
89
 private:
90 91
  OpInfoMap() = default;
  std::unordered_map<std::string, OpInfo<Dtype>> map_;
朔-望's avatar
朔-望 已提交
92

93
  //  DISABLE_COPY_AND_ASSIGN(OpInfoMap);
朔-望's avatar
朔-望 已提交
94 95
};

朔-望's avatar
朔-望 已提交
96 97
}  // namespace framework
}  // namespace paddle_mobile