pytorch

1
#include <torch/csrc/jit/frontend/name_mangler.h>
2

3
namespace torch::jit {
4

5
c10::QualifiedName NameMangler::mangle(const c10::QualifiedName& name) {
6
  static const std::string manglePrefix = "___torch_mangle_";
7
  std::vector<std::string> atoms = name.atoms();
8

9
  // Search for an already-existing mangle namespace.
10
  // If the name is already mangled, just bump the integer.
11
  for (auto& atom : atoms) {
12
    auto pos = atom.find(manglePrefix);
13
    if (pos != std::string::npos) {
14
      auto num = atom.substr(pos + manglePrefix.size());
15
      // current mangle index in the name
16
      size_t num_i = std::stoi(num);
17
      // bump the mangleIndex_ to num_i + 1
18
      mangleIndex_ = std::max(mangleIndex_, num_i + 1);
19
      std::string newAtomPrefix;
20
      newAtomPrefix.reserve(atom.size());
21
      // Append the part of the name up to the end of the prefix
22
      newAtomPrefix.append(atom, 0, pos);
23
      newAtomPrefix.append(manglePrefix);
24
      atom = newAtomPrefix + c10::to_string(mangleIndex_++);
25
      // increment mangleIndex_ until the type is not defined
26
      return c10::QualifiedName(atoms);
27
    }
28
  }
29

30
  // Otherwise add a mangle namespace right before the basename
31
  TORCH_INTERNAL_ASSERT(!atoms.empty());
32
  atoms.insert(atoms.end() - 1, manglePrefix + c10::to_string(mangleIndex_++));
33
  return c10::QualifiedName(atoms);
34
}
35

36
} // namespace torch::jit
37