Minor cleanup on Jit/code_runtime.h

Summary:
Hides the implementation, types globals and builtins as dicts, fixes up and adds
comments, and shrinks GenYieldPoint by 8 bytes.

Reviewed By: mpage

Differential Revision: D83566613

fbshipit-source-id: 4bb7635c56d40d3f4b53f2c14c3eb3201c17326d

Author: alexmalyshev

cinderx/Jit/code_runtime.cpp

@@ -2,15 +2,69 @@
 
 #include "cinderx/Jit/code_runtime.h"
 
+#include "cinderx/Common/util.h"
+
 namespace jit {
 
-const int64_t CodeRuntime::kPyCodeOffset =
-    RuntimeFrameState::codeOffset() + CodeRuntime::frameStateOffset();
+GenYieldPoint::GenYieldPoint(std::size_t deopt_idx, ptrdiff_t yield_from_offset)
+    : deopt_idx_{deopt_idx}, yield_from_offset_{yield_from_offset} {}
 
-void CodeRuntime::releaseReferences() {
-  // Serialize as we modify ref-counts which may be widely accessible.
-  ThreadedCompileSerialize guard;
-  references_.clear();
+void GenYieldPoint::setResumeTarget(uintptr_t resume_target) {
+  resume_target_ = resume_target;
+}
+
+uintptr_t GenYieldPoint::resumeTarget() const {
+  return resume_target_;
+}
+
+std::size_t GenYieldPoint::deoptIdx() const {
+  return deopt_idx_;
+}
+
+bool GenYieldPoint::isYieldFrom() const {
+  return yield_from_offset_ != kInvalidYieldFromOffset;
+}
+
+ptrdiff_t GenYieldPoint::yieldFromOffset() const {
+  return yield_from_offset_;
+}
+
+bool RuntimeFrameState::isGen() const {
+  return code()->co_flags & kCoFlagsAnyGenerator;
+}
+
+BorrowedRef<PyCodeObject> RuntimeFrameState::code() const {
+  return code_;
+}
+
+BorrowedRef<PyDictObject> RuntimeFrameState::builtins() const {
+  return builtins_;
+}
+
+BorrowedRef<PyDictObject> RuntimeFrameState::globals() const {
+  return globals_;
+}
+
+BorrowedRef<PyFunctionObject> RuntimeFrameState::func() const {
+  return func_;
+}
+
+CodeRuntime::CodeRuntime(BorrowedRef<PyFunctionObject> func)
+    : CodeRuntime{
+          BorrowedRef<PyCodeObject>{func->func_code},
+          func->func_builtins,
+          func->func_globals} {}
+
+CodeRuntime::CodeRuntime(
+    BorrowedRef<PyCodeObject> code,
+    BorrowedRef<PyDictObject> builtins,
+    BorrowedRef<PyDictObject> globals)
+    : frame_state_{code, builtins, globals} {
+  // Ensure code, globals, and builtins objects live as long as their compiled
+  // functions.
+  addReference(code);
+  addReference(builtins);
+  addReference(globals);
 }
 
 void CodeRuntime::addReference(BorrowedRef<> obj) {
@@ -19,4 +73,31 @@ void CodeRuntime::addReference(BorrowedRef<> obj) {
   references_.emplace(ThreadedRef<>::create(obj));
 }
 
+void CodeRuntime::releaseReferences() {
+  // Serialize as we modify ref-counts which may be widely accessible.
+  ThreadedCompileSerialize guard;
+  references_.clear();
+}
+
+GenYieldPoint* CodeRuntime::addGenYieldPoint(GenYieldPoint&& gen_yield_point) {
+  gen_yield_points_.emplace_back(std::move(gen_yield_point));
+  return &gen_yield_points_.back();
+}
+
+const RuntimeFrameState* CodeRuntime::frameState() const {
+  return &frame_state_;
+}
+
+int CodeRuntime::frameSize() const {
+  return frame_size_;
+}
+
+void CodeRuntime::setFrameSize(int size) {
+  frame_size_ = size;
+}
+
+DebugInfo* CodeRuntime::debugInfo() {
+  return &debug_info_;
+}
+
 } // namespace jit

cinderx/Jit/code_runtime.h

@@ -9,188 +9,132 @@
 #endif
 
 #include "cinderx/Common/ref.h"
-#include "cinderx/Common/util.h"
 #include "cinderx/Jit/debug_info.h"
-#include "cinderx/Jit/frame_header.h"
 #include "cinderx/Jit/threaded_compile.h"
 
 #include <deque>
-#include <unordered_map>
+#include <limits>
+#include <unordered_set>
 
 namespace jit {
 
+constexpr ptrdiff_t kInvalidYieldFromOffset =
+    std::numeric_limits<ptrdiff_t>::max();
+
+// Information about how a specific yield instruction should resume.
 class GenYieldPoint {
  public:
-  explicit GenYieldPoint(
-      std::size_t deopt_idx,
-      bool is_yield_from,
-      ptrdiff_t yield_from_offs)
-      : deopt_idx_(deopt_idx),
-        isYieldFrom_(is_yield_from),
-        yieldFromOffs_(yield_from_offs) {}
-
-  void setResumeTarget(uint64_t resume_target) {
-    resume_target_ = resume_target;
-  }
-
-  uint64_t resumeTarget() const {
-    return resume_target_;
-  }
-
-  std::size_t deoptIdx() const {
-    return deopt_idx_;
+  static constexpr int resumeTargetOffset() {
+    return offsetof(GenYieldPoint, resume_target_);
   }
 
-  bool isYieldFrom() const {
-    return isYieldFrom_;
-  }
+  GenYieldPoint(std::size_t deopt_idx, ptrdiff_t yield_from_offset);
 
-  ptrdiff_t yieldFromOffset() const {
-    return yieldFromOffs_;
-  }
+  // Get and set what address the yield should resume from.
+  uintptr_t resumeTarget() const;
+  void setResumeTarget(uintptr_t resume_target);
 
-  static constexpr int resumeTargetOffset() {
-    return offsetof(GenYieldPoint, resume_target_);
-  }
+  std::size_t deoptIdx() const;
+  bool isYieldFrom() const;
+  ptrdiff_t yieldFromOffset() const;
 
  private:
-  uint64_t resume_target_{0};
+  uintptr_t resume_target_{0};
   const std::size_t deopt_idx_;
-  const bool isYieldFrom_;
-  const ptrdiff_t yieldFromOffs_;
+  const ptrdiff_t yield_from_offset_;
 };
 
 class alignas(16) RuntimeFrameState {
  public:
+  static constexpr int64_t codeOffset() {
+    return offsetof(RuntimeFrameState, code_);
+  }
+
   RuntimeFrameState(
       BorrowedRef<PyCodeObject> code,
-      BorrowedRef<> builtins,
-      BorrowedRef<> globals,
+      BorrowedRef<PyDictObject> builtins,
+      BorrowedRef<PyDictObject> globals,
       BorrowedRef<PyFunctionObject> func = nullptr)
-      : code_(code), builtins_(builtins), globals_(globals), func_(func) {}
+      : code_{code}, builtins_{builtins}, globals_{globals}, func_{func} {}
 
-  bool isGen() const {
-    return code()->co_flags & kCoFlagsAnyGenerator;
-  }
-
-  BorrowedRef<PyCodeObject> code() const {
-    return code_;
-  }
-
-  BorrowedRef<> builtins() const {
-    return builtins_;
-  }
-
-  BorrowedRef<> globals() const {
-    return globals_;
-  }
+  // Check if this is a generator frame.
+  bool isGen() const;
 
-  BorrowedRef<PyFunctionObject> func() const {
-    return func_;
-  }
-
-  static constexpr int64_t codeOffset() {
-    return offsetof(RuntimeFrameState, code_);
-  }
+  BorrowedRef<PyCodeObject> code() const;
+  BorrowedRef<PyDictObject> builtins() const;
+  BorrowedRef<PyDictObject> globals() const;
+  BorrowedRef<PyFunctionObject> func() const;
 
  private:
-  // These are owned by the CodeRuntime that owns this RuntimeFrameState.
+  // All fields are owned by the CodeRuntime that owns this RuntimeFrameState.
+
   BorrowedRef<PyCodeObject> code_;
-  BorrowedRef<> builtins_;
-  BorrowedRef<> globals_;
+  BorrowedRef<PyDictObject> builtins_;
+  BorrowedRef<PyDictObject> globals_;
+  // The function is only set for inlined frames.
   BorrowedRef<PyFunctionObject> func_;
 };
 
 // Runtime data for a PyCodeObject object, containing caches and any other data
 // associated with a JIT-compiled function.
 class alignas(16) CodeRuntime {
  public:
-  CodeRuntime(
-      PyCodeObject* code,
-      PyObject* builtins,
-      PyObject* globals,
-      PyFunctionObject* func = nullptr)
-      : frame_state_(code, builtins, globals, func) {
-    // TASK(T88040922): Until we work out something smarter, force code,
-    // globals, and builtins objects for compiled functions to live as long as
-    // the JIT is initialized.
-    addReference(BorrowedRef(code));
-    addReference(builtins);
-    addReference(globals);
-    if (func != nullptr) {
-      addReference(&func->ob_base);
-    }
+  static constexpr int64_t frameStateOffset() {
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Winvalid-offsetof"
+    return offsetof(CodeRuntime, frame_state_);
+#pragma GCC diagnostic pop
   }
 
-  explicit CodeRuntime(PyFunctionObject* func)
-      : CodeRuntime(
-            reinterpret_cast<PyCodeObject*>(func->func_code),
-            func->func_builtins,
-            func->func_globals) {}
+  static constexpr int64_t codeOffset() {
+    return CodeRuntime::frameStateOffset() + RuntimeFrameState::codeOffset();
+  }
+
+  explicit CodeRuntime(BorrowedRef<PyFunctionObject> func);
+  CodeRuntime(
+      BorrowedRef<PyCodeObject> code,
+      BorrowedRef<PyDictObject> builtins,
+      BorrowedRef<PyDictObject> globals);
 
   template <typename... Args>
   RuntimeFrameState* allocateRuntimeFrameState(Args&&... args) {
-    // Serialize as we modify the globally shared runtimes data.
-    ThreadedCompileSerialize guard;
-    inlined_frame_states_.emplace_back(
-        std::make_unique<RuntimeFrameState>(std::forward<Args>(args)...));
-    return inlined_frame_states_.back().get();
-  }
-
-  const RuntimeFrameState* frameState() const {
-    return &frame_state_;
+    return inlined_frame_states_
+        .emplace_back(
+            std::make_unique<RuntimeFrameState>(std::forward<Args>(args)...))
+        .get();
   }
 
-  // Release any references this CodeRuntime holds to Python objects.
-  void releaseReferences();
-
-  // Ensure that this CodeRuntime owns a reference to the given owned object,
-  // keeping it alive for use by the compiled code. Transfer ownership of the
-  // object to the CodeRuntime.
-  void addReference(Ref<>&& obj);
-
   // Ensure that this CodeRuntime owns a reference to the given borrowed
   // object, keeping it alive for use by the compiled code. Make CodeRuntime a
   // new owner of the object.
   void addReference(BorrowedRef<> obj);
 
+  // Release any references this CodeRuntime holds to Python objects.
+  void releaseReferences();
+
   // Store meta-data about generator yield point.
-  GenYieldPoint* addGenYieldPoint(GenYieldPoint&& gen_yield_point) {
-    gen_yield_points_.emplace_back(std::move(gen_yield_point));
-    return &gen_yield_points_.back();
-  }
+  GenYieldPoint* addGenYieldPoint(GenYieldPoint&& gen_yield_point);
 
-  void set_frame_size(int size) {
-    frame_size_ = size;
-  }
-  int frame_size() const {
-    return frame_size_;
-  }
+  // Get the top-level runtime frame state for this CodeRuntime's PyCodeObject.
+  const RuntimeFrameState* frameState() const;
 
-  DebugInfo* debug_info() {
-    return &debug_info_;
-  }
+  // Get and set the total size of a stack frame for this compiled code object.
+  int frameSize() const;
+  void setFrameSize(int size);
 
-  static constexpr int64_t frameStateOffset() {
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Winvalid-offsetof"
-    return offsetof(CodeRuntime, frame_state_);
-#pragma GCC diagnostic pop
-  }
-
-  static const int64_t kPyCodeOffset;
+  DebugInfo* debugInfo();
 
  private:
   RuntimeFrameState frame_state_;
   std::vector<std::unique_ptr<RuntimeFrameState>> inlined_frame_states_;
 
+  // References owned by this CodeRuntime.
   std::unordered_set<ThreadedRef<PyObject>> references_;
 
   // Metadata about yield points. Deque so we can have raw pointers to content.
   std::deque<GenYieldPoint> gen_yield_points_;
 
   int frame_size_{-1};
-
   DebugInfo debug_info_;
 };
 

cinderx/Jit/codegen/autogen.cpp

@@ -366,9 +366,10 @@ void emitStoreGenYieldPoint(
       live_regs_input - num_live_regs,
       live_regs_input);
 
-  auto yield_from_offset = is_yield_from ? calc_spill_offset(2) : 0;
+  auto yield_from_offset =
+      is_yield_from ? calc_spill_offset(2) : kInvalidYieldFromOffset;
   GenYieldPoint* gen_yield_point = env->code_rt->addGenYieldPoint(
-      GenYieldPoint{deopt_idx, is_yield_from, yield_from_offset});
+      GenYieldPoint{deopt_idx, yield_from_offset});
 
   env->unresolved_gen_entry_labels.emplace(gen_yield_point, resume_label);
   if (yield->origin()) {