Bring a few jithelpers to new unwind plan

src/coreclr/pal/inc/unixasmmacrosarm64.inc

@@ -468,6 +468,29 @@ C_FUNC(\Name\()_End):
         add     \target, sp, #192
 .endm
 
+// Epilog for PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS - restores FP callee-saved and returns
+.macro POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN
+
+        // Restore FP callee-saved registers (d8-d15) from sp+0
+        ldp     d8,  d9,  [sp, #0]
+        ldp     d10, d11, [sp, #16]
+        ldp     d12, d13, [sp, #32]
+        ldp     d14, d15, [sp, #48]
+
+        // Deallocate space for FloatArgumentRegisters + FP callee-saved
+        EPILOG_STACK_FREE 192
+
+        // Restore callee-saved registers
+        EPILOG_RESTORE_REG_PAIR   x27, x28, 80
+        EPILOG_RESTORE_REG_PAIR   x25, x26, 64
+        EPILOG_RESTORE_REG_PAIR   x23, x24, 48
+        EPILOG_RESTORE_REG_PAIR   x21, x22, 32
+        EPILOG_RESTORE_REG_PAIR   x19, x20, 16
+
+        EPILOG_RESTORE_REG_PAIR_INDEXED   fp, lr, 176
+        ret
+.endm
+
 // ------------------------------------------------------------------
 // Macro to generate Redirection Stubs
 //

src/coreclr/pal/inc/unixasmmacrosloongarch64.inc

@@ -443,6 +443,31 @@ C_FUNC(\Name\()_End):
         addi.d  \target, $sp, 128
 .endm
 
+// ------------------------------------------------------------------
+// Epilog macro for PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS that restores all registers and returns.
+// Use this for JIT helpers that may return normally (unlike IL_Throw which never returns).
+.macro POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN
+        // Restore FP callee-saved registers (f24-f31) from offset 0
+        fld.d   $f24, $sp, 0
+        fld.d   $f25, $sp, 8
+        fld.d   $f26, $sp, 16
+        fld.d   $f27, $sp, 24
+        fld.d   $f28, $sp, 32
+        fld.d   $f29, $sp, 40
+        fld.d   $f30, $sp, 48
+        fld.d   $f31, $sp, 56
+
+        // Restore callee-saved registers from offset 128
+        RESTORE_CALLEESAVED_REGISTERS $sp, 128
+
+        // Restore fp, ra
+        EPILOG_RESTORE_REG_PAIR 22, 1, 128
+
+        // Deallocate stack and return
+        EPILOG_STACK_FREE 288
+        EPILOG_RETURN
+.endm
+
 // ------------------------------------------------------------------
 // Macro to generate Redirection Stubs
 //

src/coreclr/pal/inc/unixasmmacrosriscv64.inc

@@ -394,6 +394,35 @@ C_FUNC(\Name):
     addi  \target, sp, 160
 .endm
 
+// ------------------------------------------------------------------
+// Epilog macro for PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS that restores all registers and returns.
+// Use this for JIT helpers that may return normally (unlike IL_Throw which never returns).
+.macro POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN
+    // Restore FP callee-saved registers (fs0-fs11 = f8,f9,f18-f27) from offset 0
+    fld   fs0,  0(sp)   // f8
+    fld   fs1,  8(sp)   // f9
+    fld   fs2,  16(sp)  // f18
+    fld   fs3,  24(sp)  // f19
+    fld   fs4,  32(sp)  // f20
+    fld   fs5,  40(sp)  // f21
+    fld   fs6,  48(sp)  // f22
+    fld   fs7,  56(sp)  // f23
+    fld   fs8,  64(sp)  // f24
+    fld   fs9,  72(sp)  // f25
+    fld   fs10, 80(sp)  // f26
+    fld   fs11, 88(sp)  // f27
+
+    // Restore callee-saved registers from offset 160
+    RESTORE_CALLEESAVED_REGISTERS sp, 160
+
+    // Restore fp, ra
+    EPILOG_RESTORE_REG_PAIR fp, ra, 160
+
+    // Deallocate stack and return
+    EPILOG_STACK_FREE 352
+    EPILOG_RETURN
+.endm
+
 // ------------------------------------------------------------------
 // Macro to generate Redirection Stubs
 //

src/coreclr/vm/amd64/AsmHelpers.asm

@@ -402,12 +402,15 @@ NESTED_END OnCallCountThresholdReachedStub, _TEXT
 extern JIT_PatchpointWorkerWorkerWithPolicy:proc
 
 NESTED_ENTRY JIT_Patchpoint, _TEXT
-        PROLOG_WITH_TRANSITION_BLOCK
+        ; Use PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS to save all registers including FP callee-saved
+        ; This allows us to build a complete CONTEXT from TransitionBlock without RtlCaptureContext
+        PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS rcx
 
-        lea     rcx, [rsp + __PWTB_TransitionBlock] ; TransitionBlock *
+        ; RCX contains pointer to TransitionBlock
         call    JIT_PatchpointWorkerWorkerWithPolicy
 
-        EPILOG_WITH_TRANSITION_BLOCK_RETURN
+        ; If we return, restore all registers and return to caller
+        POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN
 NESTED_END JIT_Patchpoint, _TEXT
 
 ; first arg register holds iloffset, which needs to be moved to the second register, and the first register filled with NULL

src/coreclr/vm/amd64/AsmMacros.inc

@@ -503,13 +503,14 @@ PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS macro target
 
         PUSH_CALLEE_SAVED_REGISTERS
 
-        ; Allocate space for: shadow for call (32) + FP callee-saved (160) + float args (64) + arg regs (32) + padding (8) = 296 bytes
-        ; Shadow space at offset 0 is reserved for the call to IL_Throw_Impl etc.
-        ; This makes RSP 16-byte aligned (8 + 64 + 296 = 368, and original RSP - 368 is 16-byte aligned)
-        alloc_stack     296
+        ; Allocate space for: shadow (32) + FP callee-saved (160) + float args (64) + padding (8) = 264 bytes
+        ; This makes RSP 16-byte aligned (8 + 64 + 264 = 336, and original RSP - 336 is 16-byte aligned)
+        alloc_stack     264
 
-        ; Save argument registers at offset 256 (32 + 160 + 64)
-        SAVE_ARGUMENT_REGISTERS 256
+        ; Save argument registers into caller's shadow space
+        ; TransitionBlock is at rsp + 264, sizeof(TransitionBlock) = 72
+        ; So argument registers go at rsp + 264 + 72 = rsp + 336
+        SAVE_ARGUMENT_REGISTERS 336
 
         ; Save float argument registers at offset 192 (32 + 160)
         SAVE_FLOAT_ARGUMENT_REGISTERS 192
@@ -530,8 +531,31 @@ PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS macro target
 
         END_PROLOGUE
 
-        ; TransitionBlock pointer points to CalleeSavedRegisters at rsp + 296
-        lea     target, [rsp + 296]
+        ; TransitionBlock pointer points to CalleeSavedRegisters at rsp + 264
+        lea     target, [rsp + 264]
+
+        endm
+
+; Epilog for PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS - restores XMM callee-saved and returns
+POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN macro
+
+        ; Restore FP callee-saved registers (xmm6-xmm15) from offset 32 (after shadow space)
+        movdqa  xmm6,  [rsp + 20h]
+        movdqa  xmm7,  [rsp + 30h]
+        movdqa  xmm8,  [rsp + 40h]
+        movdqa  xmm9,  [rsp + 50h]
+        movdqa  xmm10, [rsp + 60h]
+        movdqa  xmm11, [rsp + 70h]
+        movdqa  xmm12, [rsp + 80h]
+        movdqa  xmm13, [rsp + 90h]
+        movdqa  xmm14, [rsp + 0A0h]
+        movdqa  xmm15, [rsp + 0B0h]
+
+        ; Deallocate stack space (264 bytes)
+        add     rsp, 264
+
+        POP_CALLEE_SAVED_REGISTERS
+        ret
 
         endm
 

src/coreclr/vm/arm64/asmconstants.h

@@ -101,12 +101,38 @@ ASMCONSTANTS_C_ASSERT(SIZEOF__Frame == sizeof(Frame));
 #endif
 ASMCONSTANTS_C_ASSERT(SIZEOF__CONTEXT == sizeof(T_CONTEXT));
 
+#define OFFSETOF__CONTEXT__ContextFlags 0x0
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__ContextFlags == offsetof(T_CONTEXT, ContextFlags));
+
+// CONTEXT_INTEGER_BIT is bit 1 in ContextFlags (from pal.h CONTEXT_INTEGER definition)
+#define CONTEXT_INTEGER_BIT           1
+
+#define OFFSETOF__CONTEXT__X0         0x08
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__X0 == offsetof(T_CONTEXT, X0));
+
 #define OFFSETOF__CONTEXT__X19        0xA0
 ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__X19 == offsetof(T_CONTEXT, X19));
 
 #define OFFSETOF__CONTEXT__Fp         0xF0
 ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__Fp == offsetof(T_CONTEXT, Fp));
 
