1 #pragma once 2 3 #include <string> 4 #include <unordered_map> 5 #include <unordered_set> 6 7 #include <torch/csrc/Export.h> 8 #include <torch/csrc/jit/tensorexpr/fwd_decls.h> 9 10 namespace torch::jit::tensorexpr { 11 12 class VarHandle; 13 class Var; 14 15 using VarNameMap = std::unordered_map<VarPtr, std::string>; 16 17 // A manager to get unique names from vars. 18 // It starts with the name hints of the var and append "_" + $counter until it 19 // hits a unique name. 20 class TORCH_API UniqueNameManager { 21 public: 22 const std::string& get_unique_name(const VarHandle& v); 23 24 const std::string& get_unique_name(const VarPtr& v); 25 26 private: 27 friend class ScopedVarName; 28 VarNameMap unique_name_mapping_; 29 std::unordered_map<std::string, int> unique_name_count_; 30 std::unordered_set<std::string> all_unique_names_; 31 }; 32 33 } // namespace torch::jit::tensorexpr 34