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
34namespace tvm {
35
36/*!
37 * \brief GlobalVarSupply can be used to generate unique GlobalVars.
38 */
39class 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 */
96class 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