Re-enable load method elimination for 3.12

Summary:
Re-enables the load method elimination optimization in multi-threaded compile that's currently disabled in 3.12.

We maintain our own dictionary of core built-in types which we can apply this to. To make this thread safe in 3.12 we use an `unordered_map` for the type->dict mapping as we can't safely use a map that's keyed off non-unicode keys without potentially hitting the interpreter state. We then use a normal Python dict w/ only unicode keys for the members which is thread-safe w/o the GIL held.

If a type isn't one of the built-in types but is immutable, has no meta-class and has no dictionary we also allow the optimization to be applied (although this hits very few types in the test cases - mostly just something in the lru cache).

The `unordered_map` is our first usage of a C++ type in `ModuleState` which has initialization that's something other than zero-initialization. Therefore we now run the in place constructor on the memory that the Python runtime has allocated for us to store the ModuleState.

Reviewed By: mpage

Differential Revision: D78361677

fbshipit-source-id: ec3280d28dc434bb2b2e4f1a81156030e4aab3ac

Author: DinoV

cinderx/Jit/hir/optimization.cpp

@@ -18,6 +18,7 @@
 #include "cinderx/Jit/hir/printer.h"
 #include "cinderx/Jit/threaded_compile.h"
 #include "cinderx/UpstreamBorrow/borrowed.h" // @donotremove
+#include "cinderx/module_state.h"
 
 #include <fmt/format.h>
 
@@ -824,20 +825,47 @@ struct MethodInvoke {
   CallMethod* call_method{nullptr};
 };
 
