Phase 4b: FP-relative saved-IP slot (eliminates kVectorcallArgsOffset overlap)

Store the saved-IP at [FP + saved_ip_fp_offset] instead of [SP, #8] to
eliminate the theoretical overlap between the saved-IP slot and
vectorcall args[0] at kVectorcallArgsOffset (also [SP+8]).

Key change: pre-compute saved_ip_fp_offset in getVectorcallEntry after
register allocation but before generateCode, because generateAssemblyBody
(which emits ADR+STR for saved-IP) runs before computeFrameInfo. A
JIT_CHECK in computeFrameInfo verifies the pre-computed value matches.

Files changed:
  environ.h        - add saved_ip_fp_offset field to Environ
  gen_asm.cpp      - early offset computation + JIT_CHECK assertion
  gen_asm_utils.cpp - emitCall uses FP-relative ptr_resolve
  autogen.cpp      - translateCall uses FP-relative ptr_resolve
  code_runtime.h/cpp - savedIpFpOffset accessor for CodeRuntime
  frame.cpp        - getIP reads from FP + savedIpFpOffset
  frame_shadow.cpp - shadow frame getIP uses same offset

Test results (15/15 CPU, PYTHONJITCOMPILATIONTHRESHOLD=1000):
  11 PASS, 4 FAIL (all pre-existing, verified with CINDERJIT_DISABLE=1)

Author: SonicField

cinderx/Jit/code_runtime.cpp

@@ -116,6 +116,16 @@ void CodeRuntime::setFrameSize(int size) {
   frame_size_ = size;
 }
 
+#if defined(__aarch64__)
+int CodeRuntime::savedIpFpOffset() const {
+  return saved_ip_fp_offset_;
+}
+
+void CodeRuntime::setSavedIpFpOffset(int offset) {
+  saved_ip_fp_offset_ = offset;
+}
+#endif
+
 DebugInfo* CodeRuntime::debugInfo() {
   return &debug_info_;
 }

cinderx/Jit/code_runtime.h

@@ -127,6 +127,10 @@ class alignas(16) CodeRuntime {
   // Get and set the total size of a stack frame for this compiled code object.
   int frameSize() const;
   void setFrameSize(int size);
+#if defined(__aarch64__)
+  int savedIpFpOffset() const;
+  void setSavedIpFpOffset(int offset);
+#endif
 
   DebugInfo* debugInfo();
 
@@ -162,6 +166,9 @@ class alignas(16) CodeRuntime {
 #endif
 
   int frame_size_{-1};
+#if defined(__aarch64__)
+  int saved_ip_fp_offset_{0};
+#endif
   DebugInfo debug_info_;
 };
 

cinderx/Jit/codegen/autogen.cpp

@@ -1937,8 +1937,6 @@ void translateCall(Environ* env, const Instruction* instr) {
   auto output = instr->output();
   auto input = instr->getInput(0);
 
-  // Save return address to [SP, #8] before BLR, matching x86 call semantics.
-  // Move call target to x16 first in case input register is x12 (scratch_0).
   asmjit::Label after_call = as->newLabel();
 
   if (input->isReg()) {
@@ -1947,20 +1945,29 @@ void translateCall(Environ* env, const Instruction* instr) {
       as->mov(arch::reg_scratch_br, target);
     }
     as->adr(arch::reg_scratch_0, after_call);
-    as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
+    as->str(
+        arch::reg_scratch_0,
+        arch::ptr_resolve(
+            as, arch::fp, env->saved_ip_fp_offset, arch::reg_scratch_1));
     as->blr(arch::reg_scratch_br);
   } else if (input->isImm()) {
     as->mov(arch::reg_scratch_br, input->getConstant());
     as->adr(arch::reg_scratch_0, after_call);
-    as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
+    as->str(
+        arch::reg_scratch_0,
+        arch::ptr_resolve(
+            as, arch::fp, env->saved_ip_fp_offset, arch::reg_scratch_1));
     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));
     as->adr(arch::reg_scratch_0, after_call);
-    as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
+    as->str(
+        arch::reg_scratch_0,
+        arch::ptr_resolve(
+            as, arch::fp, env->saved_ip_fp_offset, arch::reg_scratch_1));
     as->blr(arch::reg_scratch_br);
   } else {
     JIT_ABORT("Unsupported operand type for Call: {}", input->type());

cinderx/Jit/codegen/environ.h

@@ -28,6 +28,11 @@ struct Environ {
   // functions.
   int stack_frame_size{-1};
 
+#if defined(CINDER_AARCH64)
+  // FP-relative offset of the saved-IP slot. Set by computeFrameInfo.
+  int saved_ip_fp_offset{0};
+#endif
+
   // A subset of stack_frame_size: only the shadow frames and spilled values.
   int shadow_frames_and_spill_size{0};
 

cinderx/Jit/codegen/gen_asm.cpp

@@ -1297,6 +1297,24 @@ void* NativeGenerator::getVectorcallEntry() {
         *lir_func);
   }
 
+
+#if defined(CINDER_AARCH64)
+  // Pre-compute saved_ip_fp_offset before generateCode, because
+  // emitCall/translateCall in generateAssemblyBody need this offset to
+  // emit the ADR+STR saved-IP stores. generateAssemblyBody runs before
+  // computeFrameInfo, so the offset must be computed here.
+  {
+    int hdr = std::max(env_.shadow_frames_and_spill_size, kPointerSize)
+              + kStackAlign;
+    auto saved_regs = env_.changed_regs & CALLEE_SAVE_REGS;
+    int saved_regs_sz = ((saved_regs.count() + 1) / 2) * kStackAlign;
+    if ((hdr + saved_regs_sz + env_.max_arg_buffer_size) % kStackAlign) {
+      hdr += kPointerSize;
+    }
+    env_.saved_ip_fp_offset = -(hdr - 8);
+  }
+#endif
+
   lir_func_ = std::move(lir_func);
 
   try {
@@ -1329,7 +1347,8 @@ void* NativeGenerator::getVectorcallEntry() {
   // The physical frame includes kStackAlign extra bytes for the saved-IP
   // slot, but getIP() uses frame_base - frame_size - kPointerSize which
   // must point to SP + 8 (where the ADR+STR saves the return address).
-  env_.code_rt->setFrameSize(env_.stack_frame_size - kStackAlign);
+  env_.code_rt->setFrameSize(env_.stack_frame_size);
+  env_.code_rt->setSavedIpFpOffset(env_.saved_ip_fp_offset);
 #else
   env_.code_rt->setFrameSize(env_.stack_frame_size);
 #endif
@@ -1437,6 +1456,13 @@ NativeGenerator::FrameInfo NativeGenerator::computeFrameInfo() {
   env_.last_callee_saved_reg_off =
       info.header_and_spill_size + info.saved_regs_size();
   env_.stack_frame_size = info.size();
+#if defined(CINDER_AARCH64)
+  JIT_CHECK(
+      env_.saved_ip_fp_offset == -(info.header_and_spill_size - 8),
+      "saved_ip_fp_offset mismatch: pre-computed={}, computed={}",
+      env_.saved_ip_fp_offset,
+      -(info.header_and_spill_size - 8));
+#endif
   return info;
 }
 
@@ -1457,9 +1483,6 @@ int NativeGenerator::allocateHeaderAndSpillSpace(const FrameInfo& frame_info) {
       as_->sub(a64::sp, a64::sp, arch::reg_scratch_0);
     }
   }
-  // Zero-init the saved-IP slot at [SP, #8]. getIP() falls back to
-  // [FP+8] (saved LR) when this slot reads 0.
-  as_->str(a64::xzr, asmjit::arm::Mem(a64::sp, 8));
 
   // There is a difference here from x86-64, because the aarch64 stack cannot be
   // misaligned. Here we are returning the amount of space that we have added to
@@ -1505,6 +1528,12 @@ void NativeGenerator::saveCallerRegisters(
       as_->sub(a64::sp, a64::sp, arch::reg_scratch_0);
     }
   }
+  // Zero-init the saved-IP slot at FP + saved_ip_fp_offset.
+  // getIP() falls back to [FP+8] (saved LR) when this slot reads 0.
+  as_->str(
+      a64::xzr,
+      arch::ptr_resolve(
+          as_, arch::fp, env_.saved_ip_fp_offset, arch::reg_scratch_0));
 #else
   CINDER_UNSUPPORTED
 #endif

cinderx/Jit/codegen/gen_asm_utils.cpp

@@ -26,13 +26,16 @@ void emitCall(
 #if defined(CINDER_X86_64)
   env.as->call(label);
 #elif defined(CINDER_AARCH64)
-  // Save return address to [SP, #8] before BL, matching x86 call semantics.
-  // Slot is within the kStackAlign extra bytes reserved by computeFrameInfo.
-  // Using SP-relative avoids ptr_resolve scratch register issues.
+  // Save return address to stack before bl, matching x86 call semantics.
+  // Slot at [FP - (stack_frame_size - 8)] = [SP + 8], within the extra
+  // kStackAlign bytes reserved by computeFrameInfo.
   {
     asmjit::Label after_call = env.as->newLabel();
     env.as->adr(arch::reg_scratch_0, after_call);
-    env.as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
+    env.as->str(
+        arch::reg_scratch_0,
+        arch::ptr_resolve(
+            env.as, arch::fp, env.saved_ip_fp_offset, arch::reg_scratch_1));
     env.as->bl(label);
     env.as->bind(after_call);
   }
@@ -51,11 +54,14 @@ 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 return address to [SP, #8] before BLR.
+  // Save return address to stack before blr, matching x86 call semantics.
   {
     asmjit::Label after_call = env.as->newLabel();
     env.as->adr(arch::reg_scratch_0, after_call);
-    env.as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
+    env.as->str(
+        arch::reg_scratch_0,
+        arch::ptr_resolve(
+            env.as, arch::fp, env.saved_ip_fp_offset, arch::reg_scratch_1));
     env.as->blr(arch::reg_scratch_br);
     env.as->bind(after_call);
   }

cinderx/Jit/frame.cpp

@@ -118,15 +118,13 @@ uintptr_t getIP(_PyInterpreterFrame* frame, int frame_size) {
   memcpy(&ip, saved_ip, kPointerSize);
   return ip;
 #elif defined(__aarch64__)
-  // On aarch64, the JIT saves the return address to [SP, #8] before each
-  // BL/BLR. This slot is at frame_base - frame_size - kPointerSize
-  // (= SP + 8, since SP = FP - stack_frame_size and frame_size =
-  // stack_frame_size - kStackAlign). If no call has happened yet (slot
-  // is 0 from prologue zero-init), fall back to the saved LR from the
-  // STP at [frame_base + kPointerSize].
+  // On aarch64, frame_size is actually the FP-relative offset of the
+  // saved-IP slot (savedIpFpOffset, a negative value). The slot is at
+  // frame_base + offset. If no call has happened yet (slot is 0 from
+  // prologue zero-init), fall back to saved LR at [frame_base + 8].
   uintptr_t ip;
   auto saved_ip =
-      reinterpret_cast<uintptr_t*>(frame_base - frame_size - kPointerSize);
+      reinterpret_cast<uintptr_t*>(frame_base + frame_size);
   memcpy(&ip, saved_ip, kPointerSize);
   if (ip == 0) {
     auto saved_lr =
@@ -198,7 +196,11 @@ 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");
+#if defined(__aarch64__)
+  uintptr_t ip = getIP(non_inlined_sf, code_rt->savedIpFpOffset());
+#else
   uintptr_t ip = getIP(non_inlined_sf, code_rt->frameSize());
+#endif
   std::optional<UnitCallStack> locs =
       code_rt->debugInfo()->getUnitCallStack(ip);
   if (locs.has_value()) {

cinderx/Jit/frame_shadow.cpp

@@ -130,10 +130,15 @@ uintptr_t getIP(_PyShadowFrame* shadow_frame, int frame_size) {
   } else {
     frame_base = getFrameBaseFromOnStackShadowFrame(shadow_frame);
   }
-  // Read the saved IP from the stack
+  // Read the saved IP from the stack.
   uintptr_t ip;
+#if defined(__aarch64__)
+  auto saved_ip =
+      reinterpret_cast<uintptr_t*>(frame_base + frame_size);
+#else
   auto saved_ip =
       reinterpret_cast<uintptr_t*>(frame_base - frame_size - kPointerSize);
+#endif
   memcpy(&ip, saved_ip, kPointerSize);
   return ip;
 }
@@ -381,7 +386,11 @@ UnitState getUnitState(_PyShadowFrame* shadow_frame) {
   unit_state.reserve(unit_frames.size());
   _PyShadowFrame* non_inlined_sf = unit_frames[0];
   CodeRuntime* code_rt = getCodeRuntime(non_inlined_sf);
+#if defined(__aarch64__)
+  uintptr_t ip = getIP(non_inlined_sf, code_rt->savedIpFpOffset());
+#else
   uintptr_t ip = getIP(non_inlined_sf, code_rt->frameSize());
+#endif
   std::optional<UnitCallStack> locs =
       code_rt->debugInfo()->getUnitCallStack(ip);
   if (locs.has_value()) {