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/name_supply.h |
22 | * \brief NameSupply that can be used to generate unique variable names. |
23 | */ |
24 | #ifndef TVM_IR_NAME_SUPPLY_H_ |
25 | #define TVM_IR_NAME_SUPPLY_H_ |
26 | |
27 | #include <string> |
28 | #include <unordered_map> |
29 | #include <utility> |
30 | |
31 | #include "tvm/ir/expr.h" |
32 | |
33 | namespace tvm { |
34 | |
35 | /*! |
36 | * \brief NameSupply can be used to generate unique names. |
37 | */ |
38 | class NameSupplyNode : public Object { |
39 | public: |
40 | /*! |
41 | * \brief Empty constructor. Needed by the TVM_REGISTER_NODE_TYPE macro. |
42 | */ |
43 | NameSupplyNode() = default; |
44 | |
45 | /*! |
46 | * \brief Constructor. |
47 | * \param prefix The prefix to be used with this NameSupply. |
48 | * \param name_map The map used to guarantee uniqueness. |
49 | */ |
50 | NameSupplyNode(const String& prefix, std::unordered_map<std::string, int> name_map) |
51 | : prefix_(prefix), name_map(std::move(name_map)) {} |
52 | |
53 | /*! |
54 | * \brief Generates a unique name from this NameSupply. |
55 | * \param name The name from which the generated name is derived. |
56 | * \param add_prefix If set to true, then the prefix of this NameSupply will be prepended to the |
57 | * name. \return A unique name. |
58 | */ |
59 | String FreshName(const String& name, bool add_prefix = true); |
60 | |
61 | /*! |
62 | * \brief Reserves an existing name with this NameSupply. |
63 | * \param name The name to be reserved. |
64 | * \param add_prefix If set to true, then the prefix of this NameSupply will be prepended to the |
65 | * name before reserving it. \return The name that was reserved with the NameSupply. It can be |
66 | * different if a prefix is added. |
67 | */ |
68 | String ReserveName(const String& name, bool add_prefix = true); |
69 | |
70 | /*! |
71 | * \brief Checks if this NameSupply already generated a name. |
72 | * \param name The name to check. |
73 | * \param add_prefix If set to true, then the prefix of this NameSupply will be prepended to the |
74 | * name before checking for it. \return True if the name has already been generated. False |
75 | * otherwise. |
76 | */ |
77 | bool ContainsName(const String& name, bool add_prefix = true); |
78 | |
79 | void VisitAttrs(AttrVisitor* v) {} |
80 | |
81 | // Prefix for all GlobalVar names. It can be empty. |
82 | std::string prefix_; |
83 | |
84 | static constexpr const char* _type_key = "NameSupply" ; |
85 | static constexpr const bool _type_has_method_sequal_reduce = false; |
86 | static constexpr const bool _type_has_method_shash_reduce = false; |
87 | TVM_DECLARE_FINAL_OBJECT_INFO(NameSupplyNode, Object); |
88 | |
89 | private: |
90 | /*! \brief Helper function to add the NameSupply prefix to the name. */ |
91 | String add_prefix_to_name(const String& name); |
92 | |
93 | /*! |
94 | * \brief Function that will generate a unique name. |
95 | * \param name The name to be used as a base. |
96 | * \return A unique name. |
97 | */ |
98 | std::string GetUniqueName(std::string name); |
99 | |
100 | /*! \brief A map that is used to generate unique names. */ |
101 | std::unordered_map<std::string, int> name_map; |
102 | }; |
103 | |
104 | /*! |
105 | * \brief Managed reference class to NameSupplyNode. |
106 | * \sa NameSupplyNode |
107 | */ |
108 | class NameSupply : public ObjectRef { |
109 | public: |
110 | /*! |
111 | * \brief Constructor. |
112 | * \param prefix The prefix to be used with this NameSupply. |
113 | * \param name_map An optional map. |
114 | */ |
115 | TVM_DLL explicit NameSupply(const String& prefix, |
116 | std::unordered_map<std::string, int> name_map = {}); |
117 | |
118 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(NameSupply, ObjectRef, NameSupplyNode); |
119 | }; |
120 | |
121 | } // namespace tvm |
122 | |
123 | #endif // TVM_IR_NAME_SUPPLY_H_ |
124 | |