Fix up hijacking on x64 platforms in presence of runtime async (#122631)

Contributes to #122492.

Author: MichalStrehovsky

src/coreclr/nativeaot/Runtime/amd64/AsmMacros.inc

@@ -358,6 +358,7 @@ PTFF_SAVE_R15           equ 00000080h
 PTFF_SAVE_ALL_PRESERVED equ 000000F7h   ;; NOTE: RBP is not included in this set!
 PTFF_SAVE_RSP           equ 00008000h
 PTFF_SAVE_RAX           equ 00000100h   ;; RAX is saved in hijack handler - in case it contains a GC ref
+PTFF_SAVE_RCX           equ 00000200h   ;; RCX is saved in hijack handler - in case it contains a GC ref
 PTFF_SAVE_ALL_SCRATCH   equ 00007F00h
 PTFF_THREAD_HIJACK      equ 00100000h   ;; indicates that this is a frame for a hijacked call
 

src/coreclr/nativeaot/Runtime/amd64/GcProbe.S

@@ -6,8 +6,8 @@
 #include <unixasmmacros.inc>
 
 //
-// See PUSH_COOP_PINVOKE_FRAME, this macro is very similar, but also saves RAX/RDX and accepts the register
-// bitmask in RCX
+// See PUSH_COOP_PINVOKE_FRAME, this macro is very similar, but also saves RAX/RCX/RDX and accepts the register
+// bitmask in R8
 //
 // On entry:
 //  - BITMASK: bitmask describing pushes, may be volatile register or constant value
@@ -20,8 +20,9 @@
 //
 .macro PUSH_PROBE_FRAME threadReg, trashReg, BITMASK
     push_register    rdx                         // save RDX, it might contain an objectref
+    push_register    rcx                         // save RCX, it might contain an objectref (async continuation)
     push_register    rax                         // save RAX, it might contain an objectref
-    lea             \trashReg, [rsp + 0x18]
+    lea             \trashReg, [rsp + 0x20]
     push_register   \trashReg                    // save caller`s RSP
     push_nonvol_reg r15                          // save preserved registers
     push_nonvol_reg r14                          //   ..
@@ -31,12 +32,12 @@
     push_register   \BITMASK                     // save the register bitmask passed in by caller
     push_register   \threadReg                   // Thread * (unused by stackwalker)
     push_nonvol_reg rbp                          // save caller`s RBP
-    mov             \trashReg, [rsp + 11*8]      // Find the return address
+    mov             \trashReg, [rsp + 12*8]      // Find the return address
     push_register   \trashReg                    // save m_RIP
     lea             \trashReg, [rsp + 0]         // trashReg == address of frame
 
     // allocate space for xmm0, xmm1 and alignment
-    alloc_stack     0x28
+    alloc_stack     0x20 + 0
 
     // save xmm0 and xmm1 in case they are used as return values
     movdqa          [rsp + 0x10], xmm0
@@ -53,7 +54,7 @@
 .macro POP_PROBE_FRAME
     movdqa      xmm1, [rsp + 0]
     movdqa      xmm0, [rsp + 0x10]
-    add         rsp, 0x28 + 8  // skip xmm0, xmm1 and discard RIP
+    add         rsp, 0x20 + 8  // skip xmm0, xmm1 and discard RIP
     pop         rbp
     pop         rax     // discard Thread*
     pop         rax     // discard BITMASK
@@ -64,6 +65,7 @@
     pop         r15
     pop         rax     // discard caller RSP
     pop         rax
+    pop         rcx
     pop         rdx
 .endm
 
@@ -76,28 +78,38 @@
 //
 // Register state on exit:
 //  R11: thread pointer
-//  RAX, RDX preserved, other volatile regs trashed
+//  RAX, RCX, RDX preserved, other volatile regs trashed
 //
 .macro FixupHijackedCallstack
         // preserve RAX, RDX as they may contain return values
         push rax
         push rdx
 
+        // preserve RCX as it may contain async continuation return value
+        push rcx
+
+        // align stack
+        sub rsp, 0x8
+
         // rax = GetThread(), makes nested calls
         INLINE_GETTHREAD
         mov r11, rax
 
+        add rsp, 0x8
+
+        pop rcx
+
         pop rdx
         pop rax
 
         // Fix the stack by pushing the original return address
-        mov         rcx, [r11 + OFFSETOF__Thread__m_pvHijackedReturnAddress]
-        push        rcx
+        mov         r8, [r11 + OFFSETOF__Thread__m_pvHijackedReturnAddress]
+        push        r8
 
         // Clear hijack state
-        xor         ecx, ecx
-        mov         [r11 + OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation], rcx
-        mov         [r11 + OFFSETOF__Thread__m_pvHijackedReturnAddress], rcx
+        xor         r8, r8
+        mov         [r11 + OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation], r8
+        mov         [r11 + OFFSETOF__Thread__m_pvHijackedReturnAddress], r8
 .endm
 
 //
@@ -112,12 +124,12 @@ NESTED_ENTRY RhpGcProbeHijack, _TEXT, NoHandler
         ret
 
 LOCAL_LABEL(WaitForGC):
-        mov         ecx, DEFAULT_FRAME_SAVE_FLAGS  + PTFF_SAVE_RAX + PTFF_SAVE_RDX + PTFF_THREAD_HIJACK
+        mov         r8d, DEFAULT_FRAME_SAVE_FLAGS  + PTFF_SAVE_RAX + PTFF_SAVE_RCX + PTFF_SAVE_RDX + PTFF_THREAD_HIJACK
         jmp         C_FUNC(RhpWaitForGC)
 NESTED_END RhpGcProbeHijack, _TEXT
 
 NESTED_ENTRY RhpWaitForGC, _TEXT, NoHandler
-        PUSH_PROBE_FRAME r11, rax, rcx
+        PUSH_PROBE_FRAME r11, rax, r8
         END_PROLOGUE
 
         mov         rbx, r11

src/coreclr/nativeaot/Runtime/amd64/GcProbe.asm

@@ -4,12 +4,13 @@
 include AsmMacros.inc
 
 ;;
-;; See PUSH_COOP_PINVOKE_FRAME, this macro is very similar, but also saves RAX and accepts
+;; See PUSH_COOP_PINVOKE_FRAME, this macro is very similar, but also saves RAX/RCX and accepts
 ;; the register bitmask
 ;;
 ;; On entry:
 ;;  - BITMASK: bitmask describing pushes, a volatile register
 ;;  - RAX: managed function return value, may be an object or byref
+;;  - RCX: managed function return value (async continuation), may be an object
 ;;  - preserved regs: need to stay preserved, may contain objects or byrefs
 ;;
 ;; INVARIANTS
@@ -18,8 +19,9 @@ include AsmMacros.inc
 ;;
 PUSH_PROBE_FRAME macro threadReg, trashReg, BITMASK
 
+    push_vol_reg    rcx                         ; save RCX, it might contain an objectref (async continuation)
     push_vol_reg    rax                         ; save RAX, it might contain an objectref
-    lea             trashReg, [rsp + 10h]
+    lea             trashReg, [rsp + 18h]
     push_vol_reg    trashReg                    ; save caller's RSP
     push_nonvol_reg r15                         ; save preserved registers
     push_nonvol_reg r14                         ;   ..
@@ -31,12 +33,12 @@ PUSH_PROBE_FRAME macro threadReg, trashReg, BITMASK
     push_vol_reg    BITMASK                     ; save the register bitmask passed in by caller
     push_vol_reg    threadReg                   ; Thread * (unused by stackwalker)
     push_nonvol_reg rbp                         ; save caller's RBP
-    mov             trashReg, [rsp + 12*8]      ; Find the return address
+    mov             trashReg, [rsp + 13*8]      ; Find the return address
     push_vol_reg    trashReg                    ; save m_RIP
     lea             trashReg, [rsp + 0]         ; trashReg == address of frame
 
     ;; allocate scratch space and any required alignment
-    alloc_stack     20h + 10h
+    alloc_stack     20h + 10h + 8
 
     ;; save xmm0 in case it's being used as a return value
     movdqa          [rsp + 20h], xmm0
@@ -52,7 +54,7 @@ endm
 ;;
 POP_PROBE_FRAME macro
     movdqa      xmm0, [rsp + 20h]
-    add         rsp, 20h + 10h + 8  ; deallocate stack and discard saved m_RIP
+    add         rsp, 20h + 10h + 8 + 8  ; deallocate stack and discard saved m_RIP
     pop         rbp
     pop         rax     ; discard Thread*
     pop         rax     ; discard BITMASK
@@ -65,6 +67,7 @@ POP_PROBE_FRAME macro
     pop         r15
     pop         rax     ; discard caller RSP
     pop         rax
+    pop         rcx
 endm
 
 ;;
@@ -76,20 +79,20 @@ endm
 ;;
 ;; Register state on exit:
 ;;  RDX: thread pointer
-;;  RAX: preserved, other volatile regs trashed
+;;  RAX/RCX: preserved, other volatile regs trashed
 ;;
 FixupHijackedCallstack macro
-        ;; rdx <- GetThread(), TRASHES rcx
-        INLINE_GETTHREAD rdx, rcx
+        ;; rdx <- GetThread(), TRASHES r8
+        INLINE_GETTHREAD rdx, r8
 
         ;; Fix the stack by pushing the original return address
-        mov         rcx, [rdx + OFFSETOF__Thread__m_pvHijackedReturnAddress]
-        push        rcx
+        mov         r8, [rdx + OFFSETOF__Thread__m_pvHijackedReturnAddress]
+        push        r8
 
         ;; Clear hijack state
-        xor         ecx, ecx
-        mov         [rdx + OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation], rcx
-        mov         [rdx + OFFSETOF__Thread__m_pvHijackedReturnAddress], rcx
+        xor         r8, r8
+        mov         [rdx + OFFSETOF__Thread__m_ppvHijackedReturnAddressLocation], r8
+        mov         [rdx + OFFSETOF__Thread__m_pvHijackedReturnAddress], r8
 endm
 
 ;;
@@ -103,12 +106,12 @@ NESTED_ENTRY RhpGcProbeHijack, _TEXT
         jnz         @f
         ret
 @@:
-        mov         ecx, DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_RAX + PTFF_THREAD_HIJACK
+        mov         r8d, DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_RAX + PTFF_SAVE_RCX + PTFF_THREAD_HIJACK
         jmp         RhpWaitForGC
 NESTED_END RhpGcProbeHijack, _TEXT
 
 NESTED_ENTRY RhpWaitForGC, _TEXT
-        PUSH_PROBE_FRAME rdx, rax, rcx
+        PUSH_PROBE_FRAME rdx, rax, r8
         END_PROLOGUE
 
         mov         rbx, rdx
@@ -144,7 +147,7 @@ ifdef FEATURE_GC_STRESS
 ;;
 LEAF_ENTRY RhpGcStressHijack, _TEXT
         FixupHijackedCallstack
-        or          ecx, DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_RAX
+        mov         r8d, DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_RAX + PTFF_SAVE_RCX
         jmp         RhpGcStressProbe
 LEAF_END RhpGcStressHijack, _TEXT
 
@@ -155,14 +158,14 @@ LEAF_END RhpGcStressHijack, _TEXT
 ;;
 ;; Register state on entry:
 ;;  RDX: thread pointer
-;;  RCX: register bitmask
+;;  R8: register bitmask
 ;;
 ;; Register state on exit:
-;;  Scratch registers, except for RAX, have been trashed
+;;  Scratch registers, except for RAX/RCX, have been trashed
 ;;  All other registers restored as they were when the hijack was first reached.
 ;;
 NESTED_ENTRY RhpGcStressProbe, _TEXT
-        PUSH_PROBE_FRAME rdx, rax, rcx
+        PUSH_PROBE_FRAME rdx, rax, r8
         END_PROLOGUE
 
         call        RhpStressGc

src/coreclr/nativeaot/Runtime/unix/unixasmmacrosamd64.inc

@@ -243,6 +243,7 @@ C_FUNC(\Name):
 #define PTFF_SAVE_ALL_PRESERVED  0x000000F1   // NOTE: RBP is not included in this set!
 #define PTFF_SAVE_RSP            0x00008000
 #define PTFF_SAVE_RAX            0x00000100   // RAX is saved in hijack handler - in case it contains a GC ref
+#define PTFF_SAVE_RCX            0x00000200   // RCX is saved in hijack handler - in case it contains a GC ref
 #define PTFF_SAVE_RDX            0x00000400   // RDX is saved in hijack handler - in case it contains a GC ref
 #define PTFF_SAVE_ALL_SCRATCH    0x00007F00
 #define PTFF_THREAD_HIJACK       0x00100000   // indicates that this is a frame for a hijacked call