8326541: [AArch64] ZGC C2 load barrier stub should consider the length of live registers when spilling registers

Reviewed-by: eosterlund, rcastanedalo

Author: Joshua Zhu

src/hotspot/cpu/aarch64/aarch64.ad

@@ -1126,6 +1126,10 @@ extern RegMask _NO_SPECIAL_REG32_mask;
 extern RegMask _NO_SPECIAL_REG_mask;
 extern RegMask _NO_SPECIAL_PTR_REG_mask;
 
+// Figure out which register class each belongs in: rc_int, rc_float or
+// rc_stack.
+enum RC { rc_bad, rc_int, rc_float, rc_predicate, rc_stack };
+
 class CallStubImpl {
 
   //--------------------------------------------------------------
@@ -1900,10 +1904,6 @@ const Pipeline * MachEpilogNode::pipeline() const {
 
 //=============================================================================
 
-// Figure out which register class each belongs in: rc_int, rc_float or
-// rc_stack.
-enum RC { rc_bad, rc_int, rc_float, rc_predicate, rc_stack };
-
 static enum RC rc_class(OptoReg::Name reg) {
 
   if (reg == OptoReg::Bad) {

src/hotspot/cpu/aarch64/gc/z/zBarrierSetAssembler_aarch64.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2019, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -1081,14 +1081,43 @@ void ZBarrierSetAssembler::generate_c1_store_barrier_runtime_stub(StubAssembler*
 
 #ifdef COMPILER2
 
+OptoReg::Name ZBarrierSetAssembler::encode_float_vector_register_size(const Node* node, OptoReg::Name opto_reg) {
+  switch (node->ideal_reg()) {
+    case Op_RegF:
+      // No need to refine. The original encoding is already fine to distinguish.
+      assert(opto_reg % 4 == 0, "Float register should only occupy a single slot");
+      break;
+    // Use different encoding values of the same fp/vector register to help distinguish different sizes.
+    // Such as V16. The OptoReg::name and its corresponding slot value are
+    // "V16": 64, "V16_H": 65, "V16_J": 66, "V16_K": 67.
+    case Op_RegD:
+    case Op_VecD:
+      opto_reg &= ~3;
+      opto_reg |= 1;
+      break;
+    case Op_VecX:
+      opto_reg &= ~3;
+      opto_reg |= 2;
+      break;
+    case Op_VecA:
+      opto_reg &= ~3;
+      opto_reg |= 3;
+      break;
+    default:
+      assert(false, "unexpected ideal register");
+      ShouldNotReachHere();
+  }
+  return opto_reg;
+}
+
 OptoReg::Name ZBarrierSetAssembler::refine_register(const Node* node, OptoReg::Name opto_reg) {
   if (!OptoReg::is_reg(opto_reg)) {
     return OptoReg::Bad;
   }
 
   const VMReg vm_reg = OptoReg::as_VMReg(opto_reg);
   if (vm_reg->is_FloatRegister()) {
-    return opto_reg & ~1;
+    opto_reg = encode_float_vector_register_size(node, opto_reg);
   }
 
   return opto_reg;
@@ -1099,28 +1128,82 @@ OptoReg::Name ZBarrierSetAssembler::refine_register(const Node* node, OptoReg::N
 
 class ZSaveLiveRegisters {
 private:
+  struct RegisterData {
+    VMReg _reg;
+    int   _slots; // slots occupied once pushed into stack
+
+    // Used by GrowableArray::find()
+    bool operator == (const RegisterData& other) {
+      return _reg == other._reg;
+    }
+  };
+
   MacroAssembler* const _masm;
   RegSet                _gp_regs;
   FloatRegSet           _fp_regs;
+  FloatRegSet           _neon_regs;
+  FloatRegSet           _sve_regs;
   PRegSet               _p_regs;
 
 public:
   void initialize(ZBarrierStubC2* stub) {
-    // Record registers that needs to be saved/restored
+    int index = -1;
+    GrowableArray<RegisterData> registers;
+    VMReg prev_vm_reg = VMRegImpl::Bad();
+
     RegMaskIterator rmi(stub->live());
     while (rmi.has_next()) {
-      const OptoReg::Name opto_reg = rmi.next();
-      if (OptoReg::is_reg(opto_reg)) {
-        const VMReg vm_reg = OptoReg::as_VMReg(opto_reg);
-        if (vm_reg->is_Register()) {
-          _gp_regs += RegSet::of(vm_reg->as_Register());
-        } else if (vm_reg->is_FloatRegister()) {
+      OptoReg::Name opto_reg = rmi.next();
+      VMReg vm_reg = OptoReg::as_VMReg(opto_reg);
+
+      if (vm_reg->is_Register()) {
+        // GPR may have one or two slots in regmask
+        // Determine whether the current vm_reg is the same physical register as the previous one
+        if (is_same_register(vm_reg, prev_vm_reg)) {
+          registers.at(index)._slots++;
+        } else {
+          RegisterData reg_data = { vm_reg, 1 };
+          index = registers.append(reg_data);
+        }
+      } else if (vm_reg->is_FloatRegister()) {
+        // We have size encoding in OptoReg of stub->live()
+        // After encoding, float/neon/sve register has only one slot in regmask
+        // Decode it to get the actual size
+        VMReg vm_reg_base = vm_reg->as_FloatRegister()->as_VMReg();
+        int slots = decode_float_vector_register_size(opto_reg);
+        RegisterData reg_data = { vm_reg_base, slots };
+        index = registers.append(reg_data);
+      } else if (vm_reg->is_PRegister()) {
+        // PRegister has only one slot in regmask
+        RegisterData reg_data = { vm_reg, 1 };
+        index = registers.append(reg_data);
+      } else {
+        assert(false, "Unknown register type");
+        ShouldNotReachHere();
+      }
+      prev_vm_reg = vm_reg;
+    }
+
+    // Record registers that needs to be saved/restored
+    for (GrowableArrayIterator<RegisterData> it = registers.begin(); it != registers.end(); ++it) {
+      RegisterData reg_data = *it;
+      VMReg vm_reg = reg_data._reg;
+      int slots = reg_data._slots;
+      if (vm_reg->is_Register()) {
+        assert(slots == 1 || slots == 2, "Unexpected register save size");
+        _gp_regs += RegSet::of(vm_reg->as_Register());
+      } else if (vm_reg->is_FloatRegister()) {
+        if (slots == 1 || slots == 2) {
           _fp_regs += FloatRegSet::of(vm_reg->as_FloatRegister());
-        } else if (vm_reg->is_PRegister()) {
-          _p_regs += PRegSet::of(vm_reg->as_PRegister());
+        } else if (slots == 4) {
+          _neon_regs += FloatRegSet::of(vm_reg->as_FloatRegister());
         } else {
-          fatal("Unknown register type");
+          assert(slots == Matcher::scalable_vector_reg_size(T_FLOAT), "Unexpected register save size");
+          _sve_regs += FloatRegSet::of(vm_reg->as_FloatRegister());
         }
+      } else {
+        assert(vm_reg->is_PRegister() && slots == 1, "Unknown register type");
+        _p_regs += PRegSet::of(vm_reg->as_PRegister());
       }
     }
 
@@ -1130,27 +1213,84 @@ class ZSaveLiveRegisters {
     } else {
       _gp_regs -= RegSet::range(r19, r30) + RegSet::of(r8, r9);
     }
+
+    // Remove C-ABI SOE fp registers
+    _fp_regs -= FloatRegSet::range(v8, v15);
+  }
+
+  static enum RC rc_class(VMReg reg) {
+    if (reg->is_reg()) {
+      if (reg->is_Register()) {
+        return rc_int;
+      } else if (reg->is_FloatRegister()) {
+        return rc_float;
+      } else if (reg->is_PRegister()) {
+        return rc_predicate;
+      }
+    }
+    if (reg->is_stack()) {
+      return rc_stack;
+    }
+    return rc_bad;
+  }
+
+  static bool is_same_register(VMReg reg1, VMReg reg2) {
+    if (reg1 == reg2) {
+      return true;
+    }
+    if (rc_class(reg1) == rc_class(reg2)) {
+      if (reg1->is_Register()) {
+        return reg1->as_Register() == reg2->as_Register();
+      } else if (reg1->is_FloatRegister()) {
+        return reg1->as_FloatRegister() == reg2->as_FloatRegister();
+      } else if (reg1->is_PRegister()) {
+        return reg1->as_PRegister() == reg2->as_PRegister();
+      }
+    }
+    return false;
+  }
+
+  static int decode_float_vector_register_size(OptoReg::Name opto_reg) {
+    switch (opto_reg & 3) {
+      case 0:
+        return 1;
+      case 1:
+        return 2;
+      case 2:
+        return 4;
+      case 3:
+        return Matcher::scalable_vector_reg_size(T_FLOAT);
+      default:
+        ShouldNotReachHere();
+        return 0;
+    }
   }
 
   ZSaveLiveRegisters(MacroAssembler* masm, ZBarrierStubC2* stub)
     : _masm(masm),
       _gp_regs(),
       _fp_regs(),
+      _neon_regs(),
+      _sve_regs(),
       _p_regs() {
 
     // Figure out what registers to save/restore
     initialize(stub);
 
     // Save registers
     __ push(_gp_regs, sp);
-    __ push_fp(_fp_regs, sp);
+    __ push_fp(_fp_regs, sp, MacroAssembler::PushPopFp);
+    __ push_fp(_neon_regs, sp, MacroAssembler::PushPopNeon);
+    __ push_fp(_sve_regs, sp, MacroAssembler::PushPopSVE);
     __ push_p(_p_regs, sp);
   }
 
   ~ZSaveLiveRegisters() {
     // Restore registers
     __ pop_p(_p_regs, sp);
-    __ pop_fp(_fp_regs, sp);
+    __ pop_fp(_sve_regs, sp, MacroAssembler::PushPopSVE);
+    __ pop_fp(_neon_regs, sp, MacroAssembler::PushPopNeon);
+    __ pop_fp(_fp_regs, sp, MacroAssembler::PushPopFp);
 
     // External runtime call may clobber ptrue reg
     __ reinitialize_ptrue();

src/hotspot/cpu/aarch64/gc/z/zBarrierSetAssembler_aarch64.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2019, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -187,6 +187,9 @@ class ZBarrierSetAssembler : public ZBarrierSetAssemblerBase {
 #endif // COMPILER1
 
 #ifdef COMPILER2
+  OptoReg::Name encode_float_vector_register_size(const Node* node,
+                                                  OptoReg::Name opto_reg);
+
   OptoReg::Name refine_register(const Node* node,
                                 OptoReg::Name opto_reg);