[clang] [llvm] [AArch64][SME] Save VG for unwind info when changing streaming-mode (PR #83301)

[Sander de Smalen via cfe-commits]

================
@@ -3062,7 +3131,68 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
        Size = 2;
        Alignment = Align(2);
        break;
+    case RegPairInfo::VG:
+      StrOpc = AArch64::STRXui;
+      Size = 8;
+      Alignment = Align(8);
+      break;
+    }
+
+    unsigned X0Scratch = AArch64::NoRegister;
+    if (Reg1 == AArch64::VG) {
+      // Find an available register to store value of VG to.
+      Reg1 = findScratchNonCalleeSaveRegister(&MBB);
+      assert(Reg1 != AArch64::NoRegister);
+      SMEAttrs Attrs(MF.getFunction());
+
+      if (Attrs.hasStreamingBody() && !Attrs.hasStreamingInterface() &&
+          AFI->getStreamingVGIdx() == std::numeric_limits<int>::max()) {
+        // For locally-streaming functions, we need to store both the streaming
+        // & non-streaming VG. Spill the streaming value first.
+        BuildMI(MBB, MI, DL, TII.get(AArch64::RDSVLI_XI), Reg1)
+            .addImm(1)
+            .setMIFlag(MachineInstr::FrameSetup);
+        BuildMI(MBB, MI, DL, TII.get(AArch64::UBFMXri), Reg1)
+            .addReg(Reg1)
+            .addImm(3)
+            .addImm(63)
+            .setMIFlag(MachineInstr::FrameSetup);
+
+        AFI->setStreamingVGIdx(RPI.FrameIdx);
+      } else {
+        if (HasSVE)
+          BuildMI(MBB, MI, DL, TII.get(AArch64::CNTD_XPiI), Reg1)
+              .addImm(31)
+              .addImm(1)
+              .setMIFlag(MachineInstr::FrameSetup);
+        else {
+          const AArch64Subtarget &STI = MF.getSubtarget<AArch64Subtarget>();
+          for (const auto &LiveIn : MBB.liveins())
+            if (STI.getRegisterInfo()->isSuperOrSubRegisterEq(AArch64::X0,
+                                                              LiveIn.PhysReg))
+              X0Scratch = Reg1;
+
+          if (X0Scratch != AArch64::NoRegister)
+            BuildMI(MBB, MI, DL, TII.get(AArch64::ORRXrr), Reg1)
+                .addReg(AArch64::XZR)
+                .addReg(AArch64::X0, RegState::Undef)
+                .addReg(AArch64::X0, RegState::Implicit)
+                .setMIFlag(MachineInstr::FrameSetup);
+
+          const uint32_t *RegMask = TRI->getCallPreservedMask(
+              MF, CallingConv::
+                      AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1);
+          BuildMI(MBB, MI, DL, TII.get(AArch64::BL))
+              .addExternalSymbol("__arm_get_current_vg")
+              .addRegMask(RegMask)
+              .addReg(AArch64::X0, RegState::ImplicitDefine)
+              .setMIFlag(MachineInstr::FrameSetup);
+          Reg1 = AArch64::X0;
+        }
+        AFI->setVGIdx(RPI.FrameIdx);
+      }
----------------
sdesmalen-arm wrote:

nit: put `AFI->setVGIdx(RPI.FrameIdx);` in both the `if(HasSVE)` and the `else` branch, and remove a level of indentation?

```suggestion
      } else if (HasSVE) {
        BuildMI(MBB, MI, DL, TII.get(AArch64::CNTD_XPiI), Reg1)
            .addImm(31)
            .addImm(1)
            .setMIFlag(MachineInstr::FrameSetup);
        AFI->setVGIdx(RPI.FrameIdx);
      } else { 
        const AArch64Subtarget &STI = MF.getSubtarget<AArch64Subtarget>();
        for (const auto &LiveIn : MBB.liveins())
          if (STI.getRegisterInfo()->isSuperOrSubRegisterEq(AArch64::X0,
                                                            LiveIn.PhysReg))
            X0Scratch = Reg1;
              
        if (X0Scratch != AArch64::NoRegister)
          BuildMI(MBB, MI, DL, TII.get(AArch64::ORRXrr), Reg1)
              .addReg(AArch64::XZR)
              .addReg(AArch64::X0, RegState::Undef)
              .addReg(AArch64::X0, RegState::Implicit)
              .setMIFlag(MachineInstr::FrameSetup);
                
        const uint32_t *RegMask = TRI->getCallPreservedMask(
            MF, CallingConv::
                    AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1);
        BuildMI(MBB, MI, DL, TII.get(AArch64::BL))
            .addExternalSymbol("__arm_get_current_vg")
            .addRegMask(RegMask)
            .addReg(AArch64::X0, RegState::ImplicitDefine)
            .setMIFlag(MachineInstr::FrameSetup);
        Reg1 = AArch64::X0;
        AFI->setVGIdx(RPI.FrameIdx);
      }
```

https://github.com/llvm/llvm-project/pull/83301