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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
|
From: Zachary DeVito <zdevito@meta.com>
Date: Wed, 11 Jan 2023 22:27:59 +0000
Subject: Fix Python 3.11 bytecode decoding in dims (#91290)
Adds a PyInstDecoder object that handles the differences in bytecode
added in 3.11. Basically some instructions have inline caches which
change the size of the instruction, so calculating the next instruction
is slightly different.
fixes #91246
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91290
Approved by: https://github.com/albanD
---
functorch/csrc/dim/dim.cpp | 105 +++++++++++++++++++++++++++++----------------
1 file changed, 68 insertions(+), 37 deletions(-)
diff --git a/functorch/csrc/dim/dim.cpp b/functorch/csrc/dim/dim.cpp
index db59983..b1c034a 100644
--- a/functorch/csrc/dim/dim.cpp
+++ b/functorch/csrc/dim/dim.cpp
@@ -21,6 +21,11 @@
#include "arena.h"
#include "python_variable_simple.h"
+#if IS_PYTHON_3_11_PLUS
+#define Py_BUILD_CORE
+#include "internal/pycore_opcode.h"
+#undef Py_BUILD_CORE
+#endif
// C++ API functions for objects to
// * construct the object, returning a ref-counted handle
@@ -1433,33 +1438,6 @@ bool relevant_op(_Py_CODEUNIT c) {
}
}
-py::object getname(PyCodeObject* code, _Py_CODEUNIT c) {
- PyObject* names = NULL;
- switch(_Py_OPCODE(c)) {
- case STORE_NAME:
- case STORE_GLOBAL:
- names = code->co_names;
- break;
- case STORE_FAST:
-#if PY_VERSION_HEX < 0x030b0000
- names = code->co_varnames;
-#else
- names = PyCode_GetVarnames(code);
-#endif
- break;
- case STORE_DEREF:
-#if PY_VERSION_HEX < 0x030b0000
- names = code->co_cellvars;
-#else
- names = PyCode_GetCellvars(code);
-#endif
- break;
- default:
- return py::object();
- }
- return py::object::steal(PySequence_GetItem(names, _Py_OPARG(c)));
-}
-
py::object create_dim(py::object name, py::handle size) {
auto d = Dim::create(std::move(name));
if (!py::is_none(size)) {
@@ -1487,10 +1465,54 @@ py::object create_dimlist(py::object name, py::handle size) {
// Python wrappers that make new reflection primitives available for older runtimes
-#if PY_VERSION_HEX < 0x030b0000
+#if !(IS_PYTHON_3_11_PLUS)
#define _PyCode_CODE(CO) ((_Py_CODEUNIT*)PyBytes_AS_STRING((CO)->co_code))
#endif
+struct PyInstDecoder {
+ PyInstDecoder(PyCodeObject* code_object, int lasti)
+ : code_object_(code_object), code_(_PyCode_CODE(code_object)), offset_(lasti / sizeof(_Py_CODEUNIT)) {}
+ void next() {
+ #if IS_PYTHON_3_11_PLUS
+ offset_ += _PyOpcode_Caches[opcode()];
+ #endif
+ offset_ += 1;
+ }
+ int opcode() {
+ auto r = _Py_OPCODE(code_[offset_]);
+ #if IS_PYTHON_3_11_PLUS
+ r = _PyOpcode_Deopt[r];
+ #endif
+ return r;
+ }
+ int oparg() {
+ return _Py_OPARG(code_[offset_]);
+ }
+
+ py::object name() {
+ py::object names;
+ switch(opcode()) {
+ case STORE_NAME:
+ case STORE_GLOBAL:
+ names = py::object::borrow(code_object_->co_names);
+ break;
+ case STORE_FAST:
+ names = py::object::steal(PyCode_GetVarnames(code_object_));
+ break;
+ case STORE_DEREF:
+ names = py::object::steal(PyCode_GetCellvars(code_object_));
+ break;
+ default:
+ return py::object();
+ }
+ return py::object::steal(PySequence_GetItem(names.ptr(), oparg()));
+ }
+private:
+ PyCodeObject* code_object_;
+ _Py_CODEUNIT* code_;
+ int offset_;
+};
+
template<py::object (*create_object)(py::object, py::handle)>
static PyObject* _dims(PyObject *self,
PyObject *const *args,
@@ -1518,15 +1540,22 @@ static PyObject* _dims(PyObject *self,
PyThreadState* state = PyThreadState_GET();
auto f = py::obj<PyFrameObject>::steal(PyThreadState_GetFrame(state));
auto c = py::obj<PyCodeObject>::steal(PyFrame_GetCode(f.ptr()));
- auto code = _PyCode_CODE(c.ptr());
- int first = PyFrame_GetLasti(f.ptr()) / 2 + 1;
- auto unpack = code[first];
- int names_start = first;
- if (relevant_op(unpack)) {
+ auto lasti = PyFrame_GetLasti(f.ptr());
+ auto decoder = PyInstDecoder(c.ptr(), lasti);
+ #if IS_PYTHON_3_11_PLUS
+ // When py3.11 adapts bytecode lasti points to the precall
+ // rather than the call instruction after it
+ if (decoder.opcode() == PRECALL) {
+ decoder.next();
+ }
+ #endif
+ decoder.next();
+
+ if (relevant_op(decoder.opcode())) {
found_ndims = 1;
- } else if (_Py_OPCODE(unpack) == UNPACK_SEQUENCE) {
- found_ndims = _Py_OPARG(unpack);
- names_start++;
+ } else if (decoder.opcode() == UNPACK_SEQUENCE) {
+ found_ndims = decoder.oparg();
+ decoder.next();
}
if (specified_ndims == -1) {
@@ -1542,11 +1571,13 @@ static PyObject* _dims(PyObject *self,
auto genobject = [&](int i) -> py::object {
py::object name;
if (i < found_ndims) {
- name = getname(c.ptr(), code[names_start + i]);
+ name = decoder.name();
}
if (!name.ptr()) {
name = py::unicode_from_format("d%d", i);
found_ndims = 0; // once we fail at finding a name, we can find any more
+ } else {
+ decoder.next();
}
return create_object(std::move(name), sizes != -1 ? py::sequence_view(py_sizes)[i] : py::handle(Py_None));
};
|
