Lazily install the CinderX frame evaluator function

Summary:
Doing this on module init is risky, we have features in CinderX that don't need
the frame evaluator (e.g. Parallel GC) and so we don't want it to be
unconditionally enabled.  We know we have issues with PGO not being applied on
the frame evaluator in all builds, so best to only turn it on when we know its
going to be needed.

The cases where this has to be turned on:

* The JIT is enabled and is going to start compiling code.
* The Static Python bytecode compiler is used.
* Shadowcode is enabled (Python 3.10 only).

Reviewed By: DinoV, jbower-fb

Differential Revision: D92628744

fbshipit-source-id: 6acb7c4c9bc811b6ab42454d68fbadf6652d8d0b

Author: alexmalyshev

CMakeLists.txt

@@ -254,7 +254,11 @@ target_link_libraries(cached-properties PRIVATE common borrowed)
 ########################################
 # Interpreter/
 
-set(INTERP_SOURCES ${PROJECT_SOURCE_DIR}/Interpreter/cinder_opcode.c ${PROJECT_SOURCE_DIR}/Interpreter/iter_helpers.c)
+set(INTERP_SOURCES
+  ${PROJECT_SOURCE_DIR}/Interpreter/cinder_opcode.c
+  ${PROJECT_SOURCE_DIR}/Interpreter/interpreter_base.cpp
+  ${PROJECT_SOURCE_DIR}/Interpreter/iter_helpers.c
+)
 
 if (${PY_VERSION} EQUAL 3.12 OR ${PY_VERSION} EQUAL 3.14 OR ${PY_VERSION} EQUAL 3.15)
   set(SOURCE_CINDERX_OPCODE_TARGETS_H ${PROJECT_SOURCE_DIR}/Interpreter/${PY_VERSION}/cinderx_opcode_targets.h)

cinderx/Interpreter/interpreter.h

@@ -4,16 +4,19 @@
 
 #include "cinderx/python.h"
 
-#include "cinderx/Common/extra-py-flags.h"
-#include "cinderx/UpstreamBorrow/borrowed.h"
+// Exporting Ci_PyFunction_Vectorcall.
 #include "cinderx/module_c_state.h"
-#if PY_VERSION_HEX >= 0x030D0000
-#include "internal/pycore_function.h"
-#endif
+
+#include <stdbool.h>
+#include <stdint.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+/*
+ * The CinderX frame evaluator function (interpreter loop).
+ */
 #if PY_VERSION_HEX < 0x030C0000
 PyObject* _Py_HOT_FUNCTION
 Ci_EvalFrame(PyThreadState* tstate, PyFrameObject* f, int throwflag);
