Add JIT support for STORE_GLOBAL

Summary: whatcouldgowrong

Reviewed By: alexmalyshev

Differential Revision: D82694078

fbshipit-source-id: fbe27a1a72d13f369574d849e4f88947095af741

Author: jbower-fb

cinderx/Jit/hir/builder.cpp

@@ -228,6 +228,7 @@ bool isSupportedOpcode(int opcode) {
     case STORE_FAST_LOAD_FAST:
     case STORE_FAST_STORE_FAST:
     case STORE_FIELD:
+    case STORE_GLOBAL:
     case STORE_LOCAL:
     case STORE_SLICE:
     case STORE_SUBSCR:
@@ -1558,6 +1559,10 @@ void HIRBuilder::translate(
           emitLoadBuildClass(tc);
           break;
         }
+        case STORE_GLOBAL: {
+          emitStoreGlobal(tc, bc_instr);
+          break;
+        }
         case CHECK_EG_MATCH:
         case CHECK_EXC_MATCH:
         case CLEANUP_THROW:
@@ -5005,6 +5010,25 @@ void HIRBuilder::emitLoadBuildClass(TranslationContext& tc) {
   tc.frame.stack.push(result);
 }
 
+void HIRBuilder::emitStoreGlobal(
+    TranslationContext& tc,
+    const BytecodeInstruction& bc_instr) {
+  Register* globals = temps_.AllocateNonStack();
+  Register* key = temps_.AllocateNonStack();
+
+  tc.emit<LoadConst>(globals, Type::fromObject(tc.frame.globals));
+  // Starting at the preloader the JIT seems to assume globals will be a
+  // dictionary, however I'm not sure there's any guarantee of this.
+  Register* globals_dict = temps_.AllocateNonStack();
+  tc.emit<GuardType>(globals_dict, TDictExact, globals, tc.frame);
+  tc.emit<LoadConst>(
+      key,
+      Type::fromObject(PyTuple_GET_ITEM(code_->co_names, bc_instr.oparg())));
+  Register* value = tc.frame.stack.pop();
+  Register* result = temps_.AllocateNonStack();
+  tc.emit<SetDictItem>(result, globals_dict, key, value, tc.frame);
+}
+
 void HIRBuilder::insertEvalBreakerCheck(
     CFG& cfg,
     BasicBlock* check_block,

cinderx/Jit/hir/builder.h

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

cinderx/Jit/pyjit.cpp

@@ -931,14 +931,20 @@ 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.
+    // Assuming the builtins + globals 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));
+    JIT_DCHECK(
+        PyDict_CheckExact(outer_func->func_globals),
+        "Unexpected type for globals ({}) on function {}",
+        Py_TYPE(outer_func->func_globals)->tp_name,
+        funcFullname(outer_func));
     preloader = hir::Preloader::makePreloader(
         code,
         outer_func->func_builtins,

cinderx/PythonLib/test_cinderx/test_python310_bytecodes.py

@@ -24,6 +24,21 @@ class TestClass:
         self.assertEqual(result.__class__.__name__, "TestClass")
         self.assertBytecodeContains(x, "LOAD_BUILD_CLASS")
 
+    def test_STORE_GLOBAL(self):
+        global test_STORE_GLOBAL_v
+
+        test_STORE_GLOBAL_v = 1
+
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x():
+            global test_STORE_GLOBAL_v
+            test_STORE_GLOBAL_v = 42
+
+        x()
+        self.assertEqual(test_STORE_GLOBAL_v, 42)
+        self.assertBytecodeContains(x, "STORE_GLOBAL")
+
 
 if __name__ == "__main__":
     unittest.main()

cinderx/PythonLib/test_cinderx/test_python312_bytecodes.py

@@ -24,6 +24,21 @@ class TestClass:
         self.assertEqual(result.__class__.__name__, "TestClass")
         self.assertBytecodeContains(x, "LOAD_BUILD_CLASS")
 
+    def test_STORE_GLOBAL(self):
+        global test_STORE_GLOBAL_v
+
+        test_STORE_GLOBAL_v = 1
+
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x():
+            global test_STORE_GLOBAL_v
+            test_STORE_GLOBAL_v = 42
+
+        x()
+        self.assertEqual(test_STORE_GLOBAL_v, 42)
+        self.assertBytecodeContains(x, "STORE_GLOBAL")
+
 
 if __name__ == "__main__":
     unittest.main()

cinderx/PythonLib/test_cinderx/test_python314_bytecodes.py

@@ -488,6 +488,21 @@ class TestClass:
         self.assertEqual(result.__class__.__name__, "TestClass")
         self.assertBytecodeContains(x, "LOAD_BUILD_CLASS")
 
+    def test_STORE_GLOBAL(self):
+        global test_STORE_GLOBAL_v
+
+        test_STORE_GLOBAL_v = 1
+
+        @cinder_support.fail_if_deopt
+        @cinder_support.failUnlessJITCompiled
+        def x():
+            global test_STORE_GLOBAL_v
+            test_STORE_GLOBAL_v = 42
+
+        x()
+        self.assertEqual(test_STORE_GLOBAL_v, 42)
+        self.assertBytecodeContains(x, "STORE_GLOBAL")
+
 
 if __name__ == "__main__":
     unittest.main()