• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2020 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef MINDSPORE_CORE_OPS_PRIMITIVE_C_H_
18 #define MINDSPORE_CORE_OPS_PRIMITIVE_C_H_
19 #include <map>
20 #include <memory>
21 #include <string>
22 #include <vector>
23 #include "ir/primitive.h"
24 #include "ir/value.h"
25 #include "utils/hash_map.h"
26 namespace mindspore {
27 namespace ops {
28 /// \brief PrimitiveC defines the base class for c++ operators.
29 using PrimitiveC = Primitive;
30 using OpPrimCDefineFunc = std::function<std::shared_ptr<PrimitiveC>()>;
31 /// \brief OpPrimCRegister defines the singleton to save c++ operators.
32 class MS_CORE_API OpPrimCRegister {
33  public:
34   /// \brief Destructor of OpPrimCRegister.
~OpPrimCRegister()35   ~OpPrimCRegister() {}
36 
37   /// \brief Get the OpPrimCRegister singleton.
38   ///
39   /// \return The OpPrimCRegister singleton.
40   static OpPrimCRegister &GetInstance();
41 
42   /// \brief Get PrimCMap of the OpPrimCRegister singleton.
43   ///
44   /// \return The PrimCMap of the OpPrimCRegister singleton.
45   const HashMap<std::string, OpPrimCDefineFunc> &GetPrimCMap() const;
46 
47   /// \brief Add an element into the PrimCMap of the OpPrimCRegister singleton.
48   ///
49   /// param[in] kname The name of the input end side operator.
50   /// param[in] fn The input end side operator.
51   void SetPrimCMap(const std::string &kname, const OpPrimCDefineFunc &fn);
52 
53  private:
OpPrimCRegister()54   OpPrimCRegister() {}
55   HashMap<std::string, OpPrimCDefineFunc> op_primc_fns_;
56 };
57 
58 /// \brief OpPrimCRegisterHelper defines the helper class for the OpPrimCRegister singleton.
59 class MS_CORE_API OpPrimCRegisterHelper {
60  public:
61   /// \brief Constructor for OpPrimCRegisterHelper.
62   ///
63   /// param[in] kname The name of the input end side operator.
64   /// param[in] fn The input end side operator.
OpPrimCRegisterHelper(const std::string & kname,const OpPrimCDefineFunc & fn)65   OpPrimCRegisterHelper(const std::string &kname, const OpPrimCDefineFunc &fn) {
66     OpPrimCRegister::GetInstance().SetPrimCMap(kname, fn);
67     (void)id_;  // make compiler happy on macos
68   }
69 
70   /// Destructor of OpPrimCRegisterHelper.
71   ~OpPrimCRegisterHelper() = default;
72 
73  private:
74   int id_{0};
75 };
76 
77 #define REGISTER_PRIMITIVE_C(kname, primc)                    \
78   std::shared_ptr<PrimitiveC> GetDefaultPrimC##primc() {      \
79     primc out;                                                \
80     return std::dynamic_pointer_cast<PrimitiveC>(out.impl()); \
81   }                                                           \
82   OpPrimCRegisterHelper primc_gen_##kname(kname, GetDefaultPrimC##primc);
83 }  // namespace ops
84 }  // namespace mindspore
85 #endif  // MINDSPORE_CORE_OPS_PRIMITIVE_C_H_
86