+#define OFFSETOF__CONTEXT__Lr         0xF8
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__Lr == offsetof(T_CONTEXT, Lr));
+
+#define OFFSETOF__CONTEXT__Sp         0x100
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__Sp == offsetof(T_CONTEXT, Sp));
+
+#define OFFSETOF__CONTEXT__Pc         0x108
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__Pc == offsetof(T_CONTEXT, Pc));
+
+// Floating point registers V[0..31] start at 0x110, each is 16 bytes (NEON128)
+// Non-volatile FP registers are V8-V15 (d8-d15)
+#define OFFSETOF__CONTEXT__V0         0x110
+ASMCONSTANTS_C_ASSERT(OFFSETOF__CONTEXT__V0 == offsetof(T_CONTEXT, V));
+
+// CONTEXT_FLOATING_POINT_BIT is bit 2 in ContextFlags
+#define CONTEXT_FLOATING_POINT_BIT    2
+
 #define               OFFSETOF__DynamicHelperStubArgs__Constant1    0x0
 ASMCONSTANTS_C_ASSERT(OFFSETOF__DynamicHelperStubArgs__Constant1
                     == offsetof(DynamicHelperStubArgs, Constant1));

src/coreclr/vm/arm64/asmhelpers.S

@@ -586,9 +586,11 @@ NESTED_END OnCallCountThresholdReachedStub, _TEXT
 NESTED_ENTRY JIT_Patchpoint, _TEXT, NoHandler
     PROLOG_WITH_TRANSITION_BLOCK
 