@@ -51,17 +54,17 @@ PyObject* Ci_PyFunction_CallStatic(
  * This is a different function for Static Python functions versus "normal"
  * Python functions.
  */
-static inline vectorcallfunc getInterpretedVectorcall(
-    [[maybe_unused]] const PyFunctionObject* func) {
-#ifdef ENABLE_INTERPRETER_LOOP
-  const PyCodeObject* code = (const PyCodeObject*)(func->func_code);
-  return (code->co_flags & CI_CO_STATICALLY_COMPILED)
-      ? Ci_StaticFunction_Vectorcall
-      : Ci_PyFunction_Vectorcall;
-#else
-  return Ci_PyFunction_Vectorcall;
-#endif
-}
+vectorcallfunc getInterpretedVectorcall(const PyFunctionObject* func);
+
+/*
+ * Install the CinderX frame evaluator function into the runtime.
+ */
+int Ci_InitFrameEvalFunc();
+
+/*
+ * Remove the CinderX frame evaluator function from the runtime.
+ */
+void Ci_FiniFrameEvalFunc();
 
 void Ci_InitOpcodes();
 

cinderx/Interpreter/interpreter_base.cpp

@@ -0,0 +1,67 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include "cinderx/Common/extra-py-flags.h"
+#include "cinderx/Interpreter/interpreter.h"
+#include "cinderx/UpstreamBorrow/borrowed.h"
+
+#if PY_VERSION_HEX >= 0x030D0000
+#include "internal/pycore_function.h"
+#endif
+
+#if PY_VERSION_HEX < 0x030D0000 && defined(ENABLE_EVAL_HOOK)
+#include "cinder/hooks.h"
+#endif
+
+extern "C" {
+
+vectorcallfunc getInterpretedVectorcall(
+    [[maybe_unused]] const PyFunctionObject* func) {
+#ifdef ENABLE_INTERPRETER_LOOP
+  const PyCodeObject* code = (const PyCodeObject*)(func->func_code);
+  return (code->co_flags & CI_CO_STATICALLY_COMPILED)
+      ? Ci_StaticFunction_Vectorcall
+      : Ci_PyFunction_Vectorcall;
+#else
+  return Ci_PyFunction_Vectorcall;
+#endif
+}
+
+int Ci_InitFrameEvalFunc() {
+#ifdef ENABLE_INTERPRETER_LOOP
+#ifdef ENABLE_EVAL_HOOK
+  Ci_hook_EvalFrame = Ci_EvalFrame;
+#elif defined(ENABLE_PEP523_HOOK)
+  // Let borrowed.h know the eval frame pointer
+  Ci_EvalFrameFunc = Ci_EvalFrame;
+
+  auto interp = _PyInterpreterState_GET();
+  auto current_eval_frame = _PyInterpreterState_GetEvalFrameFunc(interp);
+  if (current_eval_frame == Ci_EvalFrame) {
+    return 0;
+  }
+  if (current_eval_frame != _PyEval_EvalFrameDefault) {
+    PyErr_SetString(
+        PyExc_RuntimeError,
+        "CinderX tried to set a frame evaluator function but something else "
+        "has done it first, this is not supported");
+    return -1;
+  }
+
+  _PyInterpreterState_SetEvalFrameFunc(interp, Ci_EvalFrame);
+#endif
+#endif
+
+  return 0;
+}
+
+void Ci_FiniFrameEvalFunc() {
+#ifdef ENABLE_INTERPRETER_LOOP
+#ifdef ENABLE_EVAL_HOOK
+  Ci_hook_EvalFrame = nullptr;
+#elif defined(ENABLE_PEP523_HOOK)
+  _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), nullptr);
+#endif
+#endif
+}
+
+} // extern "C"

cinderx/Jit/pyjit.cpp

