Flow thread state in and frame out

Summary:
Currently for allocating a generator frame we use `PyThreadState_GET` to acquire the thread state in our runtime helper and then return it to the JITed function and then we pull the current frame from the thread state.

But for any non-BOLTed build we're going to be able to acquire the thread state faster in JITed code through our runtime resolved address. Freeing up the return value also let's us get rid of the dependent loads to get the frame object. Overall that's less instructions as well. We end up saving 3 instructions but 1 of those is due to an inefficiency in `adr` and the other is due to a register allocator change.

We could just pass 5 arguments into this function but that'll complicate the call on Windows so we shuffle the resume entry point into the code runtime which we're already reading anyway.

Reviewed By: alexmalyshev

Differential Revision: D114815510

fbshipit-source-id: 309cb32eb8e8d95a3e48e86869b7b4c158d81fb8

Author: DinoV

cinderx/Jit/code_runtime.cpp

@@ -130,6 +130,14 @@ void CodeRuntime::setSpillWords(uint32_t words) {
   spill_words_ = words;
 }
 
+GenResumeFunc CodeRuntime::genResumeEntry() const {
+  return gen_resume_entry_;
+}
+
+void CodeRuntime::setGenResumeEntry(GenResumeFunc resume_entry) {
+  gen_resume_entry_ = resume_entry;
+}
+
 DebugInfo* CodeRuntime::debugInfo() {
   return &debug_info_;
 }

cinderx/Jit/code_runtime.h

@@ -3,6 +3,7 @@
 #pragma once
 
 #include "cinderx/Common/ref.h"
+#include "cinderx/Common/util.h"
 #include "cinderx/Jit/debug_info.h"
 #include "cinderx/Jit/deopt.h"
 #include "cinderx/Jit/threaded_compile.h"
