Changes to handle locally streaming functions with streaming mode changes:

 - Emit both the streaming and non-streaming value of VG in the prologue of
   functions with the aarch64_pstate_sm_body attribute.

 - Added the VGUnwindInfoPseudo node which expands to either .cfi_restore or a
   .cfi_offset depending on the value of the immediate used (0 or 1 respectively).

 - VGUnwindInfoPseudo nodes are emitted with the smstop/smstart pair around
   calls to streaming-mode functions from a locally-streaming caller. The
   .cfi_offset will save the streaming-VG value, whilst the restore sets the rule
   for VG to the same as it was at the beginning of the function (non-streaming).

 - The frame index used for the streaming VG value is saved in AArch64FunctionInfo
   so that it can be used to calculate the offset when expanding the pseudo.

 - Added the @vg_locally_streaming_fn() test to sme-vg-to-stack.ll

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -29,6 +29,7 @@
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DebugLoc.h"
+#include "llvm/MC/MCDwarf.h"
 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CodeGen.h"
@@ -1552,6 +1553,53 @@ bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
    case AArch64::COALESCER_BARRIER_FPR128:
      MI.eraseFromParent();
      return true;
+   case AArch64::VGUnwindInfoPseudo: {
+     MachineFunction &MF = *MBB.getParent();
+     SMEAttrs FuncAttrs(MF.getFunction());
+     const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+
+     if ((!FuncAttrs.hasStreamingBody() && FuncAttrs.hasStreamingInterface()) ||
+         !AFI->hasStreamingModeChanges())
+       return false;
+
+     int64_t StreamingVGIdx = AFI->getStreamingVGIdx();
+     assert(StreamingVGIdx != std::numeric_limits<int>::max() &&
+            "Expected FrameIdx for Streaming-VG");
+
+     const TargetSubtargetInfo &STI = MF.getSubtarget();
+     const TargetInstrInfo &TII = *STI.getInstrInfo();
+     const TargetRegisterInfo &TRI = *STI.getRegisterInfo();
+     if (MI.getOperand(0).getImm() == 1) {
+       // This pseudo has been inserted after a streaming-mode change
+       // to save the streaming value of VG before a call.
+       // Calculate and emit the CFI offset using StreamingVGIdx.
+       MachineFrameInfo &MFI = MF.getFrameInfo();
+       const AArch64FrameLowering *TFI =
+           MF.getSubtarget<AArch64Subtarget>().getFrameLowering();
+
+       int64_t Offset =
+           MFI.getObjectOffset(StreamingVGIdx) - TFI->getOffsetOfLocalArea();
+       unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createOffset(
+           nullptr, TRI.getDwarfRegNum(AArch64::VG, true), Offset));
+       BuildMI(MBB, MBBI, MBBI->getDebugLoc(),
+               TII.get(TargetOpcode::CFI_INSTRUCTION))
+           .addCFIIndex(CFIIndex)
+           .setMIFlags(MachineInstr::FrameSetup);
+     } else {
+       // This is a restore of VG after returning from the call. Emit the
+       // .cfi_restore instruction, which sets the rule for VG to the same
+       // as it was on entry to the function.
+       ++MBBI;
+       DebugLoc DL = MI.getDebugLoc();
+       unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createRestore(
+           nullptr, TRI.getDwarfRegNum(AArch64::VG, true)));
+       BuildMI(MBB, MBBI, DL, TII.get(TargetOpcode::CFI_INSTRUCTION))
+           .addCFIIndex(CFIIndex);
+     }
+
+     MI.eraseFromParent();
+     return true;
+   }
    case AArch64::LD1B_2Z_IMM_PSEUDO:
      return expandMultiVecPseudo(
          MBB, MBBI, AArch64::ZPR2RegClass, AArch64::ZPR2StridedRegClass,

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -550,6 +550,7 @@ void AArch64FrameLowering::emitCalleeSavedGPRLocations(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const {
   MachineFunction &MF = *MBB.getParent();
   MachineFrameInfo &MFI = MF.getFrameInfo();
+  AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
 
   const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
   if (CSI.empty())
@@ -561,14 +562,20 @@ void AArch64FrameLowering::emitCalleeSavedGPRLocations(
   DebugLoc DL = MBB.findDebugLoc(MBBI);
 
   for (const auto &Info : CSI) {
-    if (MFI.getStackID(Info.getFrameIdx()) == TargetStackID::ScalableVector)
+    unsigned FrameIdx = Info.getFrameIdx();
+    if (MFI.getStackID(FrameIdx) == TargetStackID::ScalableVector)
       continue;
 
     assert(!Info.isSpilledToReg() && "Spilling to registers not implemented");
     unsigned 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
+    // emit the location of the non-streaming value here.
+    if (DwarfReg == TRI.getDwarfRegNum(AArch64::VG, true) &&
+        FrameIdx == AFI->getStreamingVGIdx())
+      continue;
 
-    int64_t Offset =
-        MFI.getObjectOffset(Info.getFrameIdx()) - getOffsetOfLocalArea();
     unsigned CFIIndex = MF.addFrameInst(
         MCCFIInstruction::createOffset(nullptr, DwarfReg, Offset));
     BuildMI(MBB, MBBI, DL, TII.get(TargetOpcode::CFI_INSTRUCTION))
@@ -1348,6 +1355,20 @@ 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.
+  unsigned Opc = MBBI->getOpcode();
+  if (Opc == AArch64::CNTD_XPiI || Opc == AArch64::RDSVLI_XI) {
+    AArch64FunctionInfo AFI = *MBB.getParent()->getInfo<AArch64FunctionInfo>();
+    assert(AFI.hasStreamingModeChanges() &&
+           "Unexpected callee-save save/restore opcode!");
+    ++MBBI;
+    if (MBBI->getOpcode() == AArch64::UBFMXri)
+      ++MBBI;
+  }
+
   switch (MBBI->getOpcode()) {
   default:
     llvm_unreachable("Unexpected callee-save save/restore opcode!");
@@ -1655,13 +1676,6 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
     LiveRegs.removeReg(AArch64::LR);
   }
 
-  // If the function contains streaming mode changes, we expect the first
-  // instruction of MBB to be a CNTD. Move past this instruction if found.
-  if (AFI->hasStreamingModeChanges() && F.needsUnwindTableEntry()) {
-    assert(MBBI->getOpcode() == AArch64::CNTD_XPiI && "Unexpected instruction");
-    MBBI = std::next(MBBI);
-  }
-
   auto VerifyClobberOnExit = make_scope_exit([&]() {
     if (NonFrameStart == MBB.end())
       return;
@@ -1846,6 +1860,13 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   // pointer bump above.
   while (MBBI != End && MBBI->getFlag(MachineInstr::FrameSetup) &&
          !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!");
+      ++MBBI;
+    }
     if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*MBBI, AFI->getLocalStackSize(),
                                         NeedsWinCFI, &HasWinCFI);
@@ -2999,6 +3020,8 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
   bool NeedsWinCFI = needsWinCFI(MF);
   DebugLoc DL;
   SmallVector<RegPairInfo, 8> RegPairs;
+  bool SpilledStreamingVG = false;
+  MachineFrameInfo &MFI = MF.getFrameInfo();
 
   computeCalleeSaveRegisterPairs(MF, CSI, TRI, RegPairs, hasFP(MF));
 
@@ -3073,10 +3096,30 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
       // 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() &&
+          !SpilledStreamingVG) {
+        // 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);
 
-      BuildMI(MBB, MBB.begin(), DL, TII.get(AArch64::CNTD_XPiI), Reg1)
-          .addImm(31)
-          .addImm(1);
+        AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+        AFI->setStreamingVGIdx(RPI.FrameIdx);
+        SpilledStreamingVG = true;
+      } else {
+        BuildMI(MBB, MI, DL, TII.get(AArch64::CNTD_XPiI), Reg1)
+            .addImm(31)
+            .addImm(1)
+            .setMIFlag(MachineInstr::FrameSetup);
+      }
     }
 
     LLVM_DEBUG(dbgs() << "CSR spill: (" << printReg(Reg1, TRI);
@@ -3122,7 +3165,6 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
     MachineFrameInfo &MFI = MF.getFrameInfo();
     if (RPI.Type == RegPairInfo::ZPR || RPI.Type == RegPairInfo::PPR)
       MFI.setStackID(RPI.FrameIdx, TargetStackID::ScalableVector);
-
   }
   return true;
 }
@@ -3348,9 +3390,16 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
 
   // Increase the callee-saved stack size if the function has streaming mode
   // changes, as we will need to spill the value of the VG register.
+  // For locally streaming functions, we spill both the streaming and
+  // non-streaming VG value.
   const Function &F = MF.getFunction();
-  if (AFI->hasStreamingModeChanges() && F.needsUnwindTableEntry())
-    CSStackSize += 8;
+  SMEAttrs Attrs(F);
+  if (AFI->hasStreamingModeChanges() && F.needsUnwindTableEntry()) {
+    if (Attrs.hasStreamingBody() && !Attrs.hasStreamingInterface())
+      CSStackSize += 16;
+    else
+      CSStackSize += 8;
+  }
 
   // Save number of saved regs, so we can easily update CSStackSize later.
   unsigned NumSavedRegs = SavedRegs.count();
@@ -3491,19 +3540,29 @@ 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()) {
+    std::vector<CalleeSavedInfo> VGSaves;
+    SMEAttrs Attrs(MF.getFunction());
+
     auto VGInfo = CalleeSavedInfo(AArch64::VG);
     VGInfo.setRestored(false);
+    VGSaves.push_back(VGInfo);
+
+    // Add VG again if the function is locally-streaming, as we will spill two
+    // values.
+    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, VGInfo);
+        CSI.insert(CSI.begin() + I, VGSaves.begin(), VGSaves.end());
         break;
       }
 
     if (!InsertBeforeLR)
-      CSI.push_back(VGInfo);
+      CSI.insert(CSI.end(), VGSaves.begin(), VGSaves.end());
   }
 
   for (auto &CS : CSI) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -2436,6 +2436,7 @@ const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
   case AArch64ISD::FIRST_NUMBER:
     break;
     MAKE_CASE(AArch64ISD::COALESCER_BARRIER)
