Move outer coder funcs to jitContext

Summary: This is just a simple refactoring - these are better suited in the Jit's Context and we need to have access to them there in the future.

Reviewed By: alexmalyshev

Differential Revision: D92781350

fbshipit-source-id: c50314160d16ed752581d42029f45ac5aacf9c19

Author: DinoV

cinderx/Jit/context.h

@@ -398,6 +398,13 @@ class Context : public IJitContext {
 
   const hir::Type& typeForCommonConstant(int i) const;
 
+  // Map of all code objects to the functions that they were found in.
+  // Needed for printing the name of the code object and for preloading.
+  UnorderedMap<BorrowedRef<PyCodeObject>, BorrowedRef<PyFunctionObject>>&
+  codeOuterFunctions() {
+    return code_outer_funcs_;
+  }
+
   // Allocate all CodeRuntimes together so they can be mlocked() without
   // including any other data that happened to be on the same page.
   SlabArena<CodeRuntime> code_runtimes_;
@@ -475,6 +482,10 @@ class Context : public IJitContext {
   Ref<> cinderjit_module_;
 
   std::atomic_size_t total_compile_time_ms_;
+
+  // Map of all code objects to the functions that they were found in.
+  UnorderedMap<BorrowedRef<PyCodeObject>, BorrowedRef<PyFunctionObject>>
+      code_outer_funcs_;
 };
 
 // A CompilerContext is like a Context but it also holds a compiler object

cinderx/Jit/pyjit.cpp

@@ -870,7 +870,7 @@ std::string unitFullname(BorrowedRef<> unit) {
   if (func != nullptr) {
     return funcFullname(func);
   }
-  auto& jit_code_outer_funcs = cinderx::getModuleState()->codeOuterFunctions();
+  auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
   auto iter = jit_code_outer_funcs.find(code);
   if (iter == jit_code_outer_funcs.end()) {
     return fmt::format(
@@ -897,8 +897,7 @@ hir::Preloader* preload(BorrowedRef<> unit) {
     preloader =
         hir::Preloader::makePreloader(func, makeFrameReifier(func->func_code));
   } else {
-    auto& jit_code_outer_funcs =
-        cinderx::getModuleState()->codeOuterFunctions();
+    auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
     auto it = jit_code_outer_funcs.find(code);
     if (it == jit_code_outer_funcs.end()) {
       PyErr_Format(
@@ -1113,9 +1112,6 @@ bool compile_all(size_t workers = 0) {
 
   hir::preloaderManager().clear();
 
-  auto& jit_code_outer_funcs = cinderx::getModuleState()->codeOuterFunctions();
-  jit_code_outer_funcs.clear();
-
   return true;
 }
 
@@ -1183,6 +1179,11 @@ bool registerFunction(BorrowedRef<PyFunctionObject> func) {
   auto& jit_reg_units = cinderx::getModuleState()->registeredCompilationUnits();
   jit_reg_units.emplace(func.getObj());
 
+  // Map this function's code object to itself.
+  auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
+  BorrowedRef<PyCodeObject> func_code{func->func_code};
+  jit_code_outer_funcs.emplace(func_code, func);
+
   // If we have an active jit-list, scan this function's code object for any
   // nested functions that might be on the jit-list, and register them as well.
   if (cinderx::getModuleState()->jitList() != nullptr) {
@@ -1191,8 +1192,6 @@ bool registerFunction(BorrowedRef<PyFunctionObject> func) {
     BorrowedRef<> top_consts{top_code->co_consts};
     for (BorrowedRef<PyCodeObject> code : findNestedCodes(module, top_consts)) {
       jit_reg_units.emplace(code.getObj());
-      auto& jit_code_outer_funcs =
-          cinderx::getModuleState()->codeOuterFunctions();
       jit_code_outer_funcs.emplace(code, func);
     }
   }
@@ -2989,6 +2988,11 @@ _PyJIT_Result compile_func(BorrowedRef<PyFunctionObject> func) {
   // jitable function, resulting in a single-function compile
   hir::IsolatedPreloaders ip;
 
+  // Ensure the function's code object is mapped to itself.
+  auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
+  BorrowedRef<PyCodeObject> func_code{func->func_code};
+  jit_code_outer_funcs.emplace(func_code, func);
+
   // Collect a list of functions to compile.  If it's empty then there must have
   // been a Python error during preloading.
   std::vector<BorrowedRef<PyFunctionObject>> targets = preloadFuncAndDeps(func);
@@ -3433,7 +3437,7 @@ void finalize() {
 
   // Clear some global maps that reference Python data.
   auto mod_state = cinderx::getModuleState();
-  auto& jit_code_outer_funcs = mod_state->codeOuterFunctions();
+  auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
   auto& jit_reg_units = mod_state->registeredCompilationUnits();
   jit_code_outer_funcs.clear();
   jit_reg_units.clear();
@@ -3622,9 +3626,9 @@ void codeDestroyed(BorrowedRef<PyCodeObject> code) {
   if (isJitUsable()) {
     auto mod_state = cinderx::getModuleState();
     auto& jit_reg_units = mod_state->registeredCompilationUnits();
-    auto& jit_code_outer_funcs = mod_state->codeOuterFunctions();
+    auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
     jit_reg_units.erase(code.getObj());
-    jit_code_outer_funcs.erase(code.getObj());
+    jit_code_outer_funcs.erase(code);
     if (handle_unit_deleted_during_preload != nullptr) {
       handle_unit_deleted_during_preload(code.getObj());
     }
@@ -3643,7 +3647,7 @@ void funcDestroyed(BorrowedRef<PyFunctionObject> func) {
 
     // erase any child code objects we registered too
     if (mod_state->jitList() != nullptr) {
-      auto& jit_code_outer_funcs = mod_state->codeOuterFunctions();
+      auto& jit_code_outer_funcs = jitCtx()->codeOuterFunctions();
       PyObject* module = func->func_module;
       BorrowedRef<PyCodeObject> top_code{func->func_code};
       BorrowedRef<> top_consts{top_code->co_consts};

cinderx/module_state.cpp

@@ -101,9 +101,4 @@ jit::UnorderedSet<BorrowedRef<>>& ModuleState::registeredCompilationUnits() {
   return registered_compilation_units;
 }
 
-jit::UnorderedMap<BorrowedRef<PyCodeObject>, BorrowedRef<PyFunctionObject>>&
-ModuleState::codeOuterFunctions() {
-  return code_outer_funcs;
-}
-
 } // namespace cinderx

cinderx/module_state.h

@@ -189,9 +189,6 @@ class ModuleState {
 
   jit::UnorderedSet<BorrowedRef<>>& registeredCompilationUnits();
 
-  jit::UnorderedMap<BorrowedRef<PyCodeObject>, BorrowedRef<PyFunctionObject>>&
-  codeOuterFunctions();
-
  private:
   WatcherState watcher_state_;
 
@@ -215,13 +212,6 @@ class ModuleState {
   // Function and code objects ("units") registered for compilation.
   jit::UnorderedSet<BorrowedRef<>> registered_compilation_units;
 
-  // Map of all compiled or to-be-compiled code objects to the functions that
-  // they were found in.
-  //
-  // Needed for printing the name of the code object and for preloading.
-  jit::UnorderedMap<BorrowedRef<PyCodeObject>, BorrowedRef<PyFunctionObject>>
-      code_outer_funcs;
-
   // Function objects registered for pre-fork perf-trampoline compilation.
   jit::UnorderedSet<BorrowedRef<PyFunctionObject>> perf_trampoline_worklist_;