@@ -94,6 +95,12 @@ class alignas(16) CodeRuntime {
   uint32_t spillWords() const;
   void setSpillWords(uint32_t words);
 
+  // Get and set the address a generator resumes execution at.  Only meaningful
+  // for generators, and only resolvable once code generation has bound the
+  // resume entry label to an address.
+  GenResumeFunc genResumeEntry() const;
+  void setGenResumeEntry(GenResumeFunc resume_entry);
+
   DebugInfo* debugInfo();
 
   // Allocate a jump table for static type check dispatch.
@@ -154,6 +161,8 @@ class alignas(16) CodeRuntime {
   // Set after CompiledFunction::create() in makeCompiledFunction().
   BorrowedRef<CompiledFunction> compiled_function_;
 
+  GenResumeFunc gen_resume_entry_{nullptr};
+
   int frame_size_{-1};
   uint32_t spill_words_{0};
   DebugInfo debug_info_;

cinderx/Jit/codegen/gen_asm.cpp

@@ -1431,6 +1431,17 @@ void NativeGenerator::generateCode(
         codeholder.baseAddress());
   }
 
+  // allocateAndLinkGenAndInterpreterFrame reads the resume entry out of the
+  // CodeRuntime rather than taking it as an argument, which keeps that call
+  // down to four arguments.
+  if (getFunction()->code->co_flags & kCoFlagsAnyGenerator) {
+    auto resume_entry = static_cast<uintptr_t>(
+        codeholder.baseAddress() +
+        codeholder.labelOffsetFromBase(env_.gen_resume_entry_label));
+    env_.code_rt->setGenResumeEntry(
+        reinterpret_cast<GenResumeFunc>(resume_entry));
+  }
+
   // Resolve the static type check jump table entries now that block labels
   // have been bound to code addresses.
   for (auto& [index, block] : env_.static_typecheck_jt_entries) {

cinderx/Jit/jit_rt.cpp

@@ -716,10 +716,11 @@ void initFrameCellVars(
 #endif
 }
 
-std::pair<PyThreadState*, GenDataFooter*> allocateAndLinkGenAndInterpreterFrame(
+std::pair<_PyInterpreterFrame*, GenDataFooter*>
+allocateAndLinkGenAndInterpreterFrame(
+    PyThreadState* tstate,
     PyFunctionObject* func,
     CodeRuntime* code_rt,
-    GenResumeFunc resume_func,
     uint64_t original_frame_pointer) {
   JIT_DCHECK(
       PyCode_Check(func->func_code),
@@ -729,7 +730,6 @@ std::pair<PyThreadState*, GenDataFooter*> allocateAndLinkGenAndInterpreterFrame(
   JIT_DCHECK(co == code_rt->code(), "Code object mismatch");
 
   uint64_t spill_words = code_rt->spillWords();
-  PyThreadState* tstate = PyThreadState_GET();
   JIT_DCHECK(tstate != nullptr, "thread state cannot be null");
   auto [gen, gen_size] = cinderx::getModuleState()->jit_gen_free_list->allocate(
       co, spill_words * sizeof(uint64_t) + sizeof(GenDataFooter));
@@ -777,6 +777,11 @@ std::pair<PyThreadState*, GenDataFooter*> allocateAndLinkGenAndInterpreterFrame(
   init_and_link_interpreter_frame(
       func, co, tstate, FRAME_OWNED_BY_GENERATOR, frame, code_rt);
 
+  GenResumeFunc resume_func = code_rt->genResumeEntry();
+  JIT_DCHECK(
+      resume_func != nullptr,
+      "CodeRuntime has no resume entry: {}",
+      PyUnicode_AsUTF8(co->co_qualname));
   footer->resumeEntry = resume_func;
   footer->yieldPoint = nullptr;
   footer->gen = static_cast<PyGenObject*>(gen);
@@ -812,7 +817,7 @@ std::pair<PyThreadState*, GenDataFooter*> allocateAndLinkGenAndInterpreterFrame(
 
   PyObject_GC_Track(gen);
 
-  return {tstate, footer};
+  return {frame, footer};
 }
 
 std::pair<JitGenObject*, GenDataFooter*> unlinkGenFrameAndReturnGenDataFooter(

cinderx/Jit/jit_rt.h

@@ -4,6 +4,12 @@
 
 #include "cinderx/python.h"
 
+#include "internal/pycore_frame.h"
+
+#if PY_VERSION_HEX >= 0x030E0000
+#include "internal/pycore_interpframe_structs.h"
+#endif
+
 #include "cinderx/Common/util.h"
 #include "cinderx/StaticPython/typed-args-info.h"
 
@@ -36,10 +42,14 @@ PyThreadState* allocateAndLinkInterpreterFrame_Debug(
 
 PyThreadState* allocateAndLinkInterpreterFrame_Release(PyFunctionObject* func);
 
-std::pair<PyThreadState*, GenDataFooter*> allocateAndLinkGenAndInterpreterFrame(
+// Allocate a generator and its interpreter frame, and link the frame into
+// tstate.  Returns the newly linked interpreter frame and the generator's
+// GenDataFooter.
+std::pair<_PyInterpreterFrame*, GenDataFooter*>
+allocateAndLinkGenAndInterpreterFrame(
+    PyThreadState* tstate,
     PyFunctionObject* func,
     CodeRuntime* code_rt,
-    GenResumeFunc resume_entry,
     uint64_t original_frame_pointer);
 
 void initFrameCellVars(

cinderx/Jit/lir/generator.cpp

@@ -5494,32 +5494,30 @@ void LIRGenerator::emitLoadFrame(BasicBlockBuilder& bbb) {
   bool is_gen =
       func_->code != nullptr && (func_->code->co_flags & kCoFlagsAnyGenerator);
   if (is_gen) {
-    // Load the resume entry label address.
+    bbb.annotateNext("Allocate generator + interpreter frame: load tstate");
+    env_->asm_tstate =
+        bbb.appendInstr(OutVReg{}, Instruction::kLoadThreadState);
+    // spill_words and the resume entry address are read from CodeRuntime by
+    // the runtime function, so we don't need to pass them explicitly.
     bbb.annotateNext("Allocate generator + interpreter frame");
-    Instruction* resume_label = bbb.appendInstr(
-        OutVReg{DataType::k64bit},
-        Instruction::kLea,
-        AsmLbl{env_->gen_resume_entry_label});
-    // spill_words is read from CodeRuntime by the runtime function,
-    // so we don't need to pass it explicitly.
+    Instruction* frame;
     Instruction* footer;
     appendCall2RetValues(
         bbb,
-        env_->asm_tstate,
+        frame,
         footer,
         rt::allocateAndLinkGenAndInterpreterFrame,
+        env_->asm_tstate,
         env_->asm_func,
         Imm{reinterpret_cast<uint64_t>(env_->code_rt)},
-        VReg{resume_label},
         PhyReg{codegen::arch::reg_frame_pointer_loc});
     // Swap RBP to point at the generator data so spills go there.
     bbb.annotateNext("Set frame pointer to GenDataFooter");
     bbb.appendInstr(
         OutPhyReg{codegen::arch::reg_frame_pointer_loc},
         Instruction::kMove,
         VReg{footer});
-    bbb.annotateNext("Load current interpreter frame");
-    env_->asm_interpreter_frame = makeCurrentFrameAccessor(bbb).load();
+    env_->asm_interpreter_frame = frame;
 #if (defined(CINDER_AARCH64) || defined(Py_GIL_DISABLED)) && \
     defined(ENABLE_LIGHTWEIGHT_FRAMES)
     // Now that FP points at the heap-allocated GenDataFooter, compute the