+    MAKE_CASE(AArch64ISD::VG_UNWIND)
     MAKE_CASE(AArch64ISD::SMSTART)
     MAKE_CASE(AArch64ISD::SMSTOP)
     MAKE_CASE(AArch64ISD::RESTORE_ZA)
@@ -8326,12 +8327,22 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, MemOpChains);
 
   SDValue InGlue;
+  bool IsLocallyStreaming =
+      !CallerAttrs.hasStreamingInterface() && CallerAttrs.hasStreamingBody();
   if (RequiresSMChange) {
     SDValue NewChain = changeStreamingMode(
         DAG, DL, CalleeAttrs.hasStreamingInterface(), Chain, InGlue,
         getSMCondition(CallerAttrs, CalleeAttrs), PStateSM);
     Chain = NewChain.getValue(0);
     InGlue = NewChain.getValue(1);
+
+    if (IsLocallyStreaming && MF.getFunction().needsUnwindTableEntry()) {
+      NewChain = DAG.getNode(
+          AArch64ISD::VG_UNWIND, DL, DAG.getVTList(MVT::Other, MVT::Glue),
+          {Chain, DAG.getTargetConstant(/*Save*/ 1, DL, MVT::i64), InGlue});
+      Chain = NewChain.getValue(0);
+      InGlue = NewChain.getValue(1);
+    }
   }
 
   // Build a sequence of copy-to-reg nodes chained together with token chain
