[SOT][3.13] Use GetLocals to avoid SystemError when get locals from proxy frame (#69131)

SigureMo · web-flow · commit 40ef10fe830a · 2024-11-04T12:29:34.000+08:00
diff --git a/paddle/fluid/pybind/sot/cpython_internals.c b/paddle/fluid/pybind/sot/cpython_internals.c
@@ -171,15 +171,137 @@ void Internal_PyEvalFrameClearAndPop(PyThreadState *tstate,
   }
 }
 
-// Initialize frame free variables if needed
-static void Internal_frame_init_get_vars(_PyInterpreterFrame *frame) {
-// COPY_FREE_VARS has no quickened forms, so no need to use _PyOpcode_Deopt
-// here:
 #if PY_3_13_PLUS
+// Returns borrowed reference or NULL
+static PyObject *framelocalsproxy_getval(_PyInterpreterFrame *frame,
+                                         PyCodeObject *co,
+                                         int i) {
+  PyObject **fast = _PyFrame_GetLocalsArray(frame);
+  _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
+
+  PyObject *value = fast[i];
+  PyObject *cell = NULL;
+
+  if (value == NULL) {
+    return NULL;
+  }
+
+  if (kind == CO_FAST_FREE || kind & CO_FAST_CELL) {
+    // The cell was set when the frame was created from
+    // the function's closure.
+    assert(PyCell_Check(value));
+    cell = value;
+  }
+
+  if (cell != NULL) {
+    value = PyCell_GET(cell);
+  }
+
+  if (value == NULL) {
+    return NULL;
+  }
+
+  return value;
+}
+
+bool Internal_PyFrame_HasHiddenLocals(_PyInterpreterFrame *frame) {
+  /*
+   * This function returns if there are hidden locals introduced by PEP 709,
+   * which are the isolated fast locals for inline comprehensions
+   */
+  PyCodeObject *co = _PyFrame_GetCode(frame);
+
+  for (int i = 0; i < co->co_nlocalsplus; i++) {
+    _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
+
+    if (kind & CO_FAST_HIDDEN) {
+      PyObject *value = framelocalsproxy_getval(frame, co, i);
+
+      if (value != NULL) {
+        return true;
+      }
+    }
+  }
+
+  return false;
+}
+static PyObject *framelocalsproxy_new(PyTypeObject *type,
+                                      PyObject *args,
+                                      PyObject *kwds) {
+  if (PyTuple_GET_SIZE(args) != 1) {
+    PyErr_Format(PyExc_TypeError,
+                 "FrameLocalsProxy expected 1 argument, got %zd",
+                 PyTuple_GET_SIZE(args));
+    return NULL;
+  }
+  PyObject *item = PyTuple_GET_ITEM(args, 0);
+
+  if (!PyFrame_Check(item)) {
+    PyErr_Format(PyExc_TypeError, "expect frame, not %T", item);
+    return NULL;
+  }
+  PyFrameObject *frame = (PyFrameObject *)item;
+
+  if (kwds != NULL && PyDict_Size(kwds) != 0) {
+    PyErr_SetString(PyExc_TypeError,
+                    "FrameLocalsProxy takes no keyword arguments");
+    return 0;
+  }
+
+  PyFrameLocalsProxyObject *self =
+      (PyFrameLocalsProxyObject *)type->tp_alloc(type, 0);
+  if (self == NULL) {
+    return NULL;
+  }
+
+  ((PyFrameLocalsProxyObject *)self)->frame = (PyFrameObject *)Py_NewRef(frame);
+
+  return (PyObject *)self;
+}
+
+PyObject *Internal_PyFrameLocalsProxy_New(PyFrameObject *frame) {
+  PyObject *args = PyTuple_Pack(1, frame);
+  if (args == NULL) {
+    return NULL;
+  }
+
+  PyObject *proxy =
+      (PyObject *)framelocalsproxy_new(&PyFrameLocalsProxy_Type, args, NULL);
+  Py_DECREF(args);
+  return proxy;
+}
+
+PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame) {
+  // We should try to avoid creating the FrameObject if possible.
+  // So we check if the frame is a module or class level scope
   PyCodeObject *co = _PyFrame_GetCode(frame);
+
+  if (!(co->co_flags & CO_OPTIMIZED) &&
+      !Internal_PyFrame_HasHiddenLocals(frame)) {
+    if (frame->f_locals == NULL) {
+      // We found cases when f_locals is NULL for non-optimized code.
+      // We fill the f_locals with an empty dict to avoid crash until
+      // we find the root cause.
+      frame->f_locals = PyDict_New();
+      if (frame->f_locals == NULL) {
+        return NULL;
+      }
+    }
+    return Py_NewRef(frame->f_locals);
+  }
+
+  PyFrameObject *f = Internal_PyFrame_GetFrameObject(frame);
+
+  return Internal_PyFrameLocalsProxy_New(f);
+}
+
 #else