@@ -1472,14 +1472,20 @@ PyObject* patched_sys_settrace(
 
 #endif // PY_VERSION_HEX >= 0x030C0000
 
-void compile_after_n_calls_impl(uint32_t calls) {
+int compile_after_n_calls_impl(uint32_t calls) {
+  if (Ci_InitFrameEvalFunc() < 0) {
+    return -1;
+  }
+
   getMutableConfig().compile_after_n_calls = calls;
 
   // Schedule all pre-existing functions for compilation.
   walkFunctionObjects(
       [](BorrowedRef<PyFunctionObject> func) { scheduleJitCompile(func); });
 
   JIT_DLOG("Configuring JIT to compile functions after {} calls", calls);
+
+  return 0;
 }
 
 PyObject* compile_after_n_calls(PyObject* /* self */, PyObject* arg) {
@@ -1495,14 +1501,18 @@ PyObject* compile_after_n_calls(PyObject* /* self */, PyObject* arg) {
     return nullptr;
   }
 
-  compile_after_n_calls_impl(calls);
+  if (compile_after_n_calls_impl(calls) < 0) {
+    return nullptr;
+  }
 
   Py_RETURN_NONE;
 }
 
 PyObject* auto_jit(PyObject* /* self */, PyObject* /* arg */) {
   // Default value that works well for most applications.
-  compile_after_n_calls_impl(1000);
+  if (compile_after_n_calls_impl(1000) < 0) {
+    return nullptr;
+  }
 
   Py_RETURN_NONE;
 }
@@ -1567,6 +1577,10 @@ PyObject* force_compile(PyObject* /* self */, PyObject* arg) {
     Py_RETURN_FALSE;
   }
 
+  if (Ci_InitFrameEvalFunc() < 0) {
+    return nullptr;
+  }
+
   _PyJIT_Result result = compileFunction(func);
   switch (result) {
     case PYJIT_RESULT_OK:
@@ -1619,6 +1633,10 @@ PyObject* lazy_compile(PyObject* /* self */, PyObject* arg) {
     Py_RETURN_FALSE;
   }
 
+  if (Ci_InitFrameEvalFunc() < 0) {
+    return nullptr;
+  }
+
   func->vectorcall = forcedJitVectorcall;
   if (!registerFunction(func)) {
     func->vectorcall = getInterpretedVectorcall(func);
@@ -1884,13 +1902,19 @@ void deleteJitList() {
 
 // Reschedule all functions on the JIT list for compilation.  Run when the JIT
 // list is modified.
-void rescheduleJitList() {
+int rescheduleJitList() {
+  if (Ci_InitFrameEvalFunc() < 0) {
+    return -1;
+  }
+
   walkFunctionObjects([](BorrowedRef<PyFunctionObject> func) {
     auto jit_list = cinderx::getModuleState()->jitList();
     if (jit_list->lookupFunc(func)) {
       scheduleJitCompile(func);
     }
   });
+
+  return 0;
 }
 
 PyObject* append_jit_list(PyObject* /* self */, PyObject* arg) {
@@ -1924,7 +1948,9 @@ PyObject* append_jit_list(PyObject* /* self */, PyObject* arg) {
     return nullptr;
   }
 
-  rescheduleJitList();
+  if (rescheduleJitList() < 0) {
+    return nullptr;
+  }
 
   Py_RETURN_NONE;
 }
@@ -1960,7 +1986,9 @@ PyObject* read_jit_list(PyObject* /* self */, PyObject* arg) {
     return nullptr;
   }
 
-  rescheduleJitList();
+  if (rescheduleJitList() < 0) {
+    return nullptr;
+  }
 
   Py_RETURN_NONE;
 }
@@ -3364,9 +3392,13 @@ int initialize() {
   // startup, start scheduling functions for compilation now.
   if (auto compile_n = getConfig().compile_after_n_calls;
       compile_n.has_value()) {
-    compile_after_n_calls_impl(*compile_n);
+    if (compile_after_n_calls_impl(*compile_n) < 0) {
+      return -1;
+    }
   } else if (mod_state->jitList() != nullptr) {
-    rescheduleJitList();
+    if (rescheduleJitList() < 0) {
+      return -1;
+    }
   }
 
   return 0;

cinderx/PythonLib/cinderx/__init__.py

@@ -86,7 +86,10 @@ def is_supported_runtime() -> bool:
         get_parallel_gc_settings,
         has_parallel_gc,
         immortalize_heap,
+        install_frame_evaluator,
+        is_frame_evaluator_installed,
         is_immortal,
+        remove_frame_evaluator,
         strict_module_patch,
         strict_module_patch_delete,
         strict_module_patch_enabled,
@@ -501,11 +504,20 @@ def has_parallel_gc() -> bool:
     def immortalize_heap() -> None:
         pass
 
+    def install_frame_evaluator() -> None:
+        pass
+
+    def is_frame_evaluator_installed() -> bool:
+        return False
+
     def is_immortal(obj: object) -> bool:
         raise RuntimeError(
             "Can't answer whether an object is mortal or immortal from Python code"
         )
 
+    def remove_frame_evaluator() -> None:
+        pass
+
     def strict_module_patch(mod: object, name: str, value: object) -> None:
         pass
 

cinderx/PythonLib/cinderx/compiler/strict/loader.py

@@ -46,7 +46,7 @@
     from types import CodeType, ModuleType
     from typing import Callable, cast, Collection, final, Iterable, Mapping
 
-    from _cinderx import StrictModule, watch_sys_modules
+    from _cinderx import install_frame_evaluator, StrictModule, watch_sys_modules
     from cinderx.static import install_sp_audit_hook
 
     from ..consts import CI_CO_STATICALLY_COMPILED
@@ -795,6 +795,7 @@ def init_static_python() -> None:
 
     Should be called at least once if any Static modules/functions exist.
     """
+    install_frame_evaluator()
     watch_sys_modules()
     install_sp_audit_hook()
 

cinderx/PythonLib/test_cinderx/test_compiler/test_static/compile.py

@@ -24,7 +24,7 @@
 
 import cinderx.jit
 import xxclassloader
-from cinderx import StrictModule
+from cinderx import install_frame_evaluator, StrictModule
 from cinderx.compiler.consts import CI_CO_STATICALLY_COMPILED
 from cinderx.compiler.pycodegen import PythonCodeGenerator
 from cinderx.compiler.static import StaticCodeGenerator
@@ -107,6 +107,8 @@ def __class_getitem__(cls, elem_type):
     code = comp.compile("xxclassloader", "", tree, codestr, optimize=0)
     d = {"<builtins>": builtins.__dict__}
     add_fixed_module(d)
+
+    install_frame_evaluator()
     exec(code, d, d)
 
     xxclassloader.XXGeneric = d["XXGeneric"]