xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/tensorexpr/unique_name_manager.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <torch/csrc/jit/tensorexpr/unique_name_manager.h>
2 
3 #include <torch/csrc/jit/tensorexpr/ir.h>
4 #include <cctype>
5 
6 namespace torch::jit::tensorexpr {
7 
get_unique_name(const VarPtr & v)8 const std::string& UniqueNameManager::get_unique_name(const VarPtr& v) {
9   // Find if we have already encountered this variable.
10   auto iter = unique_name_mapping_.find(v);
11   if (iter != unique_name_mapping_.end()) {
12     return iter->second;
13   }
14 
15   // First use the name_hint as a prefix to check if there is another name
16   // with the same prefix.
17   std::string name_hint = v->name_hint();
18   if (name_hint.empty()) {
19     name_hint = "v";
20   } else if (std::isdigit(name_hint[0])) {
21     name_hint = "v" + name_hint;
22   }
23   int& count = unique_name_count_[name_hint];
24   while (true) {
25     // Even if with a new count, this name might already be used. For example
26     // ("x", 1) could collidewith ("x_1", 0)
27     int count_v = count++;
28     std::string unique_name = name_hint;
29     if (count_v > 0) {
30       unique_name += "_" + std::to_string(count_v);
31     }
32     if (all_unique_names_.count(unique_name) == 0) {
33       all_unique_names_.insert(unique_name);
34       auto result = unique_name_mapping_.insert(std::make_pair(v, unique_name));
35       return result.first->second;
36     }
37   }
38 }
39 
get_unique_name(const VarHandle & v)40 const std::string& UniqueNameManager::get_unique_name(const VarHandle& v) {
41   return get_unique_name(v.node());
42 }
43 
44 } // namespace torch::jit::tensorexpr
45