[AArch64][SME] Move saving VG for the unwinder out of frame lowering

This patch moves computing and storing of VG in functions with streaming
mode changes to a new pre-RA pass (MachineSMEABI). The goal is to make
saving VG simpler, as computing VG may require calling
`__arm_get_current_vg` -- which requires saving X0 around the call and
the LR (among complexities in frame lowering). Doing this pre-RA allows
the register allocator to handle this (rather than manual scavenging).

The MachineSMEABI saves to VG to AArch64::SAVED_STREAMING_VG_SLOT and
AArch64::SAVED_VG_SLOT target frame indices. These will be resolved to
an actual frame indices during PEI (as they are not known before then).

For the most part this does not significantly change codegen, however,
there is one downside that resolving the frame indices outside of
the prologue may need exta instructions (to step past later allocations
on the stack, such as scalable vectors).

Fixes #145635

Author: MacDue

llvm/lib/CodeGen/MachineVerifier.cpp

@@ -2529,8 +2529,8 @@ MachineVerifier::visitMachineOperand(const MachineOperand *MO, unsigned MONum) {
       }
 
       // Check that an instruction has register operands only as expected.
-      if (MCOI.OperandType == MCOI::OPERAND_REGISTER &&
-          !MO->isReg() && !MO->isFI())
+      if (MCOI.OperandType == MCOI::OPERAND_REGISTER && !MO->isReg() &&
+          !MO->isFI() && !MO->isTargetIndex())
         report("Expected a register operand.", MO, MONum);
       if (MO->isReg()) {
         if (MCOI.OperandType == MCOI::OPERAND_IMMEDIATE ||

llvm/lib/Target/AArch64/AArch64.h

@@ -60,6 +60,7 @@ FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
 FunctionPass *createAArch64CollectLOHPass();
 FunctionPass *createSMEABIPass();
 FunctionPass *createSMEPeepholeOptPass();
+FunctionPass *createMachineSMEABIPass();
 ModulePass *createSVEIntrinsicOptsPass();
 InstructionSelector *
 createAArch64InstructionSelector(const AArch64TargetMachine &,
@@ -111,8 +112,14 @@ void initializeFalkorMarkStridedAccessesLegacyPass(PassRegistry&);
 void initializeLDTLSCleanupPass(PassRegistry&);
 void initializeSMEABIPass(PassRegistry &);
 void initializeSMEPeepholeOptPass(PassRegistry &);
+void initializeMachineSMEABIPass(PassRegistry &);
 void initializeSVEIntrinsicOptsPass(PassRegistry &);
 void initializeAArch64Arm64ECCallLoweringPass(PassRegistry &);
+
+namespace AArch64 {
+enum TargetIndex { SAVED_VG_SLOT, SAVED_STREAMING_VG_SLOT };
+}
+
 } // end namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -333,7 +333,10 @@ static bool needsWinCFI(const MachineFunction &MF);
 static StackOffset getSVEStackSize(const MachineFunction &MF);
 static Register findScratchNonCalleeSaveRegister(MachineBasicBlock *MBB,
                                                  bool HasCall = false);
-static bool requiresSaveVG(const MachineFunction &MF);
+static bool requiresSaveVG(const MachineFunction &MF) {
+  return MF.getSubtarget<AArch64Subtarget>().getRegisterInfo()->requiresSaveVG(
+      MF);
+}
 
 /// Returns true if a homogeneous prolog or epilog code can be emitted
 /// for the size optimization. If possible, a frame helper call is injected.
@@ -1105,8 +1108,7 @@ bool AArch64FrameLowering::canUseAsPrologue(
 
   // May need a scratch register (for return value) if require making a special
   // call
-  if (requiresSaveVG(*MF) ||
-      windowsRequiresStackProbe(*MF, std::numeric_limits<uint64_t>::max()))
+  if (windowsRequiresStackProbe(*MF, std::numeric_limits<uint64_t>::max()))
     if (findScratchNonCalleeSaveRegister(TmpMBB, true) == AArch64::NoRegister)
       return false;
 
@@ -1391,38 +1393,6 @@ bool requiresGetVGCall(MachineFunction &MF) {
          !MF.getSubtarget<AArch64Subtarget>().hasSVE();
 }
 
-static bool requiresSaveVG(const MachineFunction &MF) {
-  const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
-  // For Darwin platforms we don't save VG for non-SVE functions, even if SME
-  // is enabled with streaming mode changes.
-  if (!AFI->hasStreamingModeChanges())
-    return false;
-  auto &ST = MF.getSubtarget<AArch64Subtarget>();
-  if (ST.isTargetDarwin())
-    return ST.hasSVE();
-  return true;
-}
-
-bool isVGInstruction(MachineBasicBlock::iterator MBBI) {
-  unsigned Opc = MBBI->getOpcode();
-  if (Opc == AArch64::CNTD_XPiI || Opc == AArch64::RDSVLI_XI ||
-      Opc == AArch64::UBFMXri)
-    return true;
-
-  if (requiresGetVGCall(*MBBI->getMF())) {
-    if (Opc == AArch64::ORRXrr)
-      return true;
-
-    if (Opc == AArch64::BL) {
-      auto Op1 = MBBI->getOperand(0);
-      return Op1.isSymbol() &&
-             (StringRef(Op1.getSymbolName()) == "__arm_get_current_vg");
-    }
-  }
-
-  return false;
-}
-
 // Convert callee-save register save/restore instruction to do stack pointer
 // decrement/increment to allocate/deallocate the callee-save stack area by
 // converting store/load to use pre/post increment version.
@@ -1434,15 +1404,6 @@ static MachineBasicBlock::iterator convertCalleeSaveRestoreToSPPrePostIncDec(
     int CFAOffset = 0) {
   unsigned NewOpc;
 
-  // If the function contains streaming mode changes, we expect instructions
-  // to calculate the value of VG before spilling. For locally-streaming
-  // functions, we need to do this for both the streaming and non-streaming
-  // vector length. Move past these instructions if necessary.
-  MachineFunction &MF = *MBB.getParent();
-  if (requiresSaveVG(MF))
-    while (isVGInstruction(MBBI))
-      ++MBBI;
-
   switch (MBBI->getOpcode()) {
   default:
     llvm_unreachable("Unexpected callee-save save/restore opcode!");
@@ -1979,9 +1940,7 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   // pointer bump above.
   while (MBBI != End && MBBI->getFlag(MachineInstr::FrameSetup) &&
          !IsSVECalleeSave(MBBI)) {
-    if (CombineSPBump &&
-        // Only fix-up frame-setup load/store instructions.
-        (!requiresSaveVG(MF) || !isVGInstruction(MBBI)))
+    if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*MBBI, AFI->getLocalStackSize(),
                                         NeedsWinCFI, &HasWinCFI);
     ++MBBI;
@@ -3403,66 +3362,19 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
       StrOpc =
           Size == 16 ? AArch64::SPILL_PPR_TO_ZPR_SLOT_PSEUDO : AArch64::STR_PXI;
       break;
-    case RegPairInfo::VG:
-      StrOpc = AArch64::STRXui;
-      break;
-    }
-
-    unsigned X0Scratch = AArch64::NoRegister;
-    if (Reg1 == AArch64::VG) {
-      // Find an available register to store value of VG to.
-      Reg1 = findScratchNonCalleeSaveRegister(&MBB, true);
-      assert(Reg1 != AArch64::NoRegister);
+    case RegPairInfo::VG: {
       SMEAttrs Attrs = AFI->getSMEFnAttrs();
-
       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);
-
+        // & non-streaming VG.
         AFI->setStreamingVGIdx(RPI.FrameIdx);
-      } else if (MF.getSubtarget<AArch64Subtarget>().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>();
-        if (llvm::any_of(
-                MBB.liveins(),
-                [&STI](const MachineBasicBlock::RegisterMaskPair &LiveIn) {
-                  return 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);
       }
+      // VG will be written to the frame indices immediately after the prologue.
+      continue;
+    }
     }
 
     LLVM_DEBUG(dbgs() << "CSR spill: (" << printReg(Reg1, TRI);
@@ -3556,13 +3468,6 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
       if (RPI.isPaired())
         MFI.setStackID(FrameIdxReg2, TargetStackID::ScalableVector);
     }
-
-    if (X0Scratch != AArch64::NoRegister)
-      BuildMI(MBB, MI, DL, TII.get(AArch64::ORRXrr), AArch64::X0)
-          .addReg(AArch64::XZR)
-          .addReg(X0Scratch, RegState::Undef)
-          .addReg(X0Scratch, RegState::Implicit)
-          .setMIFlag(MachineInstr::FrameSetup);
   }
   return true;
 }
@@ -4092,31 +3997,19 @@ bool AArch64FrameLowering::assignCalleeSavedSpillSlots(
       MaxCSFrameIndex = FrameIdx;
   }
 
-  // Insert VG into the list of CSRs, immediately before LR if saved.
   if (requiresSaveVG(MF)) {
-    std::vector<CalleeSavedInfo> VGSaves;
-    SMEAttrs Attrs = AFI->getSMEFnAttrs();
-
-    auto VGInfo = CalleeSavedInfo(AArch64::VG);
+    CalleeSavedInfo VGInfo(AArch64::VG);
     VGInfo.setRestored(false);
-    VGSaves.push_back(VGInfo);
+    SmallVector<CalleeSavedInfo, 2> VGSaves{VGInfo};
 
     // Add VG again if the function is locally-streaming, as we will spill two
     // values.
+    SMEAttrs Attrs = AFI->getSMEFnAttrs();
     if (Attrs.hasStreamingBody() && !Attrs.hasStreamingInterface())
       VGSaves.push_back(VGInfo);
 
-    bool InsertBeforeLR = false;
-
-    for (unsigned I = 0; I < CSI.size(); I++)
-      if (CSI[I].getReg() == AArch64::LR) {
-        InsertBeforeLR = true;
-        CSI.insert(CSI.begin() + I, VGSaves.begin(), VGSaves.end());
-        break;
-      }
-
-    if (!InsertBeforeLR)
-      llvm::append_range(CSI, VGSaves);
+    // Insert the VG saves at the start of the CSI (alongside GPRs).
+    CSI.insert(CSI.begin(), VGSaves.begin(), VGSaves.end());
   }
 
   Register LastReg = 0;
@@ -5135,13 +5028,28 @@ static void emitVGSaveRestore(MachineBasicBlock::iterator II,
   MI.eraseFromParent();
 }
 
+static void replaceVGTargetIndices(MachineBasicBlock::iterator II,
+                                   AArch64FunctionInfo *AFI) {
+  for (auto &MO : II->explicit_operands()) {
+    if (MO.isTargetIndex()) {
+      if (MO.getIndex() == AArch64::SAVED_STREAMING_VG_SLOT)
+        MO.ChangeToFrameIndex(AFI->getStreamingVGIdx());
+      if (MO.getIndex() == AArch64::SAVED_VG_SLOT)
+        MO.ChangeToFrameIndex(AFI->getVGIdx());
+    }
+  }
+}
+
 void AArch64FrameLowering::processFunctionBeforeFrameIndicesReplaced(
     MachineFunction &MF, RegScavenger *RS = nullptr) const {
+  bool VGSaved = requiresSaveVG(MF);
+  AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   for (auto &BB : MF)
     for (MachineBasicBlock::iterator II = BB.begin(); II != BB.end();) {
-      if (requiresSaveVG(MF))
+      if (VGSaved) {
+        replaceVGTargetIndices(II, AFI);
         emitVGSaveRestore(II++, this);
-      else if (StackTaggingMergeSetTag)
+      } else if (StackTaggingMergeSetTag)
         II = tryMergeAdjacentSTG(II, this, RS);
     }
 }

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -3688,6 +3688,14 @@ static unsigned offsetExtendOpcode(unsigned Opcode) {
   }
 }
 
+ArrayRef<std::pair<int, const char *>>
+AArch64InstrInfo::getSerializableTargetIndices() const {
+  static constexpr std::pair<int, const char *> TargetIndices[] = {
+      {AArch64::SAVED_VG_SLOT, "saved-vg-slot"},
+      {AArch64::SAVED_STREAMING_VG_SLOT, "saved-streaming-vg-slot"}};
+  return TargetIndices;
+}
+
 MachineInstr *AArch64InstrInfo::emitLdStWithAddr(MachineInstr &MemI,
                                                  const ExtAddrMode &AM) const {
 

llvm/lib/Target/AArch64/AArch64InstrInfo.h

@@ -299,6 +299,9 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
   MachineInstr *emitLdStWithAddr(MachineInstr &MemI,
                                  const ExtAddrMode &AM) const override;
 
+  ArrayRef<std::pair<int, const char *>>
+  getSerializableTargetIndices() const override;
+
   bool getMemOperandsWithOffsetWidth(
       const MachineInstr &MI, SmallVectorImpl<const MachineOperand *> &BaseOps,
       int64_t &Offset, bool &OffsetIsScalable, LocationSize &Width,

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

@@ -673,6 +673,18 @@ bool AArch64RegisterInfo::hasBasePointer(const MachineFunction &MF) const {
   return false;
 }
 
+bool AArch64RegisterInfo::requiresSaveVG(const MachineFunction &MF) const {
+  const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+  // For Darwin platforms we don't save VG for non-SVE functions, even if SME
+  // is enabled with streaming mode changes.
+  if (!AFI->hasStreamingModeChanges())
+    return false;
+  auto &ST = MF.getSubtarget<AArch64Subtarget>();
+  if (ST.isTargetDarwin())
+    return ST.hasSVE();
+  return true;
+}
+
 bool AArch64RegisterInfo::isArgumentRegister(const MachineFunction &MF,
                                              MCRegister Reg) const {
   CallingConv::ID CC = MF.getFunction().getCallingConv();

llvm/lib/Target/AArch64/AArch64RegisterInfo.h

@@ -127,6 +127,8 @@ class AArch64RegisterInfo final : public AArch64GenRegisterInfo {
   bool hasBasePointer(const MachineFunction &MF) const;
   unsigned getBaseRegister() const;
 
+  bool requiresSaveVG(const MachineFunction &MF) const;
+
   bool isArgumentRegister(const MachineFunction &MF,
                           MCRegister Reg) const override;
 

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

@@ -803,6 +803,9 @@ bool AArch64PassConfig::addILPOpts() {
 }
 
 void AArch64PassConfig::addPreRegAlloc() {
+  // Insert VG saves for the unwinder.
+  addPass(createMachineSMEABIPass());
+
   // Change dead register definitions to refer to the zero register.
   if (TM->getOptLevel() != CodeGenOptLevel::None &&
       EnableDeadRegisterElimination)

llvm/lib/Target/AArch64/CMakeLists.txt

@@ -87,6 +87,7 @@ add_llvm_target(AArch64CodeGen
   AArch64TargetObjectFile.cpp
   AArch64TargetTransformInfo.cpp
   SMEABIPass.cpp
+  MachineSMEABIPass.cpp
   SMEPeepholeOpt.cpp
   SVEIntrinsicOpts.cpp
   AArch64SIMDInstrOpt.cpp