Directly embed deopt idx instead of return address

Summary:
On ARM there's no way to reliably recover where the return address is stored by the callee. They are free to first store some additional data and then setup the frame - for example they may initially save callee saved registers to ensure that the values in their frame that they care about are closer. We could do something unsafe like search the stack for our frame pointer and then gather the IP address from there but if a callee preserved register collided with our frame pointer we'd return incorrect results.

Instead this removes storing the return address for every call and instead stores the deopt index. This takes the same amount of instructions but we do it less as not every call is a deopt. We now typically generate:
```
mov x14, #1
str x14, [x19]
```
Instead of:
```
adr x13, #0xfffe3633032c
str x13, [x29, #0x10]
```

But for example we generate none of these for decref where we previously generated 2.

We could instead save an index which is the bytecode offset for each line number as that's what we really care about currently but the free-threaded work is going to drive us to actually care about the deopt index instead so we may as well just start there.

Reviewed By: alexmalyshev

Differential Revision: D97524722

fbshipit-source-id: 2fb51c6848451496f340ccee42548829705c26f7

Author: DinoV

cinderx/Jit/code_runtime.cpp

@@ -153,4 +153,18 @@ int CodeRuntime::traverse(visitproc visit, void* arg) {
   return 0;
 }
 
+std::optional<UnitCallStack> CodeRuntime::getUnitCallStackFromDeoptIdx(
+    std::size_t deopt_idx) const {
+  if (deopt_idx >= deopt_metadatas_.size()) {
+    return std::nullopt;
+  }
+  const DeoptMetadata& meta = deopt_metadatas_[deopt_idx];
+  UnitCallStack stack;
+  stack.reserve(meta.frame_meta.size());
+  for (const auto& frame : meta.frame_meta) {
+    stack.emplace_back(frame.code, frame.cause_instr_idx);
+  }
+  return stack;
+}
+
 } // namespace jit

cinderx/Jit/code_runtime.h

@@ -136,6 +136,10 @@ class alignas(16) CodeRuntime {
   // True if the references have been cleared
   bool isCleared() const;
 
+  // Get the UnitCallStack from a deopt metadata index.
+  std::optional<UnitCallStack> getUnitCallStackFromDeoptIdx(
+      std::size_t deopt_idx) const;
+
 #if PY_VERSION_HEX >= 0x030E0000 && defined(ENABLE_LIGHTWEIGHT_FRAMES)
   void setReifier(BorrowedRef<> reifier) {
     ThreadedCompileSerialize guard;

cinderx/Jit/codegen/arch.h

@@ -86,6 +86,9 @@ constexpr auto reg_stack_pointer_loc = SP;
 constexpr auto fp = asmjit::a64::x29;
 constexpr auto lr = asmjit::a64::x30;
 
+// Size of the AArch64 frame record: saved FP + LR (two 64-bit registers).
+constexpr int kFrameRecordSize = 2 * sizeof(void*);
+
 } // namespace jit::codegen::arch
 
 #else

cinderx/Jit/codegen/autogen.cpp

@@ -1333,7 +1333,7 @@ void translateEpilogueEnd(Environ* env, const Instruction* instr) {
     restoreCalleeSavedRegsAarch64(as, saved_regs);
   }
   as->mov(a64::sp, arch::fp);
-  as->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 32));
+  as->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
   as->ret(arch::lr);
 #else
   CINDER_UNSUPPORTED
