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
33namespace tvm {
34
35/*!
36 * \brief NameSupply can be used to generate unique names.
37 */
38class 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 */
108class 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