Support running test_jit_disable.py functions multiple times

Summary:
The test functions here need to be explicitly uncompiled (and in one case,
unsuppressed) otherwise their compiled state will stick around on the next test
run.  This happens when running the refleak tests.

`force_uncompile` needs to get stronger about forcefully dropping all
compilation state for a function.  Its current behavior is effectively just a
deopt, if we create another function with the same code object we'll reopt it
immediately upon calling lazy_compile(), which is not what we want here.

Reviewed By: czardoz

Differential Revision: D79690872

fbshipit-source-id: 52f14546cf67cb0f7ac6f3a9bd1b8fe55f35d256

Author: alexmalyshev

cinderx/Jit/context.cpp

@@ -35,6 +35,11 @@ Context::CompilationResult Context::compilePreloader(
   return result;
 }
 
+void Context::uncompile(BorrowedRef<PyFunctionObject> func) {
+  deoptFuncImpl(func);
+  compiled_codes_.erase(CompilationKey{func});
+}
+
 bool Context::deoptFunc(BorrowedRef<PyFunctionObject> func) {
   if (deoptFuncImpl(func)) {
     deopted_funcs_.emplace(func);
@@ -104,6 +109,9 @@ void Context::funcModified(BorrowedRef<PyFunctionObject> func) {
 void Context::funcDestroyed(BorrowedRef<PyFunctionObject> func) {
   compiled_funcs_.erase(func);
   deopted_funcs_.erase(func);
+
+  // This doesn't modify compiled_codes_, so if this is a nested function it can
+  // easily be reopted later.
 }
 
 Context::CompilationResult Context::compilePreloaderImpl(

cinderx/Jit/context.h

@@ -26,6 +26,11 @@ struct CompilationKey {
   PyObject* builtins;
   PyObject* globals;
 
+  explicit CompilationKey(BorrowedRef<PyFunctionObject> func)
+      : code{func->func_code},
+        builtins{func->func_builtins},
+        globals{func->func_globals} {}
+
   CompilationKey(PyObject* code, PyObject* builtins, PyObject* globals)
       : code(code), builtins(builtins), globals(globals) {}
 
@@ -72,6 +77,11 @@ class Context {
       BorrowedRef<PyFunctionObject> func,
       const hir::Preloader& preloader);
 
+  /*
+   * Fully remove all effects of compilation from a function.
+   */
+  void uncompile(BorrowedRef<PyFunctionObject> func);
+
   /*
    * De-optimize a function by setting it to run through the interpreter if it
    * had been previously JIT-compiled.

cinderx/Jit/pyjit.cpp

@@ -1382,6 +1382,9 @@ PyObject* force_uncompile(PyObject* /* self */, PyObject* arg) {
   // "Destroy" the function from the perspective of the JIT, effectively erasing
   // all traces of it from the metadata.
   funcDestroyed(func);
+  if (jit_ctx != nullptr) {
+    jit_ctx->uncompile(func);
+  }
 
   Py_RETURN_TRUE;
 }

cinderx/PythonLib/test_cinderx/test_jit_disable.py

@@ -8,10 +8,12 @@
     disable as disable_jit,
     enable as enable_jit,
     force_compile,
+    force_uncompile,
     INSTALLED,
     is_enabled as is_jit_enabled,
     is_jit_compiled,
     jit_suppress,
+    jit_unsuppress,
     lazy_compile,
     pause as pause_jit,
 )
@@ -45,6 +47,8 @@ def foo(a, b):
         enable_jit()
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_suppress_and_reopt(self) -> None:
         def foo(a, b):
             return a + b
@@ -61,6 +65,8 @@ def foo(a, b):
         enable_jit()
         self.assertFalse(is_jit_compiled(foo))
 
+        jit_unsuppress(foo)
+
     def test_disable_then_deopt(self) -> None:
         def foo(a, b):
             return a + b
@@ -77,6 +83,8 @@ def foo(a, b):
         enable_jit()
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_already_disabled(self) -> None:
         def foo(a, b):
             return a + b
@@ -93,6 +101,8 @@ def foo(a, b):
         enable_jit()
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_already_enabled(self) -> None:
         def foo(a, b):
             return a + b
@@ -109,6 +119,8 @@ def foo(a, b):
         enable_jit()
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_compile_new_after_reenable(self) -> None:
         disable_jit(deopt_all=True)
 
@@ -123,6 +135,8 @@ def foo(a, b):
         force_compile(foo)
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_pause(self) -> None:
         def foo(a, b):
             return a + b
@@ -148,6 +162,8 @@ def foo(a, b):
         self.assertTrue(is_jit_enabled())
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_pause_nested(self) -> None:
         def foo(a, b):
             return a + b
@@ -169,6 +185,8 @@ def foo(a, b):
         self.assertTrue(is_jit_enabled())
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
     def test_pause_between_lazy_compile(self) -> None:
         def foo(a, b):
             return a + b
@@ -185,6 +203,8 @@ def foo(a, b):
         foo(3, 4)
         self.assertTrue(is_jit_compiled(foo))
 
+        force_uncompile(foo)
+
 
 if __name__ == "__main__":
     unittest.main()