@@ -8486,6 +8497,11 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     Result = changeStreamingMode(
         DAG, DL, !CalleeAttrs.hasStreamingInterface(), Result, InGlue,
         getSMCondition(CallerAttrs, CalleeAttrs), PStateSM);
+
+    if (IsLocallyStreaming && MF.getFunction().needsUnwindTableEntry())
+      Result = DAG.getNode(
+          AArch64ISD::VG_UNWIND, DL, MVT::Other,
+          {Result, DAG.getTargetConstant(/*Restore*/ 0, DL, MVT::i64)});
   }
 
   if (CallerAttrs.requiresEnablingZAAfterCall(CalleeAttrs))

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -60,6 +60,8 @@ enum NodeType : unsigned {
 
   COALESCER_BARRIER,
 
+  VG_UNWIND,
+
   SMSTART,
   SMSTOP,
   RESTORE_ZA,

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -212,6 +212,8 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   // on function entry to record the initial pstate of a function.
   Register PStateSMReg = MCRegister::NoRegister;
 
+  int64_t StreamingVGIdx = std::numeric_limits<int>::max();
+
 public:
   AArch64FunctionInfo(const Function &F, const AArch64Subtarget *STI);
 
@@ -223,6 +225,9 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   Register getPStateSMReg() const { return PStateSMReg; };
   void setPStateSMReg(Register Reg) { PStateSMReg = Reg; };
 
+  int64_t getStreamingVGIdx() const { return StreamingVGIdx; };
+  void setStreamingVGIdx(unsigned Idx) { StreamingVGIdx = Idx; };
+
   bool isSVECC() const { return IsSVECC; };
   void setIsSVECC(bool s) { IsSVECC = s; };
 

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

@@ -31,6 +31,9 @@ def AArch64_save_zt : SDNode<"AArch64ISD::SAVE_ZT", SDTypeProfile<0, 2,
 def AArch64CoalescerBarrier
     : SDNode<"AArch64ISD::COALESCER_BARRIER", SDTypeProfile<1, 1, []>, [SDNPOptInGlue, SDNPOutGlue]>;
 
+def AArch64VGUnwind : SDNode<"AArch64ISD::VG_UNWIND", SDTypeProfile<0, 1, []>,
+                             [SDNPHasChain, SDNPSideEffect, SDNPOptInGlue, SDNPOutGlue]>;
+
 //===----------------------------------------------------------------------===//
 // Instruction naming conventions.
 //===----------------------------------------------------------------------===//
@@ -221,6 +224,15 @@ def : Pat<(AArch64_smstop (i32 svcr_op:$pstate), (i64 /*AArch64SME::Always*/0)),
           (MSRpstatesvcrImm1 svcr_op:$pstate, 0b0)>;
 
 
+// Pseudo to insert cfi_offset/cfi_restore instructions. Used to save or restore
+// the streaming value of VG around streaming-mode changes in locally-streaming
+// functions.
+def VGUnwindInfoPseudo :
+  Pseudo<(outs), (ins timm0_1:$save_restore), []>, Sched<[]>;
+
+def : Pat<(AArch64VGUnwind (i64 timm0_1:$save_restore)),
+          (VGUnwindInfoPseudo timm0_1:$save_restore)>;
+
 //===----------------------------------------------------------------------===//
 // SME2 Instructions
 //===----------------------------------------------------------------------===//

llvm/test/CodeGen/AArch64/sme-streaming-compatible-interface.ll

@@ -447,8 +447,8 @@ define void @disable_tailcallopt() "aarch64_pstate_sm_compatible" nounwind {
 define void @call_to_non_streaming_pass_args(ptr nocapture noundef readnone %ptr, i64 %long1, i64 %long2, i32 %int1, i32 %int2, float %float1, float %float2, double %double1, double %double2) "aarch64_pstate_sm_compatible" {
 ; CHECK-LABEL: call_to_non_streaming_pass_args:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    cntd x9
 ; CHECK-NEXT:    sub sp, sp, #128
+; CHECK-NEXT:    cntd x9
 ; CHECK-NEXT:    stp d15, d14, [sp, #32] // 16-byte Folded Spill
 ; CHECK-NEXT:    stp d13, d12, [sp, #48] // 16-byte Folded Spill
 ; CHECK-NEXT:    stp d11, d10, [sp, #64] // 16-byte Folded Spill