-    add     x0, sp, #__PWTB_TransitionBlock // TransitionBlock *
+    // x0 = pointer to TransitionBlock
+    add     x0, sp, #__PWTB_TransitionBlock
     bl      C_FUNC(JIT_PatchpointWorkerWorkerWithPolicy)
 
+    // If we return, restore all registers and return to caller
     EPILOG_WITH_TRANSITION_BLOCK_RETURN
 NESTED_END JIT_Patchpoint, _TEXT
 
@@ -2805,3 +2807,62 @@ NESTED_ENTRY IL_Rethrow, _TEXT, NoHandler
     brk     #0
 NESTED_END IL_Rethrow, _TEXT
 
+// ------------------------------------------------------------------
+// ClrRestoreNonvolatileContextWorker
+//
+// Restores registers based on ContextFlags and jumps to the target IP.
+// When CONTEXT_INTEGER is set, restores ALL integer registers (x0-x28)
+// because exception handling needs x0 (exception object).
+// When CONTEXT_FLOATING_POINT is set, restores non-volatile FP regs (d8-d15).
+//
+// Arguments:
+//   x0 - pointer to CONTEXT structure
+//
+// Does not return - jumps directly to the Pc stored in the CONTEXT
+// ------------------------------------------------------------------
+LEAF_ENTRY ClrRestoreNonvolatileContextWorker, _TEXT
+
+    // Save context pointer in x16 since we'll overwrite x0
+    mov     x16, x0
+
+    // Check ContextFlags to see if we should restore floating point registers
+    ldr     w17, [x16, #OFFSETOF__CONTEXT__ContextFlags]
+    tbz     w17, #CONTEXT_FLOATING_POINT_BIT, LOCAL_LABEL(SkipFloatingPointRestore)
+
+    // Restore non-volatile FP registers d8-d15 (lower 64 bits of v8-v15)
+    // V8 is at OFFSETOF__CONTEXT__V0 + 8*16 = 0x110 + 0x80 = 0x190
+    ldp     q8, q9,   [x16, #(OFFSETOF__CONTEXT__V0 + 8*16)]
+    ldp     q10, q11, [x16, #(OFFSETOF__CONTEXT__V0 + 10*16)]
+    ldp     q12, q13, [x16, #(OFFSETOF__CONTEXT__V0 + 12*16)]
+    ldp     q14, q15, [x16, #(OFFSETOF__CONTEXT__V0 + 14*16)]
+
+LOCAL_LABEL(SkipFloatingPointRestore):
+    // Check ContextFlags to see if we should restore integer registers
+    tbz     w17, #CONTEXT_INTEGER_BIT, LOCAL_LABEL(SkipIntegerRestore)
+
+    // Restore argument registers x0-x7 (exception handling needs x0 for exception object)
+    // and non-volatile registers x19-x28
+    ldp     x0, x1,   [x16, #OFFSETOF__CONTEXT__X0]
+    ldp     x2, x3,   [x16, #(OFFSETOF__CONTEXT__X0 + 16)]
+    ldp     x4, x5,   [x16, #(OFFSETOF__CONTEXT__X0 + 32)]
+    ldp     x6, x7,   [x16, #(OFFSETOF__CONTEXT__X0 + 48)]
+    // Skip x8-x18: x8-x15 are scratch, x16-x17 we're using, x18 is platform reserved
+    ldp     x19, x20, [x16, #OFFSETOF__CONTEXT__X19]
+    ldp     x21, x22, [x16, #(OFFSETOF__CONTEXT__X19 + 16)]
+    ldp     x23, x24, [x16, #(OFFSETOF__CONTEXT__X19 + 32)]
+    ldp     x25, x26, [x16, #(OFFSETOF__CONTEXT__X19 + 48)]
+    ldp     x27, x28, [x16, #(OFFSETOF__CONTEXT__X19 + 64)]
+
+LOCAL_LABEL(SkipIntegerRestore):
+    // Restore fp (x29) and lr (x30)
+    ldp     fp, lr, [x16, #OFFSETOF__CONTEXT__Fp]
+
+    // Load Sp and Pc into scratch registers (after all other loads)
+    ldr     x17, [x16, #OFFSETOF__CONTEXT__Sp]
+    ldr     x16, [x16, #OFFSETOF__CONTEXT__Pc]
+
+    // Set sp and jump to target
+    mov     sp, x17
+    br      x16
+
+LEAF_END ClrRestoreNonvolatileContextWorker, _TEXT

src/coreclr/vm/arm64/asmhelpers.asm

@@ -1013,9 +1013,11 @@ __HelperNakedFuncName SETS "$helper":CC:"Naked"
     NESTED_ENTRY JIT_Patchpoint
         PROLOG_WITH_TRANSITION_BLOCK
 
-        add     x0, sp, #__PWTB_TransitionBlock ; TransitionBlock *
+        ; x0 = pointer to TransitionBlock
+        add     x0, sp, #__PWTB_TransitionBlock
         bl      JIT_PatchpointWorkerWorkerWithPolicy
 
+        ; If we return, restore all registers and return to caller
         EPILOG_WITH_TRANSITION_BLOCK_RETURN
     NESTED_END
 
@@ -3058,5 +3060,63 @@ CopyLoop
         brk     #0
     NESTED_END IL_Rethrow
 
+; ------------------------------------------------------------------
+; ClrRestoreNonvolatileContextWorker
+;
+; Restores registers based on ContextFlags and jumps to PC.
+; When CONTEXT_INTEGER is set, restores ALL integer registers (x0-x28)
+; because exception handling needs x0 (exception object).
+; When CONTEXT_FLOATING_POINT is set, restores non-volatile FP regs (d8-d15).
+;
+; x0 - pointer to CONTEXT structure
+; x1 - unused (SSP, not used on ARM64)
+; ------------------------------------------------------------------
+    LEAF_ENTRY ClrRestoreNonvolatileContextWorker
+
+        ; Save CONTEXT pointer in x16 before we potentially clobber x0
+        mov     x16, x0
+
+        ; Check if CONTEXT_FLOATING_POINT bit is set (bit 2)
+        ldr     w17, [x16, #OFFSETOF__CONTEXT__ContextFlags]
+        tbz     w17, #2, SkipFloatingPointRestore
+
+        ; Restore non-volatile FP registers d8-d15 (full q8-q15)
+        ; V8 is at OFFSETOF__CONTEXT__V0 + 8*16 = 0x110 + 0x80 = 0x190
+        ldp     q8, q9,   [x16, #(OFFSETOF__CONTEXT__V0 + 128)]
+        ldp     q10, q11, [x16, #(OFFSETOF__CONTEXT__V0 + 160)]
+        ldp     q12, q13, [x16, #(OFFSETOF__CONTEXT__V0 + 192)]
+        ldp     q14, q15, [x16, #(OFFSETOF__CONTEXT__V0 + 224)]
+
+SkipFloatingPointRestore
+        ; Check if CONTEXT_INTEGER bit is set
+        tbz     w17, #1, SkipIntegerRestore  ; CONTEXT_INTEGER_BIT = 1
+
+        ; Restore argument registers x0-x7 (exception handling needs x0 for exception object)
+        ; and non-volatile registers x19-x28
+        ldp     x0, x1,   [x16, #OFFSETOF__CONTEXT__X0]
+        ldp     x2, x3,   [x16, #(OFFSETOF__CONTEXT__X0 + 16)]
+        ldp     x4, x5,   [x16, #(OFFSETOF__CONTEXT__X0 + 32)]
+        ldp     x6, x7,   [x16, #(OFFSETOF__CONTEXT__X0 + 48)]
+        ; Skip x8-x18: x8-x15 are scratch, x16-x17 we're using, x18 is platform reserved
+        ldp     x19, x20, [x16, #OFFSETOF__CONTEXT__X19]
+        ldp     x21, x22, [x16, #(OFFSETOF__CONTEXT__X19 + 16)]
+        ldp     x23, x24, [x16, #(OFFSETOF__CONTEXT__X19 + 32)]
+        ldp     x25, x26, [x16, #(OFFSETOF__CONTEXT__X19 + 48)]
+        ldp     x27, x28, [x16, #(OFFSETOF__CONTEXT__X19 + 64)]
+
+SkipIntegerRestore
+        ; Restore fp and lr
+        ldp     fp, lr, [x16, #OFFSETOF__CONTEXT__Fp]
+
+        ; Load Sp and Pc (x16 will be overwritten)
+        ldr     x17, [x16, #OFFSETOF__CONTEXT__Pc]
+        ldr     x16, [x16, #OFFSETOF__CONTEXT__Sp]
+
+        ; Set sp and jump
+        mov     sp, x16
+        br      x17
+
+    LEAF_END ClrRestoreNonvolatileContextWorker
+
 ; Must be at very end of file
     END

src/coreclr/vm/arm64/asmmacros.h

@@ -253,6 +253,29 @@ OFFSETOF__ee_alloc_context  EQU OFFSETOF__RuntimeThreadLocals__ee_alloc_context
         add     $Target, sp, #192
     MEND
 
+; Epilog for PUSH_COOP_PINVOKE_FRAME_WITH_FLOATS - restores FP callee-saved and returns
+    MACRO
+    POP_COOP_PINVOKE_FRAME_WITH_FLOATS_RETURN
+
+        ; Restore FP callee-saved registers (d8-d15) from sp+0
+        ldp     d8,  d9,  [sp, #0]
+        ldp     d10, d11, [sp, #16]
+        ldp     d12, d13, [sp, #32]
+        ldp     d14, d15, [sp, #48]
+
+        ; Deallocate space for FloatArgumentRegisters + FP callee-saved
+        EPILOG_STACK_FREE 192
+
+        ; Restore callee-saved registers
+        EPILOG_RESTORE_REG_PAIR   x27, x28, #80
+        EPILOG_RESTORE_REG_PAIR   x25, x26, #64
+        EPILOG_RESTORE_REG_PAIR   x23, x24, #48
+        EPILOG_RESTORE_REG_PAIR   x21, x22, #32
+        EPILOG_RESTORE_REG_PAIR   x19, x20, #16
+
+        EPILOG_RESTORE_REG_PAIR_RET   fp, lr, #176!
+    MEND
+
 #define GC_ALLOC_FINALIZE 1
 
 ;-----------------------------------------------------------------------------