Support dynamically enabling the JIT with a new call count threshold

Summary:
Adds new Python API functions for configuring the JIT to compile functions
automatically.

`cinderx.jit.compile_after_n_calls(N)`: As it sounds, this tells the JIT to
compile functions after they have been called `N` times.  This is equivalent to
`PYTHONJITAUTO=N`, but it can be called multiple times as a program runs.

`cinderx.jit.auto()`: This tells the JIT to automatically compile functions, but
leaves all the decision-making to the JIT.  It's the recommended behavior for
users.

Reviewed By: jbower-fb

Differential Revision: D81356873

fbshipit-source-id: 94145aa1b4d00d29301ae8c5d9651c932e2e9584

Author: alexmalyshev

cinderx/Jit/pyjit.cpp

@@ -191,6 +191,18 @@ _PyJIT_Result tryCompile(BorrowedRef<PyFunctionObject> func) {
   return result;
 }
 
+void incrementShadowcodeCall([[maybe_unused]] BorrowedRef<PyCodeObject> code) {
+#if SHADOWCODE_SUPPORTED
+  // The interpreter will only increment up to the shadowcode threshold
+  // PYSHADOW_INIT_THRESHOLD. After that, it will stop incrementing. If someone
+  // sets -X jit-auto above the PYSHADOW_INIT_THRESHOLD, we still have to keep
+  // counting.
+  if (code->co_mutable->ncalls > PYSHADOW_INIT_THRESHOLD) {
+    code->co_mutable->ncalls++;
+  }
+#endif
+}
+
 // Python function entry point when AutoJIT is enabled.
 PyObject* autoJITVectorcall(
     PyObject* func_obj,
@@ -207,16 +219,10 @@ PyObject* autoJITVectorcall(
 
   // Interpret function as usual until it passes the call count threshold.
 
-  // The interpreter will only increment up to the shadowcode threshold
-  // PYSHADOW_INIT_THRESHOLD. After that, it will stop incrementing. If someone
-  // sets -X jit-auto above the PYSHADOW_INIT_THRESHOLD, we still have to keep
-  // counting.
-#if SHADOWCODE_SUPPORTED
-  if (code->co_mutable->ncalls > PYSHADOW_INIT_THRESHOLD) {
-    code->co_mutable->ncalls++;
-  }
-#endif
-  if (countCalls(code) < jit::getConfig().auto_jit_threshold) {
+  auto const calls = countCalls(code);
+
+  if (calls < jit::getConfig().auto_jit_threshold) {
+    incrementShadowcodeCall(code);
     auto entry = getInterpretedVectorcall(func);
     return entry(func_obj, stack, nargsf, kwnames);
   }
@@ -226,6 +232,7 @@ PyObject* autoJITVectorcall(
     return nullptr;
   }
   if (result == PYJIT_RESULT_RETRY) {
+    incrementShadowcodeCall(code);
     auto entry = getInterpretedVectorcall(func);
     return entry(func_obj, stack, nargsf, kwnames);
   }
@@ -1332,6 +1339,44 @@ PyObject* enable_jit(PyObject* /* self */, PyObject* /* arg */) {
   Py_RETURN_NONE;
 }
 
+PyObject* compile_after_n_calls(PyObject* /* self */, PyObject* arg) {
+  Py_ssize_t calls = -1;
+  if (!PyArg_Parse(arg, "n:compile_after_n_calls", &calls)) {
+    return nullptr;
+  }
+  if (calls < 0) {
+    PyErr_Format(
+        PyExc_ValueError,
+        "Cannot configure JIT to compile functions after '%zd' calls",
+        calls);
+    return nullptr;
+  }
+  if (calls == 0) {
+    PyErr_Format(
+        PyExc_ValueError,
+        "compile_after_n_calls(0) not supported yet, use PYTHONJITALL=1");
+    return nullptr;
+  }
+
+  getMutableConfig().auto_jit_threshold = calls;
+  JIT_DLOG("Configuring JIT to compile functions after {} calls", calls);
+
+  Py_RETURN_NONE;
+}
+
+PyObject* auto_jit(PyObject* /* self */, PyObject* /* arg */) {
+  // Default value that works well for most applications.
+  constexpr size_t kThreshold = 1000;
+
+  getMutableConfig().auto_jit_threshold = kThreshold;
+
+  JIT_DLOG(
+      "Configuring JIT to compile functions automatically using default "
+      "behavior");
+
+  Py_RETURN_NONE;
+}
+
 PyObject* get_batch_compilation_time_ms(PyObject*, PyObject*) {
   return PyLong_FromLong(g_batch_compilation_time.count());
 }
@@ -2300,6 +2345,16 @@ PyMethodDef jit_methods[] = {
      METH_NOARGS,
      PyDoc_STR("Re-enable the JIT and re-attach compiled onto previously "
                "JIT-compiled functions.")},
+    {"auto",
+     auto_jit,
+     METH_NOARGS,
+     PyDoc_STR("Configure the JIT to automatically compile functions, using "
+               "default settings")},
+    {"compile_after_n_calls",
+     compile_after_n_calls,
+     METH_O,
+     PyDoc_STR("Configure the JIT to automatically compile functions after "
+               "they are called a set number of times.")},
     {"disassemble", disassemble, METH_O, "Disassemble JIT compiled functions."},
     {"dump_elf",
      dump_elf,

cinderx/PythonLib/cinderx/jit.py

@@ -20,8 +20,10 @@
     from cinderjit import (
         _deopt_gen,
         append_jit_list,
+        auto,
         auto_jit_threshold,
         clear_runtime_stats,
+        compile_after_n_calls,
         count_interpreted_calls,
         disable,
         disable_emit_type_annotation_guards,
@@ -78,12 +80,18 @@ def _deopt_gen(
     def append_jit_list(entry: str) -> None:
         return None
 
+    def auto() -> None:
+        return None
+
     def auto_jit_threshold() -> int:
         return 0
 
     def clear_runtime_stats() -> None:
         return None
 
+    def compile_after_n_calls(calls: int) -> None:
+        return None
+
     def count_interpreted_calls(func: FuncAny) -> int:
         return 0
 

cinderx/PythonLib/test_cinderx/test_cinderjit.py

@@ -30,6 +30,7 @@
 import cinderx.test_support as cinder_support
 from cinderx.compiler.consts import CO_FUTURE_BARRY_AS_BDFL, CO_SUPPRESS_JIT
 from cinderx.jit import (
+    compile_after_n_calls,
     force_compile,
     force_uncompile,
     is_jit_compiled,
@@ -2564,6 +2565,16 @@ def f(x: int) -> int:
 
 @unittest.skipUnless(cinderx.jit.is_enabled(), "Testing the cinderjit module itself")
 class BadArgumentTests(unittest.TestCase):
+    def test_compile_after_n_calls(self) -> None:
+        with self.assertRaises(TypeError):
+            compile_after_n_calls(None)
+        with self.assertRaises(TypeError):
+            compile_after_n_calls(is_jit_compiled)
+        with self.assertRaises(ValueError):
+            compile_after_n_calls(-1)
+        with self.assertRaises(ValueError):
+            compile_after_n_calls(0)
+
     def test_is_compiled(self) -> None:
         with self.assertRaises(TypeError):
             is_jit_compiled(None)

cinderx/PythonLib/test_cinderx/test_jit_disable.py

@@ -2,7 +2,12 @@
 
 # pyre-unsafe
 
+import subprocess
+import sys
+import tempfile
+import textwrap
 import unittest
+from pathlib import Path
 
 from cinderx.jit import (
     disable as disable_jit,
@@ -17,6 +22,7 @@
     lazy_compile,
     pause as pause_jit,
 )
+from cinderx.test_support import CINDERX_PATH
 
 
 @unittest.skipUnless(is_jit_enabled(), "Tests functionality on the JIT")
@@ -201,6 +207,100 @@ def foo(a, b):
         # detection.
         force_uncompile(foo)
 
+    def test_default(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            code = textwrap.dedent("""
+            import cinderx.jit
+
+            def inc(x):
+                return x + 1
+
+            assert not cinderx.jit.is_jit_compiled(inc)
+            cinderx.jit.force_compile(inc)
+            assert cinderx.jit.is_jit_compiled(inc)
+            """)
+
+            test_file = Path(tmp_dir) / "mod.py"
+            test_file.write_text(code)
+
+            subprocess.run(
+                [sys.executable, str(test_file)],
+                check=True,
+                env={"PYTHONPATH": CINDERX_PATH},
+            )
+
+    def test_auto(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            code = textwrap.dedent("""
+            import cinderx.jit
+
+            cinderx.jit.auto()
+
+            def inc(x):
+                return x + 1
+
+            assert not cinderx.jit.is_jit_compiled(inc)
+            for i in range(1000):
+                inc(i)
+            assert not cinderx.jit.is_jit_compiled(inc)
+
+            inc(1001)
+            assert cinderx.jit.is_jit_compiled(inc)
+            """)
+
+            test_file = Path(tmp_dir) / "mod.py"
+            test_file.write_text(code)
+
+            subprocess.run(
+                [sys.executable, str(test_file)],
+                check=True,
+                env={"PYTHONPATH": CINDERX_PATH},
+            )
+
+    def test_compile_after_n_calls(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            code = textwrap.dedent("""
+            import cinderx.jit
+
+            cinderx.jit.compile_after_n_calls(2)
+
+            def inc(x):
+                return x + 1
+
+            assert not cinderx.jit.is_jit_compiled(inc)
+            inc(1)
+            inc(2)
+            assert not cinderx.jit.is_jit_compiled(inc)
+
+            inc(3)
+            assert cinderx.jit.is_jit_compiled(inc)
+
+            cinderx.jit.compile_after_n_calls(5)
+
+            def dec(x):
+                return x - 1
+
+            assert not cinderx.jit.is_jit_compiled(dec)
+            dec(1)
+            dec(2)
+            dec(3)
+            dec(4)
+            dec(5)
+            assert not cinderx.jit.is_jit_compiled(dec)
+
+            dec(6)
+            assert cinderx.jit.is_jit_compiled(dec)
+            """)
+
+            test_file = Path(tmp_dir) / "mod.py"
+            test_file.write_text(code)
+
+            subprocess.run(
+                [sys.executable, str(test_file)],
+                check=True,
+                env={"PYTHONPATH": CINDERX_PATH},
+            )
+
 
 if __name__ == "__main__":
     unittest.main()