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