- Removed calls to needsUnwindTableEntry() so that VG is always

  spilled if there are streaming mode changes in the function.

- Added requiresVGSpill() to AArch64MachineFunctionInfo which returns
  true if the function has streaming mode changes and hasSVE is true.
  With this change, we will no longer spill VG at the beginning of
  functions if the target does not also have SVE.

- Removed SpilledStreamingVG flag from spillCalleeSavedRegisters.

- Rebased to include recent changes to the changeStreamingMode interface.

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -1559,7 +1559,7 @@ bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
      const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
 
      if ((!FuncAttrs.hasStreamingBody() && FuncAttrs.hasStreamingInterface()) ||
-         !AFI->hasStreamingModeChanges())
+         !AFI->requiresVGSpill(MF))
        return false;
 
      int64_t StreamingVGIdx = AFI->getStreamingVGIdx();

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -321,7 +321,7 @@ bool AArch64FrameLowering::homogeneousPrologEpilog(
     return false;
 
   auto *AFI = MF.getInfo<AArch64FunctionInfo>();
-  if (AFI->hasSwiftAsyncContext() || AFI->hasStreamingModeChanges())
+  if (AFI->hasSwiftAsyncContext() || AFI->requiresVGSpill(MF))
     return false;
 
   // If there are an odd number of GPRs before LR and FP in the CSRs list,
@@ -567,7 +567,7 @@ void AArch64FrameLowering::emitCalleeSavedGPRLocations(
       continue;
 
     assert(!Info.isSpilledToReg() && "Spilling to registers not implemented");
-    unsigned DwarfReg = TRI.getDwarfRegNum(Info.getReg(), true);
+    int64_t DwarfReg = TRI.getDwarfRegNum(Info.getReg(), true);
     int64_t Offset = MFI.getObjectOffset(FrameIdx) - getOffsetOfLocalArea();
 
     // Locally streaming functions save two values for VG, but we should only
@@ -1361,8 +1361,9 @@ static MachineBasicBlock::iterator convertCalleeSaveRestoreToSPPrePostIncDec(
   // vector length. Move past these instructions if necessary.
   unsigned Opc = MBBI->getOpcode();
   if (Opc == AArch64::CNTD_XPiI || Opc == AArch64::RDSVLI_XI) {
-    AArch64FunctionInfo AFI = *MBB.getParent()->getInfo<AArch64FunctionInfo>();
-    assert(AFI.hasStreamingModeChanges() &&
+    MachineFunction &MF = *MBB.getParent();
+    AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+    assert(AFI->requiresVGSpill(MF) &&
            "Unexpected callee-save save/restore opcode!");
     ++MBBI;
     if (MBBI->getOpcode() == AArch64::UBFMXri)
@@ -1862,10 +1863,11 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
          !IsSVECalleeSave(MBBI)) {
     unsigned Opc = MBBI->getOpcode();
     // Move past instructions generated to calculate VG
-    if (Opc == AArch64::CNTD_XPiI || Opc == AArch64::RDSVLI_XI ||
-        Opc == AArch64::UBFMXri) {
-      assert(AFI->hasStreamingModeChanges() && "Unexpected opcode!");
+    if (Opc == AArch64::CNTD_XPiI || Opc == AArch64::RDSVLI_XI) {
+      assert(AFI->requiresVGSpill(MF) && "Unexpected opcode!");
       ++MBBI;
+      if (MBBI->getOpcode() == AArch64::UBFMXri)
+        ++MBBI;
     }
     if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*MBBI, AFI->getLocalStackSize(),
@@ -3017,11 +3019,10 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
     ArrayRef<CalleeSavedInfo> CSI, const TargetRegisterInfo *TRI) const {
   MachineFunction &MF = *MBB.getParent();
   const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   bool NeedsWinCFI = needsWinCFI(MF);
   DebugLoc DL;
   SmallVector<RegPairInfo, 8> RegPairs;
-  bool SpilledStreamingVG = false;
-  MachineFrameInfo &MFI = MF.getFrameInfo();
 
   computeCalleeSaveRegisterPairs(MF, CSI, TRI, RegPairs, hasFP(MF));
 
@@ -3099,7 +3100,7 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
       SMEAttrs Attrs(MF.getFunction());
 
       if (Attrs.hasStreamingBody() && !Attrs.hasStreamingInterface() &&
-          !SpilledStreamingVG) {
+          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)
@@ -3111,9 +3112,7 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
             .addImm(63)
             .setMIFlag(MachineInstr::FrameSetup);
 
-        AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
         AFI->setStreamingVGIdx(RPI.FrameIdx);
-        SpilledStreamingVG = true;
       } else {
         BuildMI(MBB, MI, DL, TII.get(AArch64::CNTD_XPiI), Reg1)
             .addImm(31)
@@ -3394,7 +3393,7 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
   // non-streaming VG value.
   const Function &F = MF.getFunction();
   SMEAttrs Attrs(F);
-  if (AFI->hasStreamingModeChanges() && F.needsUnwindTableEntry()) {
+  if (AFI->requiresVGSpill(MF)) {
     if (Attrs.hasStreamingBody() && !Attrs.hasStreamingInterface())
       CSStackSize += 16;
     else
@@ -3538,8 +3537,7 @@ bool AArch64FrameLowering::assignCalleeSavedSpillSlots(
   }
 
   // Insert VG into the list of CSRs, immediately before LR if saved.
-  const Function &F = MF.getFunction();
-  if (AFI->hasStreamingModeChanges() && F.needsUnwindTableEntry()) {
+  if (AFI->requiresVGSpill(MF)) {
     std::vector<CalleeSavedInfo> VGSaves;
     SMEAttrs Attrs(MF.getFunction());
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -8336,7 +8336,7 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     Chain = NewChain.getValue(0);
     InGlue = NewChain.getValue(1);
 
-    if (IsLocallyStreaming && MF.getFunction().needsUnwindTableEntry()) {
+    if (IsLocallyStreaming && Subtarget->hasSVE()) {
       NewChain = DAG.getNode(
           AArch64ISD::VG_UNWIND, DL, DAG.getVTList(MVT::Other, MVT::Glue),
           {Chain, DAG.getTargetConstant(/*Save*/ 1, DL, MVT::i64), InGlue});
@@ -8498,7 +8498,7 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
         DAG, DL, !CalleeAttrs.hasStreamingInterface(), Result, InGlue,
         getSMCondition(CallerAttrs, CalleeAttrs), PStateSM);
 
-    if (IsLocallyStreaming && MF.getFunction().needsUnwindTableEntry())
+    if (IsLocallyStreaming && Subtarget->hasSVE())
       Result = DAG.getNode(
           AArch64ISD::VG_UNWIND, DL, MVT::Other,
           {Result, DAG.getTargetConstant(/*Restore*/ 0, DL, MVT::i64)});

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 #define LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 
+#include "AArch64Subtarget.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -488,6 +489,11 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
     HasStreamingModeChanges = HasChanges;
   }
 
+  bool requiresVGSpill(const MachineFunction &MF) const {
+    const AArch64Subtarget &STI = MF.getSubtarget<AArch64Subtarget>();
+    return STI.hasSVE() && HasStreamingModeChanges;
+  }
+
   bool hasStackProbing() const { return StackProbeSize != 0; }
 
   int64_t getStackProbeSize() const { return StackProbeSize; }

llvm/test/CodeGen/AArch64/sme-call-streaming-compatible-to-normal-fn-wihout-sme-attr.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
-; RUN: llc -mattr=+sve < %s | FileCheck %s
+; RUN: llc < %s | FileCheck %s
 
 ; Verify that the following code can be compiled without +sme, because if the
 ; call is not entered in streaming-SVE mode at runtime, the codepath leading