1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file tvm/ir/global_var_supply.h |
22 | * \brief GlobalVarSupply that can be used to generate unique \class GlobalVar. |
23 | */ |
24 | #ifndef TVM_IR_GLOBAL_VAR_SUPPLY_H_ |
25 | #define TVM_IR_GLOBAL_VAR_SUPPLY_H_ |
26 | |
27 | #include <string> |
28 | #include <unordered_map> |
29 | |
30 | #include "tvm/ir/expr.h" |
31 | #include "tvm/ir/module.h" |
32 | #include "tvm/ir/name_supply.h" |
33 | |
34 | namespace tvm { |
35 | |
36 | /*! |
37 | * \brief GlobalVarSupply can be used to generate unique GlobalVars. |
38 | */ |
39 | class GlobalVarSupplyNode : public Object { |
40 | public: |
41 | /*! |
42 | * \brief Empty constructor. Will use an empty NameSupply. |
43 | */ |
44 | GlobalVarSupplyNode() : GlobalVarSupplyNode(NameSupply("" )) {} |
45 | |
46 | /*! |
47 | * \brief Constructor. |
48 | * \param name_supply The NameSupply to use for generating the names of fresh GlobalVars. |
49 | * \param name_to_var_map An optional map. |
50 | */ |
51 | explicit GlobalVarSupplyNode(NameSupply name_supply, |
52 | std::unordered_map<std::string, GlobalVar> name_to_var_map = {}); |
53 | |
54 | /*! |
55 | * \brief Generates a unique GlobalVar from this supply. |
56 | * \param name The name from which the name of the GlobalVar is derived. |
57 | * \param add_prefix If set to true, then the prefix of the contained NameSupply will be prepended |
58 | * to the name. \return A unique GlobalVar. |
59 | */ |
60 | GlobalVar FreshGlobal(String name, bool add_prefix = true); |
61 | |
62 | /*! |
63 | * \brief Looks up for a GlobalVar with the given name in this supply. |
64 | * If no entry is found, creates one, places it in the cache and returns it. |
65 | * \param name The name of the GlobalVar to search for. |
66 | * \param add_prefix If set to true, the prefix of the contained NameSupply will be prepended to |
67 | * the name before performing the search. \return A cached GlobalVar. |
68 | */ |
69 | GlobalVar UniqueGlobalFor(const String& name, bool add_prefix = true); |
70 | |
71 | /*! |
72 | * \brief Reserves an existing GlobalVar with this supply. |
73 | * \param var The GlobalVar to be registered. |
74 | * \param allow_conflict Allow conflict with other GlobalVars that have the same name. |
75 | */ |
76 | void ReserveGlobalVar(const GlobalVar& var, bool allow_conflict = false); |
77 | |
78 | void VisitAttrs(AttrVisitor* v) {} |
79 | |
80 | /*! \brief The NameSupply used to generate unique name hints to GlobalVars. */ |
81 | NameSupply name_supply_; |
82 | |
83 | static constexpr const char* _type_key = "GlobalVarSupply" ; |
84 | static constexpr const bool _type_has_method_sequal_reduce = false; |
85 | static constexpr const bool _type_has_method_shash_reduce = false; |
86 | TVM_DECLARE_FINAL_OBJECT_INFO(GlobalVarSupplyNode, Object); |
87 | |
88 | private: |
89 | std::unordered_map<std::string, GlobalVar> name_to_var_map_; |
90 | }; |
91 | |
92 | /*! |
93 | * \brief Managed reference class to GlobalVarSupplyNode. |
94 | * \sa GlobalVarSupplyNode |
95 | */ |
96 | class GlobalVarSupply : public ObjectRef { |
97 | public: |
98 | /*! |
99 | * \brief Constructor. |
100 | * \param name_supply The NameSupply to be used when generating new GlobalVars. |
101 | * \param name_to_var_map An optional map. |
102 | */ |
103 | TVM_DLL explicit GlobalVarSupply(const NameSupply& name_supply, |
104 | std::unordered_map<std::string, GlobalVar> name_to_var_map = {}); |
105 | |
106 | /*! |
107 | * \brief Constructs a supply from an array of IRModules. GlobalVars generated by this supply are |
108 | * guaranteed not to conflict with any GlobalVars that belong to the modules. \param modules Array |
109 | * of IRModules. |
110 | */ |
111 | TVM_DLL explicit GlobalVarSupply(const Array<IRModule>& modules); |
112 | |
113 | /*! |
114 | * \brief Constructs a GlobalVarSupply from an IRModule. GlobalVars generated by this supply are |
115 | * guaranteed not to conflict with GlobalVars that belong to the modules. \param module The |
116 | * IRModule. |
117 | */ |
118 | TVM_DLL explicit GlobalVarSupply(const IRModule module); |
119 | |
120 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(GlobalVarSupply, ObjectRef, GlobalVarSupplyNode); |
121 | }; |
122 | |
123 | } // namespace tvm |
124 | |
125 | #endif // TVM_IR_GLOBAL_VAR_SUPPLY_H_ |
126 | |