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));
     };