Fix getIP on aarch64: SP-relative saved-IP slot [SP,#8]

Save JIT return address to [SP, #8] before each BL/BLR instruction,
and read it back in getIP() at frame_base - frame_size - kPointerSize.
Falls back to saved LR at [FP+8] when slot reads zero (no call yet).

Root causes fixed:
1. Double-counting: old FP-relative STR offset added kPointerSize on top
   of stack_frame_size which already includes kStackAlign, writing below SP.
2. Register clobbering: translateCall isReg case could have call target in
   x12/x13 which saveReturnAddress overwrote. Now moves target to x16 first.
3. ptr_resolve scratch: large frame offsets triggered ptr_resolve paths
   using x13 as scratch, clobbering the call target. SP-relative [SP,#8]
   uses simple immediate addressing with no scratch register.

Validated: 15/15 CPU modules pass (default + PYTHONJITLISTALL=1),
8/8 GPU modules pass with zero regressions, GPU JIT smoke test pass
(add/matmul/relu on NVIDIA GB200).

Author: SonicField

cinderx/Jit/codegen/autogen.cpp

@@ -1937,20 +1937,35 @@ 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()) {
-    as->blr(AT::getGp(input));
+    auto target = AT::getGp(input);
+    if (target.id() != arch::reg_scratch_br.id()) {
+      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->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->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->blr(arch::reg_scratch_br);
   } else {
     JIT_ABORT("Unsupported operand type for Call: {}", input->type());
   }
+  as->bind(after_call);
 
   if (instr->origin()) {
     asmjit::Label label = as->newLabel();

cinderx/Jit/codegen/gen_asm.cpp

@@ -1457,6 +1457,9 @@ 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

cinderx/Jit/codegen/gen_asm_utils.cpp

@@ -26,16 +26,13 @@ void emitCall(
 #if defined(CINDER_X86_64)
   env.as->call(label);
 #elif defined(CINDER_AARCH64)
-  // 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.
+  // 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.
   {
     asmjit::Label after_call = env.as->newLabel();
-    int offset = -(env.stack_frame_size - kPointerSize);
     env.as->adr(arch::reg_scratch_0, after_call);
-    env.as->str(
-        arch::reg_scratch_0,
-        arch::ptr_resolve(env.as, arch::fp, offset, arch::reg_scratch_1));
+    env.as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
     env.as->bl(label);
     env.as->bind(after_call);
   }
@@ -54,14 +51,11 @@ 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 stack before blr, matching x86 call semantics.
+  // Save return address to [SP, #8] before BLR.
   {
     asmjit::Label after_call = env.as->newLabel();
-    int offset = -(env.stack_frame_size - kPointerSize);
     env.as->adr(arch::reg_scratch_0, after_call);
-    env.as->str(
-        arch::reg_scratch_0,
-        arch::ptr_resolve(env.as, arch::fp, offset, arch::reg_scratch_1));
+    env.as->str(arch::reg_scratch_0, asmjit::arm::Mem(asmjit::a64::sp, 8));
     env.as->blr(arch::reg_scratch_br);
     env.as->bind(after_call);
   }

cinderx/Jit/frame.cpp

@@ -118,20 +118,21 @@ uintptr_t getIP(_PyInterpreterFrame* frame, int frame_size) {
   memcpy(&ip, saved_ip, kPointerSize);
   return ip;
 #elif defined(__aarch64__)
-  // On aarch64, the JIT prologue does:
-  //   stp x29, x30, [sp, #-16]!  ; saves FP and LR
-  //   mov x29, sp                ; sets up frame pointer
-  // In lightweight frame mode, frame_base == x29 (verified from
-  // frameHeaderSize() computation). The saved LR (return address)
-  // is at [frame_base + 8].
-  //
-  // We read it directly rather than walking the FP chain, because
-  // intermediate C frames may be compiled without frame pointers
-  // (e.g. pytest/pluggy extensions), breaking the chain walk.
+  // 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].
   uintptr_t ip;
   auto saved_ip =
-      reinterpret_cast<uintptr_t*>(frame_base + kPointerSize);
+      reinterpret_cast<uintptr_t*>(frame_base - frame_size - kPointerSize);
   memcpy(&ip, saved_ip, kPointerSize);
+  if (ip == 0) {
+    auto saved_lr =
+        reinterpret_cast<uintptr_t*>(frame_base + kPointerSize);
+    memcpy(&ip, saved_lr, kPointerSize);
+  }
   return ip;
 #else
   // Unsupported architecture.
@@ -444,6 +445,7 @@ _PyInterpreterFrame* convertInterpreterFrameFromStackToSlab(
   }
 
   jitFramePopulateFrame(frame);
+  updatePrevInstr(frame);
   jitFrameRemoveReifier(frame);
 
   memcpy(new_frame, frame, code->co_framesize * sizeof(PyObject*));