8360775: Fix Shenandoah GC test failures when APX is enabled

Reviewed-by: sviswanathan, jbhateja, epeter

Author: Srinivas Vamsi Parasa

src/hotspot/cpu/x86/assembler_x86.cpp

@@ -15761,6 +15761,8 @@ void Assembler::pusha_uncached() { // 64bit
     // Push pair of original stack pointer along with remaining registers
     // at 16B aligned boundary.
     push2p(rax, r31);
+    // Restore the original contents of RAX register.
+    movq(rax, Address(rax));
     push2p(r30, r29);
     push2p(r28, r27);
     push2p(r26, r25);

src/hotspot/cpu/x86/gc/shenandoah/shenandoahBarrierSetAssembler_x86.cpp

@@ -353,7 +353,7 @@ void ShenandoahBarrierSetAssembler::load_reference_barrier(MacroAssembler* masm,
 
   // The rest is saved with the optimized path
 
-  uint num_saved_regs = 4 + (dst != rax ? 1 : 0) + 4;
+  uint num_saved_regs = 4 + (dst != rax ? 1 : 0) + 4 + (UseAPX ? 16 : 0);
   __ subptr(rsp, num_saved_regs * wordSize);
   uint slot = num_saved_regs;
   if (dst != rax) {
@@ -367,6 +367,25 @@ void ShenandoahBarrierSetAssembler::load_reference_barrier(MacroAssembler* masm,
   __ movptr(Address(rsp, (--slot) * wordSize), r9);
   __ movptr(Address(rsp, (--slot) * wordSize), r10);
   __ movptr(Address(rsp, (--slot) * wordSize), r11);
+  // Save APX extended registers r16–r31 if enabled
+  if (UseAPX) {
+    __ movptr(Address(rsp, (--slot) * wordSize), r16);
+    __ movptr(Address(rsp, (--slot) * wordSize), r17);
+    __ movptr(Address(rsp, (--slot) * wordSize), r18);
+    __ movptr(Address(rsp, (--slot) * wordSize), r19);
+    __ movptr(Address(rsp, (--slot) * wordSize), r20);
+    __ movptr(Address(rsp, (--slot) * wordSize), r21);
+    __ movptr(Address(rsp, (--slot) * wordSize), r22);
+    __ movptr(Address(rsp, (--slot) * wordSize), r23);
+    __ movptr(Address(rsp, (--slot) * wordSize), r24);
+    __ movptr(Address(rsp, (--slot) * wordSize), r25);
+    __ movptr(Address(rsp, (--slot) * wordSize), r26);
+    __ movptr(Address(rsp, (--slot) * wordSize), r27);
+    __ movptr(Address(rsp, (--slot) * wordSize), r28);
+    __ movptr(Address(rsp, (--slot) * wordSize), r29);
+    __ movptr(Address(rsp, (--slot) * wordSize), r30);
+    __ movptr(Address(rsp, (--slot) * wordSize), r31);
+  }
   // r12-r15 are callee saved in all calling conventions
   assert(slot == 0, "must use all slots");
 
@@ -398,6 +417,25 @@ void ShenandoahBarrierSetAssembler::load_reference_barrier(MacroAssembler* masm,
     __ super_call_VM_leaf(CAST_FROM_FN_PTR(address, ShenandoahRuntime::load_reference_barrier_phantom), arg0, arg1);
   }
 
+  // Restore APX extended registers r31–r16 if previously saved
+  if (UseAPX) {
+    __ movptr(r31, Address(rsp, (slot++) * wordSize));
+    __ movptr(r30, Address(rsp, (slot++) * wordSize));
+    __ movptr(r29, Address(rsp, (slot++) * wordSize));
+    __ movptr(r28, Address(rsp, (slot++) * wordSize));
+    __ movptr(r27, Address(rsp, (slot++) * wordSize));
+    __ movptr(r26, Address(rsp, (slot++) * wordSize));
+    __ movptr(r25, Address(rsp, (slot++) * wordSize));
+    __ movptr(r24, Address(rsp, (slot++) * wordSize));
+    __ movptr(r23, Address(rsp, (slot++) * wordSize));
+    __ movptr(r22, Address(rsp, (slot++) * wordSize));
+    __ movptr(r21, Address(rsp, (slot++) * wordSize));
+    __ movptr(r20, Address(rsp, (slot++) * wordSize));
+    __ movptr(r19, Address(rsp, (slot++) * wordSize));
+    __ movptr(r18, Address(rsp, (slot++) * wordSize));
+    __ movptr(r17, Address(rsp, (slot++) * wordSize));
+    __ movptr(r16, Address(rsp, (slot++) * wordSize));
+  }
   __ movptr(r11, Address(rsp, (slot++) * wordSize));
   __ movptr(r10, Address(rsp, (slot++) * wordSize));
   __ movptr(r9,  Address(rsp, (slot++) * wordSize));