Reset function entrypoints when we change compile_after_n_calls

Summary:
Making sure that this affects existing functions and not just functions that are
created afterwards.

Consolidating code between `auto` and `compile_after_n_calls` here, added a
check that the call limit isn't out-of-bounds.

Reviewed By: czardoz

Differential Revision: D82479038

fbshipit-source-id: 909d9f4d044fa985bfab2849a967cc1a3ca3bb88

Author: alexmalyshev

cinderx/Jit/pyjit.cpp

@@ -1339,12 +1339,22 @@ PyObject* enable_jit(PyObject* /* self */, PyObject* /* arg */) {
   Py_RETURN_NONE;
 }
 
+void compile_after_n_calls_impl(uint32_t calls) {
+  getMutableConfig().auto_jit_threshold = calls;
+
+  // Schedule all pre-existing functions for compilation.
+  walkFunctionObjects(
+      [](BorrowedRef<PyFunctionObject> func) { scheduleJitCompile(func); });
+
+  JIT_DLOG("Configuring JIT to compile functions after {} calls", calls);
+}
+
 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) {
+  if (calls < 0 || calls > std::numeric_limits<uint32_t>::max()) {
     PyErr_Format(
         PyExc_ValueError,
         "Cannot configure JIT to compile functions after '%zd' calls",
@@ -1358,21 +1368,14 @@ PyObject* compile_after_n_calls(PyObject* /* self */, PyObject* arg) {
     return nullptr;
   }
 
-  getMutableConfig().auto_jit_threshold = calls;
-  JIT_DLOG("Configuring JIT to compile functions after {} calls", calls);
+  compile_after_n_calls_impl(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");
+  compile_after_n_calls_impl(1000);
 
   Py_RETURN_NONE;
 }

cinderx/PythonLib/test_cinderx/test_cinderjit.py

@@ -2572,6 +2572,8 @@ def test_compile_after_n_calls(self) -> None:
             compile_after_n_calls(is_jit_compiled)
         with self.assertRaises(ValueError):
             compile_after_n_calls(-1)
+        with self.assertRaises(ValueError):
+            compile_after_n_calls(10_000_000_000)
         with self.assertRaises(ValueError):
             compile_after_n_calls(0)
 

cinderx/PythonLib/test_cinderx/test_jit_disable.py

@@ -207,7 +207,12 @@ def foo(a, b):
         # detection.
         force_uncompile(foo)
 
-    def test_default(self) -> None:
+    def test_compile_no_config(self) -> None:
+        """
+        Test how code behaves when it forces compilation without any other
+        configuration or options enabled.
+        """
+
         with tempfile.TemporaryDirectory() as tmp_dir:
             code = textwrap.dedent("""
             import cinderx.jit
@@ -230,10 +235,17 @@ def inc(x):
             )
 
     def test_auto(self) -> None:
+        """
+        Basic test for cinderx.jit.auto().
+        """
+
         with tempfile.TemporaryDirectory() as tmp_dir:
             code = textwrap.dedent("""
             import cinderx.jit
 
+            def predefined(x):
+                return x + x
+
             cinderx.jit.auto()
 
             def inc(x):
@@ -257,7 +269,44 @@ def inc(x):
                 env={"PYTHONPATH": CINDERX_PATH},
             )
 
+    def test_auto_predefined(self) -> None:
+        """
+        Test that cinderx.jit.auto() works for functions that were defined
+        before it was called.
+        """
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            code = textwrap.dedent("""
+            import cinderx.jit
+
+            def predefined(x):
+                return x + x
+
+            cinderx.jit.auto()
+
+            assert not cinderx.jit.is_jit_compiled(predefined)
+            for i in range(1000):
+                predefined(i)
+            assert not cinderx.jit.is_jit_compiled(predefined)
+
+            predefined(1001)
+            assert cinderx.jit.is_jit_compiled(predefined)
+            """)
+
+            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:
+        """
+        Basic test for cinderx.jit.compile_after_n_calls().
+        """
+
         with tempfile.TemporaryDirectory() as tmp_dir:
             code = textwrap.dedent("""
             import cinderx.jit
@@ -275,6 +324,8 @@ def inc(x):
             inc(3)
             assert cinderx.jit.is_jit_compiled(inc)
 
+            # Change the setting and see it takes affect.
+
             cinderx.jit.compile_after_n_calls(5)
 
             def dec(x):
@@ -301,6 +352,39 @@ def dec(x):
                 env={"PYTHONPATH": CINDERX_PATH},
             )
 
+    def test_compile_after_n_calls_predefined(self) -> None:
+        """
+        Test that cinderx.jit.compile_after_n_calls() works for functions that
+        were defined before it was called.
+        """
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            code = textwrap.dedent("""
+            import cinderx.jit
+
+            def predefined(x):
+                return x + x
+
+            cinderx.jit.compile_after_n_calls(2)
+
+            assert not cinderx.jit.is_jit_compiled(predefined)
+            predefined(1)
+            predefined(2)
+            assert not cinderx.jit.is_jit_compiled(predefined)
+
+            predefined(3)
+            assert cinderx.jit.is_jit_compiled(predefined)
+            """)
+
+            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()