24 #ifndef TVM_IR_NAME_SUPPLY_H_
25 #define TVM_IR_NAME_SUPPLY_H_
27 #include <tvm/ffi/reflection/registry.h>
33 #include <unordered_map>
53 NameSupplyNode(
const ffi::String& prefix, std::unordered_map<std::string, int> name_map)
54 :
prefix_(prefix), name_map(std::move(name_map)) {}
64 ffi::String
FreshName(
const ffi::String& name,
bool add_prefix =
true,
65 bool add_underscore =
true);
74 ffi::String
ReserveName(
const ffi::String& name,
bool add_prefix =
true);
83 bool ContainsName(
const ffi::String& name,
bool add_prefix =
true);
93 ffi::String add_prefix_to_name(
const ffi::String& name);
101 std::string GetUniqueName(std::string name,
bool add_underscore =
true);
104 std::unordered_map<std::string, int> name_map;
119 std::unordered_map<std::string, int> name_map = {});
127 template <
typename Iter,
typename Lambda>
129 :
NameSupply(
"", GetNameMap(begin, end, f)) {}
134 template <
typename Iter,
typename Lambda>
135 static std::unordered_map<std::string, int> GetNameMap(Iter begin, Iter end, Lambda f) {
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); }),
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()) {
153 name_map[prefix] =
std::max(name_map[prefix], std::stoi(name.substr(idx_last_first_num)));
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)
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