tvm
name_supply.h
Go to the documentation of this file.
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 
24 #ifndef TVM_IR_NAME_SUPPLY_H_
25 #define TVM_IR_NAME_SUPPLY_H_
26 
27 #include <tvm/ffi/reflection/registry.h>
28 #include <tvm/ir/expr.h>
29 
30 #include <algorithm>
31 #include <cctype>
32 #include <string>
33 #include <unordered_map>
34 #include <utility>
35 
36 namespace tvm {
37 
41 class NameSupplyNode : public Object {
42  public:
46  NameSupplyNode() = default;
47 
53  NameSupplyNode(const String& prefix, std::unordered_map<std::string, int> name_map)
54  : prefix_(prefix), name_map(std::move(name_map)) {}
55 
64  String FreshName(const String& name, bool add_prefix = true, bool add_underscore = true);
65 
73  String ReserveName(const String& name, bool add_prefix = true);
74 
82  bool ContainsName(const String& name, bool add_prefix = true);
83 
84  // Prefix for all GlobalVar names. It can be empty.
85  std::string prefix_;
86 
87  static constexpr const char* _type_key = "ir.NameSupply";
89 
90  private:
92  String add_prefix_to_name(const String& name);
93 
100  std::string GetUniqueName(std::string name, bool add_underscore = true);
101 
103  std::unordered_map<std::string, int> name_map;
104 };
105 
110 class NameSupply : public ObjectRef {
111  public:
117  TVM_DLL explicit NameSupply(const String& prefix = "",
118  std::unordered_map<std::string, int> name_map = {});
119 
126  template <typename Iter, typename Lambda>
127  TVM_DLL explicit NameSupply(Iter begin, Iter end, Lambda f)
128  : NameSupply("", GetNameMap(begin, end, f)) {}
129 
131 
132  private:
133  template <typename Iter, typename Lambda>
134  static std::unordered_map<std::string, int> GetNameMap(Iter begin, Iter end, Lambda f) {
135  // static_assert is more reader-friendly than SFINAE when template specialization is not needed.
136  static_assert(std::is_convertible<decltype(f(*begin)), std::string>::value,
137  "Lambda f must has a signature of [?](*it) -> string {}");
138  std::unordered_map<std::string, int> name_map;
139  for (auto it = begin; it != end; ++it) {
140  const std::string& name = f(*it);
141  const size_t idx_last_first_num = std::distance(
142  std::find_if(name.rbegin(), name.rend(), [](char c) { return !std::isdigit(c); }),
143  name.rend());
144  // name = {O = others}{D = consecutive digits}
145  // let O -> prefix;
146  std::string prefix = name.substr(0, idx_last_first_num);
147  ICHECK(prefix.size() > 0 && std::isalpha(prefix[0])) << "Invalid variable name: " << name;
148  if (0 == name_map.count(prefix)) name_map[prefix] = 0;
149  if (idx_last_first_num < name.size()) { // has some digits.
150  // let D's nearest natural number -> idx;
151  // note: stoul("000123") = 123;
152  name_map[prefix] = std::max(name_map[prefix], std::stoi(name.substr(idx_last_first_num)));
153  }
154  }
155  return name_map;
156  }
157 };
158 
159 } // namespace tvm
160 
161 #endif // TVM_IR_NAME_SUPPLY_H_
NameSupply can be used to generate unique names.
Definition: name_supply.h:41
NameSupplyNode(const String &prefix, std::unordered_map< std::string, int > name_map)
Constructor.
Definition: name_supply.h:53
String FreshName(const String &name, bool add_prefix=true, bool add_underscore=true)
Generates a unique name from this NameSupply.
TVM_DECLARE_FINAL_OBJECT_INFO(NameSupplyNode, Object)
static constexpr const char * _type_key
Definition: name_supply.h:87
String ReserveName(const String &name, bool add_prefix=true)
Reserves an existing name with this NameSupply.
std::string prefix_
Definition: name_supply.h:85
NameSupplyNode()=default
Empty constructor. Needed by the TVM_REGISTER_NODE_TYPE macro.
bool ContainsName(const String &name, bool add_prefix=true)
Checks if this NameSupply already generated a name.
Managed reference class to NameSupplyNode.
Definition: name_supply.h:110
NameSupply(Iter begin, Iter end, Lambda f)
Construct NameSupply with a name map created from the given iterator range and the functor.
Definition: name_supply.h:127
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(NameSupply, ObjectRef, NameSupplyNode)
NameSupply(const String &prefix="", std::unordered_map< std::string, int > name_map={})
Constructor.
Base expr nodes in TVM.
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
PrimExpr max(PrimExpr a, PrimExpr b, Span span=Span())
take maximum of two values