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 ffi::String& prefix, std::unordered_map<std::string, int> name_map)
54  : prefix_(prefix), name_map(std::move(name_map)) {}
55 
64  ffi::String FreshName(const ffi::String& name, bool add_prefix = true,
65  bool add_underscore = true);
66 
74  ffi::String ReserveName(const ffi::String& name, bool add_prefix = true);
75 
83  bool ContainsName(const ffi::String& name, bool add_prefix = true);
84 
85  // Prefix for all GlobalVar names. It can be empty.
86  std::string prefix_;
87 
88  static constexpr const bool _type_mutable = true;
90 
91  private:
93  ffi::String add_prefix_to_name(const ffi::String& name);
94 
101  std::string GetUniqueName(std::string name, bool add_underscore = true);
102 
104  std::unordered_map<std::string, int> name_map;
105 };
106 
111 class NameSupply : public ObjectRef {
112  public:
118  TVM_DLL explicit NameSupply(const ffi::String& prefix = "",
119  std::unordered_map<std::string, int> name_map = {});
120 
127  template <typename Iter, typename Lambda>
128  TVM_DLL explicit NameSupply(Iter begin, Iter end, Lambda f)
129  : NameSupply("", GetNameMap(begin, end, f)) {}
130 
132 
133  private:
134  template <typename Iter, typename Lambda>
135  static std::unordered_map<std::string, int> GetNameMap(Iter begin, Iter end, Lambda f) {
136  // static_assert is more reader-friendly than SFINAE when template specialization is not needed.
137  static_assert(std::is_convertible<decltype(f(*begin)), std::string>::value,
138  "Lambda f must has a signature of [?](*it) -> string {}");
139  std::unordered_map<std::string, int> name_map;
140  for (auto it = begin; it != end; ++it) {
141  const std::string& name = f(*it);
142  const size_t idx_last_first_num = std::distance(
143  std::find_if(name.rbegin(), name.rend(), [](char c) { return !std::isdigit(c); }),
144  name.rend());
145  // name = {O = others}{D = consecutive digits}
146  // let O -> prefix;
147  std::string prefix = name.substr(0, idx_last_first_num);
148  ICHECK(prefix.size() > 0 && std::isalpha(prefix[0])) << "Invalid variable name: " << name;
149  if (0 == name_map.count(prefix)) name_map[prefix] = 0;
150  if (idx_last_first_num < name.size()) { // has some digits.
151  // let D's nearest natural number -> idx;
152  // note: stoul("000123") = 123;
153  name_map[prefix] = std::max(name_map[prefix], std::stoi(name.substr(idx_last_first_num)));
154  }
155  }
156  return name_map;
157  }
158 };
159 
160 } // namespace tvm
161 
162 #endif // TVM_IR_NAME_SUPPLY_H_
NameSupply can be used to generate unique names.
Definition: name_supply.h:41
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.NameSupply", NameSupplyNode, Object)
ffi::String ReserveName(const ffi::String &name, bool add_prefix=true)
Reserves an existing name with this NameSupply.
NameSupplyNode(const ffi::String &prefix, std::unordered_map< std::string, int > name_map)
Constructor.
Definition: name_supply.h:53
bool ContainsName(const ffi::String &name, bool add_prefix=true)
Checks if this NameSupply already generated a name.
std::string prefix_
Definition: name_supply.h:86
static constexpr const bool _type_mutable
Definition: name_supply.h:88
NameSupplyNode()=default
Empty constructor. Needed by the TVM_REGISTER_NODE_TYPE macro.
ffi::String FreshName(const ffi::String &name, bool add_prefix=true, bool add_underscore=true)
Generates a unique name from this NameSupply.
Managed reference class to NameSupplyNode.
Definition: name_supply.h:111
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:128
NameSupply(const ffi::String &prefix="", std::unordered_map< std::string, int > name_map={})
Constructor.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NameSupply, ObjectRef, NameSupplyNode)
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