Add support for LOAD_COMMON_CONSTANT

Summary: We pre-compute the `Type` objects so we don't have to worry about getting an interperter instance during JIT compilation (not compatible with multi-threaded compile).

Reviewed By: DinoV

Differential Revision: D82054730

fbshipit-source-id: 1ef69c5b9c662fd6ee5d7c9c0673fefc7ccbcb03

Author: jbower-fb

cinderx/Common/opcode_stubs.h

@@ -69,6 +69,7 @@
   X(JUMP_BACKWARD)                     \
   X(JUMP_BACKWARD_NO_INTERRUPT)        \
   X(KW_NAMES)                          \
+  X(LOAD_COMMON_CONSTANT)              \
   X(LOAD_FAST_AND_CLEAR)               \
   X(LOAD_FAST_CHECK)                   \
   X(LOAD_FAST_BORROW)                  \
@@ -183,6 +184,7 @@ enum {
   X(LOAD_ATTR_S_MODULE)                \
   X(LOAD_ATTR_TYPE)                    \
   X(LOAD_ATTR_UNCACHABLE)              \
+  X(LOAD_COMMON_CONSTANT)              \
   X(LOAD_FAST_BORROW)                  \
   X(LOAD_FAST_BORROW_LOAD_FAST_BORROW) \
   X(LOAD_FAST_LOAD_FAST)               \

cinderx/Jit/hir/builder.cpp

@@ -15,6 +15,7 @@
 #include "cinderx/Jit/containers.h"
 #include "cinderx/Jit/hir/ssa.h"
 #include "cinderx/Jit/hir/type.h"
+#include "cinderx/Jit/runtime.h"
 #include "cinderx/StaticPython/checked_dict.h"
 #include "cinderx/StaticPython/checked_list.h"
 #include "cinderx/StaticPython/classloader.h"
@@ -149,6 +150,7 @@ bool isSupportedOpcode(int opcode) {
     case LOAD_ATTR:
     case LOAD_ATTR_SUPER:
     case LOAD_CLOSURE:
+    case LOAD_COMMON_CONSTANT:
     case LOAD_CONST:
     case LOAD_DEREF:
     case LOAD_FAST:
@@ -1529,6 +1531,10 @@ void HIRBuilder::translate(
           emitFormatSimple(irfunc.cfg, tc);
           break;
         }
+        case LOAD_COMMON_CONSTANT: {
+          emitLoadCommonConstant(tc, bc_instr);
+          break;
+        }
         case CHECK_EG_MATCH:
         case CHECK_EXC_MATCH:
         case CLEANUP_THROW:
@@ -4889,6 +4895,15 @@ void HIRBuilder::emitFormatSimple(CFG& cfg, TranslationContext& tc) {
   stack.push(out);
 }
 
+void HIRBuilder::emitLoadCommonConstant(
+    TranslationContext& tc,
+    const BytecodeInstruction& bc_instr) {
+  Register* out = temps_.AllocateStack();
+  tc.emit<LoadConst>(
+      out, Runtime::get()->typeForCommonConstant(bc_instr.oparg()));
+  tc.frame.stack.push(out);
+}
+
 void HIRBuilder::insertEvalBreakerCheck(
     CFG& cfg,
     BasicBlock* check_block,

cinderx/Jit/hir/builder.h

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

cinderx/Jit/runtime.cpp

@@ -13,6 +13,16 @@
 
 namespace jit {
 
+Runtime::Runtime() : zero_(Ref<>::create(PyLong_FromLong(0))) {
+#if PY_VERSION_HEX >= 0x030E0000
+  PyObject** common_consts = PyThreadState_GET()->interp->common_consts;
+  for (int i = 0; i < NUM_COMMON_CONSTANTS; i++) {
+    common_constant_types_.emplace_back(
+        hir::Type::fromObject(common_consts[i]));
+  }
+#endif
+}
+
 void Builtins::init() {
   ThreadedCompileSerialize guard;
   if (is_initialized_) {

cinderx/Jit/runtime.h

@@ -102,9 +102,8 @@ class Builtins {
 // Runtime owns all metadata created by the JIT.
 class Runtime : public IRuntime {
  public:
-  Runtime() {
-    zero_ = Ref<>::create(PyLong_FromLong(0));
-  }
+  Runtime();
+
   // Return the singleton Runtime, creating it first if necessary.
   static Runtime* get() {
     return static_cast<Runtime*>(cinderx::getModuleState()->runtime());
@@ -278,7 +277,13 @@ class Runtime : public IRuntime {
   void watchPendingTypes();
   void fixupFunctionEntryCachePostMultiThreadedCompile();
 
- private:
+  const hir::Type& typeForCommonConstant([[maybe_unused]] int i) const {
+#if PY_VERSION_HEX >= 0x030E0000
+    return common_constant_types_.at(i);
+#endif
+    JIT_ABORT("Common constants are a feature of 3.14+");
+  }
+
   // 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_;
@@ -311,6 +316,8 @@ class Runtime : public IRuntime {
 
   Ref<> zero_;
   std::unordered_set<BorrowedRef<PyTypeObject>> pending_watches_;
+
+  std::vector<hir::Type> common_constant_types_;
 };
 
 } // namespace jit

cinderx/PythonLib/test_cinderx/test_python314_bytecodes.py

@@ -431,6 +431,15 @@ def x():
             locals,
         )
 
+    def test_LOAD_COMMON_CONSTANT(self):
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x(a):
+            assert a
+
+        x(True)
+        self._assertBytecodeContains(x, "LOAD_COMMON_CONSTANT")
+
 
 if __name__ == "__main__":
     unittest.main()