@@ -1935,32 +1935,21 @@ void translateCall(Environ* env, const Instruction* instr) {
   auto output = instr->output();
   auto input = instr->getInput(0);
 
-  // Load call target into a register, then save the return address at
-  // [fp + 16] before calling. This allows getIP() to find the return address
-  // at a fixed offset from the frame base, which is needed for cross-thread
-  // frame inspection (e.g. sys._current_frames()).
-  a64::Gp target;
   if (input->isReg()) {
-    target = AT::getGp(input);
+    as->blr(AT::getGp(input));
   } else if (input->isImm()) {
     as->mov(arch::reg_scratch_br, input->getConstant());
-    target = arch::reg_scratch_br;
+    as->blr(arch::reg_scratch_br);
   } else if (input->isStack()) {
     auto loc = input->getStackSlot().loc;
     as->ldr(
         arch::reg_scratch_br,
         arch::ptr_resolve(as, arch::fp, loc, arch::reg_scratch_0));
-    target = arch::reg_scratch_br;
+    as->blr(arch::reg_scratch_br);
   } else {
     JIT_ABORT("Unsupported operand type for Call: {}", input->type());
   }
 
-  asmjit::Label after_call = as->newLabel();
-  as->adr(arch::reg_scratch_0, after_call);
-  as->str(arch::reg_scratch_0, a64::ptr(arch::fp, 16));
-  as->blr(target);
-  as->bind(after_call);
-
   if (instr->origin()) {
     asmjit::Label label = as->newLabel();
     as->bind(label);

cinderx/Jit/codegen/frame_asm.cpp

@@ -681,13 +681,35 @@ void FrameAsm::linkLightWeightFunctionFrame(
   arch::sub_immediate(as_, arch::reg_scratch_0, arch::fp, frame_header_size);
   as_->str(a64::xzr, a64::ptr(arch::reg_scratch_0));
   env_.addAnnotation("Store rtfs state to 0", store_func_cursor);
+  // Initialize deopt_idx to 0 in FrameHeader.
+  // reg_scratch_0 already holds the frame header base address from above.
+  {
+    asmjit::BaseNode* store_deopt_idx_cursor_arm = as_->cursor();
+    as_->str(
+        a64::xzr,
+        a64::ptr(
+            arch::reg_scratch_0,
+            static_cast<int>(offsetof(FrameHeader, deopt_idx))));
+    env_.addAnnotation("Store deopt_idx to 0", store_deopt_idx_cursor_arm);
+  }
 #else
   // Initialize the fields minus previous.
   // Store func before the header
   arch::sub_immediate(as_, arch::reg_scratch_0, arch::fp, frame_header_size);
   as_->str(func_reg, a64::ptr(arch::reg_scratch_0));
   incRef(func_reg, ref_cnt, ref_cnt_scratch, tstate_reg);
   env_.addAnnotation("Store func before frame header", store_func_cursor);
+  // Initialize deopt_idx to 0 in FrameHeader (3.12 arm64).
+  // reg_scratch_0 still holds the frame header base address.
+  {
+    asmjit::BaseNode* store_deopt_idx_cursor_312_arm = as_->cursor();
+    as_->str(
+        a64::xzr,
+        a64::ptr(
+            arch::reg_scratch_0,
+            static_cast<int>(offsetof(FrameHeader, deopt_idx))));
+    env_.addAnnotation("Store deopt_idx to 0", store_deopt_idx_cursor_312_arm);
+  }
 #endif
 
   asmjit::BaseNode* store_f_code_cursor = as_->cursor();
@@ -821,7 +843,7 @@ void FrameAsm::linkLightWeightFunctionFrame(
 
   // Then finally link in our frame to thread state
   asmjit::BaseNode* update_linkage_cursor = as_->cursor();
-  int size = -frame_header_size + sizeof(PyObject*);
+  int size = -frame_header_size + sizeof(FrameHeader);
   arch::add_signed_immediate(as_, scratch, arch::fp, size);
 
 #if PY_VERSION_HEX >= 0x030D0000

cinderx/Jit/codegen/gen_asm.cpp

@@ -850,7 +850,7 @@ void* generateDeoptTrampoline(bool generator_mode) {
   a.ldr(arch::reg_scratch_br, epilogue_addr);
   // Remove our frame from the stack
   a.mov(a64::sp, arch::fp);
-  a.ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 16));
+  a.ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
   a.br(arch::reg_scratch_br);
   annot.add("Jump to real epilogue", &a, annot_cursor);
 
@@ -910,7 +910,7 @@ void* generateFailedDeferredCompileTrampoline() {
 #elif defined(CINDER_AARCH64)
   auto annot_cursor = a.cursor();
 
-  a.stp(arch::fp, arch::lr, a64::ptr_pre(a64::sp, -16));
+  a.stp(arch::fp, arch::lr, a64::ptr_pre(a64::sp, -arch::kFrameRecordSize));
   a.mov(arch::fp, a64::sp);
 
   // save incoming arg registers
@@ -925,7 +925,7 @@ void* generateFailedDeferredCompileTrampoline() {
   a.mov(arch::reg_scratch_br, JITRT_FailedDeferredCompileShim);
   a.blr(arch::reg_scratch_br);
   a.mov(a64::sp, arch::fp);
-  a.ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 16));
+  a.ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
   a.ret(arch::lr);
 #else
   CINDER_UNSUPPORTED
@@ -1276,7 +1276,7 @@ void NativeGenerator::generateFunctionEntry() {
   as_->push(x86::rbp);
   as_->mov(x86::rbp, x86::rsp);
 #elif defined(CINDER_AARCH64)
-  as_->stp(arch::fp, arch::lr, a64::ptr_pre(a64::sp, -32));
+  as_->stp(arch::fp, arch::lr, a64::ptr_pre(a64::sp, -arch::kFrameRecordSize));
   as_->mov(arch::fp, a64::sp);
 #else
   CINDER_UNSUPPORTED
@@ -1289,7 +1289,7 @@ void NativeGenerator::generateFunctionExit() {
   as_->ret();
 #elif defined(CINDER_AARCH64)
   as_->mov(a64::sp, arch::fp);
-  as_->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 32));
+  as_->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
   as_->ret(arch::lr);
 #else
   CINDER_UNSUPPORTED
@@ -2496,9 +2496,7 @@ void NativeGenerator::generateStaticEntryPoint(
       // into the generator object when we link the frame. We need to
       // capture the incoming arguments first, which will mean we'll
       // need to save and restore the register.
-      // Extra args are above the 32-byte frame (saved fp, lr, savedReturnIP,
-      // padding), so the offset is 32.
-      as_->add(a64::x10, arch::fp, 32);
+      as_->add(a64::x10, arch::fp, arch::kFrameRecordSize);
       save_regs.emplace_back(a64::x10, a64::x10);
       need_extra_args_load = false;
     }
@@ -2514,9 +2512,7 @@ void NativeGenerator::generateStaticEntryPoint(
       a64::x(INITIAL_FUNC_REG.loc), a64::x(INITIAL_TSTATE_REG.loc), save_regs);
 
   if (need_extra_args_load) {
-    // Extra args are above the 32-byte frame (saved fp, lr, savedReturnIP,
-    // padding), so the offset is 32.
-    as_->add(a64::x10, arch::fp, 32);
+    as_->add(a64::x10, arch::fp, arch::kFrameRecordSize);
   }
   as_->b(finish_frame_setup);
   env_.addAnnotation("StaticLinkFrame", static_link_cursor);
@@ -2633,7 +2629,8 @@ void NativeGenerator::generateCode(CodeHolder& codeholder) {
 
     // leave + ret equivalent on aarch64
     as_->mov(a64::sp, arch::fp);
-    as_->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 32));
+    as_->ldp(
+        arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
     as_->ret(arch::lr);
 #else
     CINDER_UNSUPPORTED
@@ -2915,7 +2912,8 @@ NativeGenerator::generateBoxedReturnWrapper() {
     as_->bind(error);
     as_->mov(a64::x0, 0);
     as_->mov(a64::sp, arch::fp);
-    as_->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 32));
+    as_->ldp(
+        arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
     as_->ret(arch::lr);
   }
 #else
@@ -3022,7 +3020,8 @@ void NativeGenerator::generateArgcountCheckPrologue(Label correct_arg_count) {
     }
     as_->blr(arch::reg_scratch_br);
     as_->mov(a64::sp, arch::fp);
-    as_->ldp(arch::fp, arch::lr, a64::ptr_post(a64::sp, 32));
+    as_->ldp(
+        arch::fp, arch::lr, a64::ptr_post(a64::sp, arch::kFrameRecordSize));
     as_->ret(arch::lr);
     env_.addAnnotation(
         "Check if called with correct argcount", arg_check_cursor);

cinderx/Jit/codegen/gen_asm_utils.cpp

@@ -44,13 +44,7 @@ void emitCall(Environ& env, uint64_t func, const jit::lir::Instruction* instr) {
   // https://github.com/asmjit/asmjit/issues/499, but as of writing is not yet
   // available.
   env.as->mov(arch::reg_scratch_br, func);
-  // Save the return address at [fp + 16] so that getIP() can find it at a
-  // fixed offset from the frame base for cross-thread frame inspection.
-  asmjit::Label after_call = env.as->newLabel();
-  env.as->adr(arch::reg_scratch_0, after_call);
-  env.as->str(arch::reg_scratch_0, asmjit::a64::ptr(arch::fp, 16));
   env.as->blr(arch::reg_scratch_br);
-  env.as->bind(after_call);
 #else
   CINDER_UNSUPPORTED
 #endif

cinderx/Jit/compiled_function.h

@@ -18,7 +18,7 @@ extern "C" {
 // jmp correct-bound-args
 #define JITRT_CALL_REENTRY_OFFSET (-6)
 #elif defined(__aarch64__)
-// stp fp, lr, [sp, #-32]!
+// stp fp, lr, [sp, #-16]!
 // mov fp, sp
 // b correct-bound-args
 #define JITRT_CALL_REENTRY_OFFSET (-12)

cinderx/Jit/frame.cpp

@@ -11,7 +11,9 @@
 #include "cinderx/Common/util.h"
 #include "cinderx/Jit/frame_header.h"
 #include "cinderx/Jit/gen_data_footer.h"
+#if defined(CINDER_X86_64)
 #include "cinderx/Jit/symbolizer.h"
+#endif
 #include "cinderx/UpstreamBorrow/borrowed.h"
 #include "cinderx/module_state.h"
 
@@ -83,6 +85,7 @@ bool isInlined(_PyInterpreterFrame* frame) {
   return hasRtfsFunction(frame);
 }
 
+#if defined(CINDER_X86_64)
 // Return the base of the stack frame given its frame.
 uintptr_t getFrameBaseFromOnStackFrame(_PyInterpreterFrame* frame) {
   // The frame is embedded in the frame header at the beginning of the
@@ -100,17 +103,9 @@ uintptr_t getIP(_PyInterpreterFrame* frame, int frame_size) {
     auto footer = jitGenDataFooter(gen);
     if (footer->yieldPoint == nullptr) {
       // The generator is running.
-#if defined(__x86_64__)
       // On x86, we read the return address from a fixed offset on the real
       // stack relative to the resume function's RBP.
       frame_base = footer->originalFramePointer;
-#elif defined(__aarch64__)
-      // On ARM64, the JIT stores the return address at [fp + 16] before
-      // each call. During generator execution, FP is set to the
-      // GenDataFooter pointer (gi_jit_data), so [footer + 16] is the
-      // savedReturnIP field.
-      frame_base = reinterpret_cast<uintptr_t>(footer);
-#endif
     } else {
       // The generator is suspended.
       return footer->yieldPoint->resumeTarget();
@@ -119,31 +114,18 @@ uintptr_t getIP(_PyInterpreterFrame* frame, int frame_size) {
     frame_base = getFrameBaseFromOnStackFrame(frame);
   }
   // Read the saved IP from the stack.
-#if defined(__x86_64__)
   // On x86, `call` pushes the return address on the stack at a fixed
   // location relative to the caller's frame pointer.
   uintptr_t ip;
   auto saved_ip =
       reinterpret_cast<uintptr_t*>(frame_base - frame_size - kPointerSize);
   memcpy(&ip, saved_ip, kPointerSize);
   return ip;
-#elif defined(__aarch64__)
-  // On ARM64, `blr` stores the return address in lr (x30) rather than
-  // pushing it on the stack. The JIT explicitly saves the return address
-  // at [fp + 16] before each call, so we can read it from a fixed offset
-  // from the frame base.
-  uintptr_t ip;
-  auto saved_ip = reinterpret_cast<uintptr_t*>(frame_base + 2 * kPointerSize);
-  memcpy(&ip, saved_ip, kPointerSize);
-  return ip;
-#else
-  // Unsupported architecture.
-  JIT_ABORT("getIP: unsupported architecture");
-#endif
 #else
   throw std::runtime_error{"getIP: Lightweight frames are not supported"};
 #endif
 }
+#endif
 
 // Collect all the frames in the unit, with the frame for the
 // non-inlined function as the first element in the return vector.
@@ -200,9 +182,14 @@ UnitState getUnitState(_PyInterpreterFrame* frame) {
   _PyInterpreterFrame* non_inlined_sf = unit_frames[0];
   CodeRuntime* code_rt = getCodeRuntime(non_inlined_sf);
   JIT_CHECK(code_rt != nullptr, "failed to find code runtime");
-  uintptr_t ip = getIP(non_inlined_sf, code_rt->frameSize());
+
+#if defined(CINDER_AARCH64)
+  // On ARM64, look up bytecode offsets using the deopt index stored in the
+  // frame header. The JIT updates this index before each instruction that
+  // can deopt, so it always reflects the current position in the bytecode.
+  std::size_t deopt_idx = jitFrameGetHeader(non_inlined_sf)->deopt_idx;
   std::optional<UnitCallStack> locs =
-      code_rt->debugInfo()->getUnitCallStack(ip);
+      code_rt->getUnitCallStackFromDeoptIdx(deopt_idx);
   if (locs.has_value()) {
     // We may have a different number of unit_frames than locs, this happens
     // when we're updating the outer frame while we're in an inlined function,
@@ -216,11 +203,35 @@ UnitState getUnitState(_PyInterpreterFrame* frame) {
       unit_state.emplace_back(unit_frames[i], locs->at(i));
     }
   } else {
-    // We might not have debug info for a number of reasons (e.g. we've read
-    // the return address incorrectly or there's a bug with how we're
-    // generating the information). The consequences of getting this wrong
-    // (incorrect line numbers) don't warrant aborting in production, but it is
-    // worth investigating. Leave some breadcrumbs to help with debugging.
+    // We might not have debug info for a number of reasons.
+    // The consequences of getting this wrong (incorrect line numbers) don't
+    // warrant aborting in production, but it is worth investigating.
+    JIT_LOG(
+        "No debug info for deopt_idx {} in {}",
+        deopt_idx,
+        PyUnicode_AsUTF8(code_rt->frameState()->func()->func_qualname));
+    logUnitFrames();
+    JIT_DABORT("No debug info for deopt_idx {}", deopt_idx);
+    for (_PyInterpreterFrame* unit_frame : unit_frames) {
+      unit_state.emplace_back(
+          unit_frame, CodeObjLoc{_PyFrame_GetCode(unit_frame), BCOffset{-1}});
+    }
+  }
+#elif defined(CINDER_X86_64)
+  // On x86-64, look up bytecode offsets using the IP-based symbolizer.
+  uintptr_t ip = getIP(non_inlined_sf, code_rt->frameSize());
+  std::optional<UnitCallStack> locs =
+      code_rt->debugInfo()->getUnitCallStack(ip);
+  if (locs.has_value()) {
+    for (std::size_t i = 0; i < unit_frames.size(); i++) {
+      JIT_DCHECK(
+          _PyFrame_GetCode(unit_frames[i]) == locs->at(i).code,
+          "code mismatch {} vs {}",
+          codeName(_PyFrame_GetCode(unit_frames[i])),
+          codeName(locs->at(i).code));
+      unit_state.emplace_back(unit_frames[i], locs->at(i));
+    }
+  } else {
     JIT_LOG(
         "No debug info for addr {:x} {}",
         ip,
@@ -232,6 +243,9 @@ UnitState getUnitState(_PyInterpreterFrame* frame) {
           unit_frame, CodeObjLoc{_PyFrame_GetCode(unit_frame), BCOffset{-1}});
     }
   }
+#else
+  CINDER_UNSUPPORTED
+#endif
 
   return unit_state;
 }
@@ -526,6 +540,9 @@ void jitFrameInitLightweight(
   setFrameCode(frame, reifier);
   setFrameFunction(frame, (PyObject*)Py_NewRef(func));
   jitFrameGetHeader(frame)->rtfs = 0;
+#if defined(CINDER_AARCH64)
+  jitFrameGetHeader(frame)->deopt_idx = 0;
+#endif
 #else
   frame->stacktop = 0;
   setFrameInstruction(frame, _PyCode_CODE(code) - 1);