+
+// Initialize frame free variables if needed
+static void Internal_frame_init_get_vars(_PyInterpreterFrame *frame) {
+  // COPY_FREE_VARS has no quickened forms, so no need to use _PyOpcode_Deopt
+  // here:
   PyCodeObject *co = frame->f_code;
-#endif
   int lasti = _PyInterpreterFrame_LASTI(frame);
   if (!(lasti < 0 && _PyCode_CODE(co)->op.code == COPY_FREE_VARS &&
         PyFunction_Check(frame->f_funcobj))) {
@@ -189,21 +311,13 @@ static void Internal_frame_init_get_vars(_PyInterpreterFrame *frame) {
 
   /* Free vars have not been initialized -- Do that */
   PyObject *closure = ((PyFunctionObject *)frame->f_funcobj)->func_closure;
-#if PY_3_13_PLUS
-  int offset = PyUnstable_Code_GetFirstFree(co);
-#else
   int offset = PyCode_GetFirstFree(co);
-#endif
   for (int i = 0; i < co->co_nfreevars; ++i) {
     PyObject *o = PyTuple_GET_ITEM(closure, i);
     frame->localsplus[offset + i] = Py_NewRef(o);
   }
-// COPY_FREE_VARS doesn't have inline CACHEs, either:
-#if PY_3_13_PLUS
-  frame->instr_ptr = _PyCode_CODE(_PyFrame_GetCode(frame));
-#else
+  // COPY_FREE_VARS doesn't have inline CACHEs, either:
   frame->prev_instr = _PyCode_CODE(frame->f_code);
-#endif
 }
 
 static int Internal_frame_get_var(_PyInterpreterFrame *frame,
@@ -254,7 +368,6 @@ static int Internal_frame_get_var(_PyInterpreterFrame *frame,
   return 1;
 }
 
-#if !PY_3_13_PLUS
 PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame,
                                      int include_hidden) {
   /* Merge fast locals into f->f_locals */
@@ -337,9 +450,7 @@ PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame,
   Py_XDECREF(hidden);
   return NULL;
 }
-#endif  // !PY_3_13_PLUS
 
-#if !PY_3_13_PLUS
 int Internal_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame) {
   PyObject *locals = Internal_PyFrame_GetLocals(frame, 0);
   if (locals == NULL) {
diff --git a/paddle/fluid/pybind/sot/cpython_internals.h b/paddle/fluid/pybind/sot/cpython_internals.h
@@ -39,7 +39,9 @@ int Internal_PyInterpreterFrame_GetLine(_PyInterpreterFrame *frame);
 static int Internal_PyFrame_OpAlreadyRan(_PyInterpreterFrame *frame,
                                          int opcode,
                                          int oparg);
-#if !PY_3_13_PLUS
+#if PY_3_13_PLUS
+PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame);
+#else
 int Internal_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame);
 #endif
 PyFrameObject *Internal_PyFrame_New_NoTrack(PyCodeObject *code);
diff --git a/paddle/fluid/pybind/sot/eval_frame.c b/paddle/fluid/pybind/sot/eval_frame.c
@@ -67,7 +67,16 @@ DECLARE_PROXY_PROPERTY(f_executable)
 #else
 DECLARE_PROXY_PROPERTY(f_code)
 #endif
+#if PY_3_13_PLUS
+static PyObject *PyInterpreterFrameProxy_property_f_locals(
+    PyInterpreterFrameProxy *self, void *closure) {
+  PyObject *f_locals = Internal_PyFrame_GetLocals(self->frame);
+  Py_XINCREF(f_locals);
+  return f_locals;
+}
+#else
 DECLARE_PROXY_PROPERTY(f_locals)
+#endif
 DECLARE_PROXY_PROPERTY(f_globals)
 DECLARE_PROXY_PROPERTY(f_builtins)
 
@@ -339,7 +348,9 @@ static PyObject *_custom_eval_frame(PyThreadState *tstate,
   // original frame. So we pass a PyInterpreterFrame to
   // _PyFrame_FastToLocalsWithError directly. But this is an internal API, so we
   // copy many code from CPython project into our project.
-#if !PY_3_13_PLUS
+#if PY_3_13_PLUS
+  Internal_PyFrame_GetLocals(frame);
+#else
   if (Internal_PyFrame_FastToLocalsWithError(frame) < 0) {
     return NULL;
   }
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py
@@ -46,7 +46,7 @@ class PopJumpCond(Enum):
 
 def _get_pyopcode_cache_size() -> dict[str, int]:
     if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
-        # Cache for some opcodes, it's for Python 3.11+
+        # Cache for some opcodes, it's for Python 3.11
         # https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
         return {
             "BINARY_SUBSCR": 4,
@@ -62,7 +62,7 @@ def _get_pyopcode_cache_size() -> dict[str, int]:
             "CALL": 4,
         }
     elif sys.version_info >= (3, 12) and sys.version_info < (3, 13):
-        # Cache for some opcodes, it's for Python 3.12+
+        # Cache for some opcodes, it's for Python 3.12
         # https://github.com/python/cpython/blob/3.12/Include/internal/pycore_opcode.h#L34-L47
         return {
             "BINARY_SUBSCR": 1,
@@ -78,7 +78,31 @@ def _get_pyopcode_cache_size() -> dict[str, int]:
             "LOAD_SUPER_ATTR": 1,
             "CALL": 3,
         }
-    elif sys.version_info >= (3, 13):
+    elif sys.version_info >= (3, 13) and sys.version_info < (3, 14):
+        # Cache for some opcodes, it's for Python 3.13
+        # https://github.com/python/cpython/blob/3.13/Include/internal/pycore_opcode_metadata.h#L1598-L1618
+        return {
+            "JUMP_BACKWARD": 1,
+            "TO_BOOL": 3,
+            "BINARY_SUBSCR": 1,
+            "STORE_SUBSCR": 1,
+            "SEND": 1,
+            "UNPACK_SEQUENCE": 1,
+            "STORE_ATTR": 4,
+            "LOAD_GLOBAL": 4,
+            "LOAD_SUPER_ATTR": 1,
+            "LOAD_ATTR": 9,
+            "COMPARE_OP": 1,
+            "CONTAINS_OP": 1,
+            "POP_JUMP_IF_TRUE": 1,
+            "POP_JUMP_IF_FALSE": 1,
+            "POP_JUMP_IF_NONE": 1,
+            "POP_JUMP_IF_NOT_NONE": 1,
+            "FOR_ITER": 1,
+            "CALL": 3,
+            "BINARY_OP": 1,
+        }
+    elif sys.version_info >= (3, 14):
         raise NotImplementedError(
             f"Need to supplement cache operation code, for Python {sys.version_info}"
         )
diff --git a/test/sot/test_01_basic.py b/test/sot/test_01_basic.py
@@ -17,7 +17,6 @@
 from test_case_base import TestCaseBase
 
 import paddle
-from paddle.jit.sot.utils import strict_mode_guard
 
 
 def foo(x: int, y: paddle.Tensor):
@@ -29,19 +28,6 @@ def test_simple(self):
         self.assert_results(foo, 1, paddle.to_tensor(2))
 
 
-def numpy_add(x, y):
-    out = paddle.to_tensor(x.numpy() + y.numpy())
-    return out
-
-
-class TestNumpyAdd(TestCaseBase):
-    @strict_mode_guard(False)
-    def test_numpy_add(self):
-        x = paddle.to_tensor([2])
-        y = paddle.to_tensor([3])
-        self.assert_results(numpy_add, x, y)
-
-
 if __name__ == "__main__":
     unittest.main()
 
diff --git a/test/sot/test_numpy.py b/test/sot/test_numpy.py
@@ -21,24 +21,35 @@
 from paddle.jit.sot.utils import strict_mode_guard
 
 
-def foo(x, y):
+def numpy_add(x, y):
+    out = paddle.to_tensor(x.numpy() + y.numpy())
+    return out
+
+
+def tensor_add_numpy(x, y):
     ret = x + y
     return ret
 
 
 class TestNumpy(TestCaseBase):
+    @strict_mode_guard(False)
+    def test_numpy_add(self):
+        x = paddle.to_tensor([2])
+        y = paddle.to_tensor([3])
+        self.assert_results(numpy_add, x, y)
+
     def test_tensor_add_numpy_number(self):
         x = paddle.to_tensor([1.0])
         y = np.int64(2)
-        self.assert_results(foo, x, y)
-        self.assert_results(foo, y, x)
+        self.assert_results(tensor_add_numpy, x, y)
+        self.assert_results(tensor_add_numpy, y, x)
 
     @strict_mode_guard(False)
     def test_tensor_add_numpy_array(self):
         x = paddle.to_tensor([1.0])
         y = np.array(2.0)
-        self.assert_results(foo, x, y)
-        self.assert_results(foo, y, x)
+        self.assert_results(tensor_add_numpy, x, y)
+        self.assert_results(tensor_add_numpy, y, x)
 
 
 if __name__ == "__main__":