-// Returns true if LoadMethod/CallMethod/GetSecondOutput were removed.
-// Returns false if they could not be removed.
-static bool tryEliminateLoadMethod(Function& irfunc, MethodInvoke& invoke) {
-  // This isn't safe in the multi-threaded compilation on 3.12 because we
-  // don't hold the GIL which is required for PyType_Lookup.
-  RETURN_MULTITHREADED_COMPILE(false);
+#if PY_VERSION_HEX >= 0x030C0000
+BorrowedRef<> immutableMultithreadedTypeLookup(
+    BorrowedRef<PyTypeObject> type,
+    BorrowedRef<> name) {
+  BorrowedRef<> mro = type->tp_mro;
+  for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(mro.get()); i++) {
+    PyTypeObject* mro_type =
+        reinterpret_cast<PyTypeObject*>(PyTuple_GET_ITEM(mro.get(), i));
+    if (PyType_HasFeature(mro_type, _Py_TPFLAGS_STATIC_BUILTIN)) {
+      auto& builtins = cinderx::getModuleState()->builtinMembers();
+
+      auto members = builtins.find(mro_type);
+      if (members == builtins.end()) {
+        // We don't know anything about this builtin type.
+        return nullptr;
+      }
+      // We load all of the members from the MRO in the builtins
+      // cache so it's completely authorative.
+      return PyDict_GetItemWithError(members->second, name);
+    } else if (
+        !PyType_HasFeature(mro_type, Py_TPFLAGS_IMMUTABLETYPE) ||
+        !PyType_CheckExact(mro_type)) {
+      // We can't trust anything about this base type
+      return nullptr;
+    }
 
-  ThreadedCompileSerialize guard;
-  PyCodeObject* code = invoke.load_method->frameState()->code;
-  PyObject* names = code->co_names;
-  PyObject* name = PyTuple_GetItem(names, invoke.load_method->name_idx());
-  JIT_DCHECK(name != nullptr, "name must not be null");
-  Register* receiver = invoke.load_method->receiver();
-  Type receiver_type = receiver->type();
+    BorrowedRef<> method_obj =
+        PyDict_GetItemWithError(_PyType_GetDict(mro_type), name);
+    if (method_obj != nullptr) {
+      return method_obj;
+    }
+  }
+  return nullptr;
+}
+#endif
+
+// Gets a directly invokable method object from a JIT Type. This only succeeds
+// if we know the type can be directly invoked.
+static BorrowedRef<> getMethodObjectFromType(
+    Type receiver_type,
+    BorrowedRef<> name) {
   // This is a list of common builtin types whose methods cannot be overwritten
   // from managed code and for which looking up the methods is guaranteed to
   // not do anything "weird" that needs to happen at runtime, like make a
@@ -846,26 +874,85 @@ static bool tryEliminateLoadMethod(Function& irfunc, MethodInvoke& invoke) {
   // loading and invoking methods off an instance (e.g. {}.fromkeys(...)) is
   // resolved and called differently than from the type (e.g.
   // dict.fromkeys(...)). The code below handles the instance case only.
+#if PY_VERSION_HEX < 0x030C0000
+
   if (!(receiver_type <= TArray || receiver_type <= TBool ||
         receiver_type <= TBytesExact || receiver_type <= TCode ||
         receiver_type <= TDictExact || receiver_type <= TFloatExact ||
         receiver_type <= TListExact || receiver_type <= TLongExact ||
         receiver_type <= TNoneType || receiver_type <= TSetExact ||
         receiver_type <= TTupleExact || receiver_type <= TUnicodeExact)) {
-    return false;
+    return nullptr;
   }
   PyTypeObject* type = receiver_type.runtimePyType();
   if (type == nullptr) {
     // This might happen for a variety of reasons, such as encountering a
     // method load on a maybe-defined value where the definition occurs in a
-    // block of code that isn't seen by the compiler (e.g. in an except block).
+    // block of code that isn't seen by the compiler (e.g. in an except
+    // block).
     JIT_DCHECK(
         receiver_type == TBottom,
         "Type {} expected to have PyTypeObject*",
         receiver_type);
-    return false;
+    return nullptr;
+  }
+  return _PyType_Lookup(type, name);
+#else
+  if (!receiver_type.hasTypeExactSpec()) {
+    return nullptr;
+  }
+  PyTypeObject* type = receiver_type.runtimePyType();
+  if (type == nullptr) {
+    // This might happen for a variety of reasons, such as encountering a
+    // method load on a maybe-defined value where the definition occurs in a
+    // block of code that isn't seen by the compiler (e.g. in an except
+    // block).
+    JIT_DCHECK(
+        receiver_type == TBottom,
+        "Type {} expected to have PyTypeObject*",
+        receiver_type);
+    return nullptr;
   }
-  auto method_obj = Ref<>::create(_PyType_Lookup(type, name));
+
+  BorrowedRef<> method_obj = nullptr;
+  // In 3.12 we can't do PyType_Lookup because for built-in types it needs
+  // access to the current runtime, and in multi-threaded compile we don't
+  // have it. So we instead have a cache of all of the builtin types that we
+  // support this for.
+  auto& builtins = cinderx::getModuleState()->builtinMembers();
+
+  if (PyType_HasFeature(type, _Py_TPFLAGS_STATIC_BUILTIN)) {
+    auto it = builtins.find(receiver_type.runtimePyType());
+    if (it != builtins.end()) {
+      method_obj = PyDict_GetItemWithError(it->second, name);
+    }
+  } else if (
+      PyType_HasFeature(type, Py_TPFLAGS_IMMUTABLETYPE) &&
+      PyType_CheckExact(type) && type->tp_dictoffset == 0) {
+    method_obj = immutableMultithreadedTypeLookup(type, name);
+    if (Py_TYPE(method_obj) != &PyClassMethodDescr_Type &&
+        Py_TYPE(method_obj) != &PyMethodDescr_Type &&
+        Py_TYPE(method_obj) != &PyWrapperDescr_Type &&
+        Py_TYPE(method_obj) != &PyFunction_Type) {
+      method_obj = nullptr;
+    }
+  }
+  return method_obj;
+#endif
+}
+
+// Returns true if LoadMethod/CallMethod/GetSecondOutput were removed.
+// Returns false if they could not be removed.
+static bool tryEliminateLoadMethod(Function& irfunc, MethodInvoke& invoke) {
+  ThreadedCompileSerialize guard;
+  PyCodeObject* code = invoke.load_method->frameState()->code;
+  PyObject* names = code->co_names;
+  PyObject* name = PyTuple_GetItem(names, invoke.load_method->name_idx());
+  JIT_DCHECK(name != nullptr, "name must not be null");
+
+  Register* receiver = invoke.load_method->receiver();
+  Type receiver_type = receiver->type();
+  BorrowedRef<> method_obj = getMethodObjectFromType(receiver_type, name);
   if (method_obj == nullptr) {
     // No such method. Let the LoadMethod fail at runtime. _PyType_Lookup does
     // not raise an exception.
@@ -889,14 +976,17 @@ static bool tryEliminateLoadMethod(Function& irfunc, MethodInvoke& invoke) {
     // Pass the type as the first argument (e.g. dict.fromkeys).
     Register* type_reg = irfunc.env.AllocateRegister();
     auto load_type = LoadConst::create(
-        type_reg, Type::fromObject(reinterpret_cast<PyObject*>(type)));
+        type_reg,
+        Type::fromObject(
+            reinterpret_cast<PyObject*>(receiver_type.runtimePyType())));
     load_type->setBytecodeOffset(invoke.load_method->bytecodeOffset());
     load_type->InsertBefore(*invoke.call_method);
     call_static->SetOperand(1, type_reg);
   } else {
     JIT_DCHECK(
         Py_TYPE(method_obj) == &PyMethodDescr_Type ||
-            Py_TYPE(method_obj) == &PyWrapperDescr_Type,
+            Py_TYPE(method_obj) == &PyWrapperDescr_Type ||
+            Py_TYPE(method_obj) == &PyFunction_Type,
         "unexpected type");
     // Pass the instance as the first argument (e.g. str.join, str.__mod__).
     call_static->SetOperand(1, receiver);

cinderx/RuntimeTests/hir_tests/all_passes_test.txt

@@ -496,15 +496,26 @@ fun jittestmodule:test {
         Locals<1> v11
       }
     }
+    v23:Object = LoadMethodCached<0; "append"> v15 {
+      LiveValues<1> o:v15
+      FrameState {
+        CurInstrOffset 8
+        Locals<1> v15
+      }
+    }
+    v18:OptObject = GetSecondOutput<OptObject> v23
     v19:ImmortalLongExact[1] = LoadConst<ImmortalLongExact[1]>
-    v24:CInt32 = ListAppend v15 v19 {
-      LiveValues<2> o:v15 unc:v19
+    v20:Object = CallMethod<3> v23 v18 v19 {
+      LiveValues<4> o:v15 o:v18 unc:v19 o:v23
       FrameState {
         CurInstrOffset 30
         Locals<1> v15
       }
     }
     Decref v15
+    XDecref v18
+    Decref v20
+    Decref v23
     v21:ImmortalNoneType = LoadConst<ImmortalNoneType>
     Return<ImmortalNoneType> v21
   }

cinderx/RuntimeTests/hir_tests/builtin_load_method_elimination_test.txt

@@ -226,11 +226,15 @@ fun jittestmodule:test {
   bb 3 (preds 1, 2) {
     v28:ImmortalUnicodeExact = Phi<1, 2> v22 v24
     Snapshot
-    UseType<ImmortalUnicodeExact> v28
-    v32:MortalObjectUser[method_descriptor:0xdeadbeef] = LoadConst<MortalObjectUser[method_descriptor:0xdeadbeef]>
-    v30:ImmortalUnicodeExact = Assign v28
+    v32:Object = LoadMethodCached<0; "join"> v28 {
+      FrameState {
+        CurInstrOffset 18
+        Locals<3> v12 v13 v28
+      }
+    }
+    v30:OptObject = GetSecondOutput<OptObject> v32
     Snapshot
-    v31:UnicodeExact = VectorCall<2, static> v32 v28 v12 {
+    v31:Object = CallMethod<3> v32 v30 v12 {
       FrameState {
         CurInstrOffset 40
         Locals<3> v12 v13 v28
@@ -402,11 +406,14 @@ fun jittestmodule:test {
       }
     }
     Snapshot
-    UseType<UnicodeExact> v16
-    v21:MortalObjectUser[method_descriptor:0xdeadbeef] = LoadConst<MortalObjectUser[method_descriptor:0xdeadbeef]>
-    v18:UnicodeExact = Assign v16
+    v21:Object = LoadMethodCached<1; "upper"> v16 {
+      FrameState {
+        CurInstrOffset 32
+      }
+    }
+    v18:OptObject = GetSecondOutput<OptObject> v21
     Snapshot
-    v19:UnicodeExact = VectorCall<1, static> v21 v16 {
+    v19:Object = CallMethod<2> v21 v18 {
       FrameState {
         CurInstrOffset 52
       }

cinderx/_cinderx-lib.cpp

@@ -1221,7 +1221,8 @@ PyMethodDef _cinderx_methods[] = {
     {nullptr, nullptr, 0, nullptr}};
 
 static int _cinderx_exec(PyObject* m) {
-  auto state = (cinderx::ModuleState*)PyModule_GetState(m);
+  void* state_mem = PyModule_GetState(m);
+  auto state = new (state_mem) cinderx::ModuleState();
   auto cache_manager = new (std::nothrow) jit::GlobalCacheManager();
   if (cache_manager == nullptr) {
     return -1;
@@ -1343,6 +1344,12 @@ static int _cinderx_exec(PyObject* m) {
   }
   state->setSymbolizer(symbolizer);
 
+  if constexpr (PY_VERSION_HEX >= 0x030C0000) {
+    if (!state->initBuiltinMembers()) {
+      return -1;
+    }
+  }
+
   cinderx::setModule(m);
 
   CiExc_StaticTypeError =

cinderx/module_state.cpp

@@ -2,6 +2,8 @@
 
 #include "cinderx/module_state.h"
 
+#include "internal/pycore_object.h"
+
 #include "cinderx/Common/log.h"
 
 namespace cinderx {
@@ -43,6 +45,58 @@ ModuleState* getModuleState() {
   return s_cinderx_state;
 }
 
+bool ModuleState::initBuiltinMembers() {
+#if PY_VERSION_HEX >= 0x030C0000
+  constexpr PyTypeObject* types[] = {
+      &PyBool_Type,
+      &PyBytes_Type,
+      &PyByteArray_Type,
+      &PyComplex_Type,
+      &PyCode_Type,
+      &PyDict_Type,
+      &PyFloat_Type,
+      &PyFrozenSet_Type,
+      &PyList_Type,
+      &PyLong_Type,
+      &_PyNone_Type,
+      &PyProperty_Type,
+      &PySet_Type,
+      &PyTuple_Type,
+      &PyUnicode_Type,
+  };
+
+  for (auto type : types) {
+    PyObject* mro = type->tp_mro;
+    if (mro == nullptr) {
+      continue;
+    }
+
+    Ref<> type_members = Ref<>::steal(PyDict_New());
+    if (type_members == nullptr) {
+      return false;
+    }
+    for (Py_ssize_t i = 0; i < Py_SIZE(mro); i++) {
+      PyTypeObject* base =
+          reinterpret_cast<PyTypeObject*>(PyTuple_GetItem(mro, i));
+      Py_ssize_t cur_mem = 0;
+      PyObject *key, *value;
+      Ref<> tp_dict = Ref<>::steal(PyType_GetDict(base));
+      while (PyDict_Next(tp_dict, &cur_mem, &key, &value)) {
+        if (PyDict_Contains(type_members, key)) {
+          continue;
+        }
+        if (PyDict_SetItem(type_members, key, value) < 0) {
+          return false;
+        }
+      }
+    }
+
+    builtin_members_.emplace(type, std::move(type_members));
+  }
+#endif
+  return true;
+}
+
 } // namespace cinderx
 
 extern "C" {

cinderx/module_state.h

@@ -16,6 +16,7 @@
 #include "cinderx/async_lazy_value_iface.h"
 
 #include <memory>
+#include <unordered_map>
 
 namespace cinderx {
 
@@ -166,6 +167,14 @@ class ModuleState {
         std::unique_ptr<jit::IJitGenFreeList>(jit_gen_free_list);
   }
 
+  // Returns a dictionary of type->dict[name, members] for standard builtin
+  // types.
+  std::unordered_map<PyTypeObject*, Ref<>>& builtinMembers() {
+    return builtin_members_;
+  }
+
+  bool initBuiltinMembers();
+
  private:
   std::unique_ptr<jit::IGlobalCacheManager> cache_manager_;
   std::unique_ptr<jit::ICodeAllocator> code_allocator_;
@@ -175,6 +184,7 @@ class ModuleState {
   std::unique_ptr<IAsyncLazyValueState> async_lazy_value_;
   std::unique_ptr<jit::IJitGenFreeList> jit_gen_free_list_;
   Ref<PyTypeObject> coro_type_, gen_type_, anext_awaitable_type_;
+  std::unordered_map<PyTypeObject*, Ref<>> builtin_members_;
 #ifdef ENABLE_LIGHTWEIGHT_FRAMES
   Ref<> frame_reifier_;
 #endif