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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
#include <c10/util/flat_hash_map.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/python_dimname.h>
#include <torch/csrc/utils/python_strings.h>
namespace torch {
struct InternedStringsTable {
InternedStringsTable() = default;
// NOLINTNEXTLINE(bugprone-exception-escape)
~InternedStringsTable();
InternedStringsTable(const InternedStringsTable&) = delete;
InternedStringsTable& operator=(InternedStringsTable const&) = delete;
InternedStringsTable(InternedStringsTable&&) = delete;
InternedStringsTable& operator=(InternedStringsTable&&) = delete;
std::optional<at::Dimname> lookup(PyObject* obj);
// Precondition: obj is an interned python string.
void addMapping(PyObject* obj, at::Dimname dimname);
private:
ska::flat_hash_map<PyObject*, at::Dimname> py_interned_string_to_dimname_;
};
InternedStringsTable kPyInternedStringToDimname;
// NOLINTNEXTLINE(bugprone-exception-escape)
InternedStringsTable::~InternedStringsTable() {
// If python is already dead, leak the wrapped python objects
if (Py_IsInitialized()) {
pybind11::gil_scoped_acquire gil;
for (auto it = py_interned_string_to_dimname_.begin();
it != py_interned_string_to_dimname_.end();
++it) {
// See Note [References to python interned strings]
Py_DECREF(it->first);
}
}
}
std::optional<at::Dimname> InternedStringsTable::lookup(PyObject* obj) {
auto it = py_interned_string_to_dimname_.find(obj);
if (it == py_interned_string_to_dimname_.end()) {
return std::nullopt;
}
return it->second;
}
void InternedStringsTable::addMapping(PyObject* obj, at::Dimname dimname) {
// Note [References to python interned strings]
// If a Python interned string has no references to it, then it gets
// deallocated, invalidating this mapping. Let's immortalize the string by
// holding a refcount to it and releasing it in the destructor
Py_INCREF(obj);
py_interned_string_to_dimname_.emplace(obj, dimname);
}
} // namespace torch
bool THPUtils_checkDimname(PyObject* obj) {
return obj == Py_None || THPUtils_checkString(obj);
}
// To avoid ambiguity with IntArrayRef, we parse obj as a DimnameList if
// it is a list or tuple and its first elt is a Dimname
bool THPUtils_checkDimnameList(PyObject* obj) {
auto tuple = PyTuple_Check(obj);
if (!tuple && !PyList_Check(obj)) {
return false;
}
// NOLINTNEXTLINE(bugprone-branch-clone)
const auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj);
if (size == 0) {
return true;
}
PyObject* first_elt =
tuple ? PyTuple_GET_ITEM(obj, 0) : PyList_GET_ITEM(obj, 0);
return THPUtils_checkDimname(first_elt);
}
at::Dimname THPDimname_parse(PyObject* obj) {
if (obj == Py_None) {
return at::Dimname::wildcard();
}
TORCH_CHECK_TYPE(
THPUtils_checkString(obj),
"expected None or string for Dimname but got ",
Py_TYPE(obj)->tp_name);
if (!THPUtils_isInterned(obj)) {
// internStringInPlace decrefs obj and increfs the result. Because we're
// not actually returning the result to the user, we need to undo these.
// See
// https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_InternInPlace
Py_INCREF(obj);
THPUtils_internStringInPlace(&obj);
Py_DECREF(obj);
}
auto maybeDimname = torch::kPyInternedStringToDimname.lookup(obj);
if (maybeDimname) {
return *maybeDimname;
}
const auto name = THPUtils_unpackString(obj);
auto dimname = at::Dimname::fromSymbol(at::Symbol::dimname(name));
torch::kPyInternedStringToDimname.addMapping(obj, dimname);
return dimname;
}
|
