Add JIT support for LOAD_BUILD_CLASS

Summary:
Apparently we've never supported this?! Quite surprising as it must mean we were at least missing quite a few functions from JIT during test.

The code assumes the builtins are always a dictionary, whereas the interpreters bytecode can handle a non-dictionary `PyObject*`. I think this assumption is safe as the preloader seems to always have the builtins as a dictionary.

Reviewed By: alexmalyshev

Differential Revision: D82689837

fbshipit-source-id: 6d10820eea803a365d14ec5268dbfe56aa6b502d

Author: jbower-fb

cinderx/Jit/hir/builder.cpp

@@ -152,6 +152,7 @@ bool isSupportedOpcode(int opcode) {
     case LOAD_ASSERTION_ERROR:
     case LOAD_ATTR:
     case LOAD_ATTR_SUPER:
+    case LOAD_BUILD_CLASS:
     case LOAD_CLOSURE:
     case LOAD_COMMON_CONSTANT:
     case LOAD_CONST:
@@ -1553,6 +1554,10 @@ void HIRBuilder::translate(
           emitSetFunctionAttribute(tc, bc_instr);
           break;
         }
+        case LOAD_BUILD_CLASS: {
+          emitLoadBuildClass(tc);
+          break;
+        }
         case CHECK_EG_MATCH:
         case CHECK_EXC_MATCH:
         case CLEANUP_THROW:
@@ -4986,6 +4991,20 @@ void HIRBuilder::emitSetFunctionAttribute(
   stack.push(func);
 }
 
+void HIRBuilder::emitLoadBuildClass(TranslationContext& tc) {
+  Register* result = temps_.AllocateStack();
+  Register* builtins = temps_.AllocateNonStack();
+  Register* key = temps_.AllocateNonStack();
+  tc.emit<LoadConst>(builtins, Type::fromObject(tc.frame.builtins));
+  // Starting at the preloader the JIT seems to assume builtins will be a
+  // dictionary, however I'm not sure there's any guarantee of this.
+  Register* builtins_dict = temps_.AllocateNonStack();
+  tc.emit<GuardType>(builtins_dict, TDictExact, builtins, tc.frame);
+  tc.emit<LoadConst>(key, Type::fromObject(Runtime::get()->strBuildClass()));
+  tc.emit<DictSubscr>(result, builtins_dict, key, tc.frame);
+  tc.frame.stack.push(result);
+}
+
 void HIRBuilder::insertEvalBreakerCheck(
     CFG& cfg,
     BasicBlock* check_block,

cinderx/Jit/hir/builder.h

@@ -477,6 +477,8 @@ class HIRBuilder {
       TranslationContext& tc,
       const BytecodeInstruction& bc_instr);
 
+  void emitLoadBuildClass(TranslationContext& tc);
+
   BorrowedRef<> constArg(const jit::BytecodeInstruction& bc_instr);
 
   ExecutionBlock popBlock(CFG& cfg, TranslationContext& tc);

cinderx/Jit/pyjit.cpp

@@ -931,6 +931,14 @@ hir::Preloader* preload(BorrowedRef<> unit) {
       return nullptr;
     }
     BorrowedRef<PyFunctionObject>& outer_func = it->second;
+    // Assuming the builtins will always be a dictionary goes way back in the
+    // JIT's history. I'm not sure what guarantees this though. Tread carefully
+    // but try not to blow things up if this happens in production code.
+    JIT_DCHECK(
+        PyDict_CheckExact(outer_func->func_builtins),
+        "Unexpected type for builtins ({}) on function {}",
+        Py_TYPE(outer_func->func_builtins)->tp_name,
+        funcFullname(outer_func));
     preloader = hir::Preloader::makePreloader(
         code,
         outer_func->func_builtins,

cinderx/Jit/runtime.cpp

@@ -13,7 +13,15 @@
 
 namespace jit {
 
-Runtime::Runtime() : zero_(Ref<>::create(PyLong_FromLong(0))) {
+Runtime::Runtime()
+    : zero_(Ref<>::create(PyLong_FromLong(0))),
+#if PY_VERSION_HEX >= 0x030C0000
+      str_build_class_(Ref<>::create(&_Py_ID(__build_class__)))
+#else
+      str_build_class_(
+          Ref<>::create(PyUnicode_InternFromString("__build_class__")))
+#endif
+{
 #if PY_VERSION_HEX >= 0x030E0000
   PyObject** common_consts = PyThreadState_GET()->interp->common_consts;
   for (int i = 0; i < NUM_COMMON_CONSTANTS; i++) {

cinderx/Jit/runtime.h

@@ -274,6 +274,10 @@ class Runtime : public IRuntime {
     return zero_.get();
   }
 
+  BorrowedRef<> strBuildClass() {
+    return str_build_class_.get();
+  }
+
   void watchPendingTypes();
   void fixupFunctionEntryCachePostMultiThreadedCompile();
 
@@ -315,6 +319,7 @@ class Runtime : public IRuntime {
       type_deopt_patchers_;
 
   Ref<> zero_;
+  Ref<> str_build_class_;
   std::unordered_set<BorrowedRef<PyTypeObject>> pending_watches_;
 
   std::vector<hir::Type> common_constant_types_;

cinderx/PythonLib/cinderx/test_support.py

@@ -1,7 +1,9 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # pyre-strict
 
+import abc
 import ctypes
+import dis
 import importlib
 import multiprocessing
 import os.path
@@ -212,3 +214,49 @@ def wrapped(*args: object) -> TRet:
         return value
 
     return wrapped
+
+
+class AssertBytecodeContainsMixin(abc.ABC):
+    @abc.abstractmethod
+    def assertIn(
+        self, expected: object, actual: Sequence[object], msg: str | None = None
+    ) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def assertTrue(self, expr: object, msg: str | None = None) -> None:
+        raise NotImplementedError
+
+    def assertBytecodeContains(
+        self,
+        func: object,
+        expected_opcode: str,
+        expected_oparg: int | None = None,
+    ) -> None:
+        try:
+            # pyre-ignore[16] - for things wrapped by fail_if_deopt()
+            inner_function = func.inner_function
+        except AttributeError:
+            pass
+        else:
+            func = inner_function
+
+        bytecode_instructions = dis.get_instructions(func)
+
+        if expected_oparg is None:
+            opcodes = [instr.opname for instr in bytecode_instructions]
+            self.assertIn(
+                expected_opcode,
+                opcodes,
+                f"{expected_opcode} opcode should be present in {func.__name__} bytecode",
+            )
+        else:
+            matching_instructions = [
+                instr
+                for instr in bytecode_instructions
+                if instr.opname == expected_opcode and instr.arg == expected_oparg
+            ]
+            self.assertTrue(
+                len(matching_instructions) > 0,
+                f"{expected_opcode} opcode with oparg {expected_oparg} should be present in {func.__name__} bytecode",
+            )

cinderx/PythonLib/test_cinderx/test_python310_bytecodes.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import sys
+import unittest
+
+import cinderx.test_support as cinder_support
+
+
+@unittest.skipUnless(
+    sys.version_info >= (3, 10) and sys.version_info < (3, 11), "Python 3.10 only"
+)
+class Python310Bytecodes(unittest.TestCase, cinder_support.AssertBytecodeContainsMixin):
+    def test_LOAD_BUILD_CLASS(self):
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x():
+            class TestClass:
+                pass
+
+            return TestClass()
+
+        result = x()
+        self.assertIsNotNone(result)
+        self.assertEqual(result.__class__.__name__, "TestClass")
+        self.assertBytecodeContains(x, "LOAD_BUILD_CLASS")
+
+
+if __name__ == "__main__":
+    unittest.main()

cinderx/PythonLib/test_cinderx/test_python312_bytecodes.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import sys
+import unittest
+
+import cinderx.test_support as cinder_support
+
+
+@unittest.skipUnless(
+    sys.version_info >= (3, 12) and sys.version_info < (3, 13), "Python 3.12 only"
+)
+class Python312Bytecodes(unittest.TestCase, cinder_support.AssertBytecodeContainsMixin):
+    def test_LOAD_BUILD_CLASS(self):
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x():
+            class TestClass:
+                pass
+
+            return TestClass()
+
+        result = x()
+        self.assertIsNotNone(result)
+        self.assertEqual(result.__class__.__name__, "TestClass")
+        self.assertBytecodeContains(x, "LOAD_BUILD_CLASS")
+
+
+if __name__ == "__main__":
+    unittest.main()