[Hexagon] Adjust handling of stack with variable-size and extra alignment

Make the stack alignment register (AP) reserved in the given function. This
will make it available everywhere in the function, and allow aligned access
to vector register spill slots.

Author: Krzysztof Parzyszek

llvm/lib/Target/Hexagon/HexagonFrameLowering.cpp

@@ -1275,7 +1275,7 @@ HexagonFrameLowering::getFrameIndexReference(const MachineFunction &MF, int FI,
   unsigned FrameSize = MFI.getStackSize();
   Register SP = HRI.getStackRegister();
   Register FP = HRI.getFrameRegister();
-  Register AP = HMFI.getStackAlignBasePhysReg();
+  Register AP = HMFI.getStackAlignBaseReg();
   // It may happen that AP will be absent even HasAlloca && HasExtraAlign
   // is true. HasExtraAlign may be set because of vector spills, without
   // aligned locals or aligned outgoing function arguments. Since vector
@@ -1290,9 +1290,6 @@ HexagonFrameLowering::getFrameIndexReference(const MachineFunction &MF, int FI,
   // assume that missing AP will be replaced by FP.
   // (A better fix would be to rematerialize AP from FP and always align
   // vector spills.)
-  if (AP == 0)
-    AP = FP;
-
   bool UseFP = false, UseAP = false;  // Default: use SP (except at -O0).
   // Use FP at -O0, except when there are objects with extra alignment.
   // That additional alignment requirement may cause a pad to be inserted,
@@ -1517,65 +1514,13 @@ void HexagonFrameLowering::processFunctionBeforeFrameFinalized(
   if (!HasAlloca || !NeedsAlign)
     return;
 
-  SmallSet<int, 4> DealignSlots;
-  unsigned LFS = MFI.getLocalFrameSize();
-  for (int i = 0, e = MFI.getObjectIndexEnd(); i != e; ++i) {
-    if (!MFI.isSpillSlotObjectIndex(i) || MFI.isDeadObjectIndex(i))
-      continue;
-    unsigned S = MFI.getObjectSize(i);
-    // Reduce the alignment to at most 8. This will require unaligned vector
-    // stores if they happen here.
-    Align A = std::max(MFI.getObjectAlign(i), Align(8));
-    MFI.setObjectAlignment(i, Align(8));
-    LFS = alignTo(LFS+S, A);
-    MFI.mapLocalFrameObject(i, -static_cast<int64_t>(LFS));
-    DealignSlots.insert(i);
-  }
-
-  MFI.setLocalFrameSize(LFS);
-  Align A = MFI.getLocalFrameMaxAlign();
-  assert(A <= 8 && "Unexpected local frame alignment");
-  if (A == 1)
-    MFI.setLocalFrameMaxAlign(Align(8));
-  MFI.setUseLocalStackAllocationBlock(true);
-
-  // Go over all MachineMemOperands in the code, and change the ones that
-  // refer to the dealigned stack slots to reflect the new alignment.
-  if (!DealignSlots.empty()) {
-    for (MachineBasicBlock &BB : MF) {
-      for (MachineInstr &MI : BB) {
-        bool KeepOld = true;
-        ArrayRef<MachineMemOperand*> memops = MI.memoperands();
-        SmallVector<MachineMemOperand*,1> new_memops;
-        for (MachineMemOperand *MMO : memops) {
-          auto *PV = MMO->getPseudoValue();
-          if (auto *FS = dyn_cast_or_null<FixedStackPseudoSourceValue>(PV)) {
-            int FI = FS->getFrameIndex();
-            if (DealignSlots.count(FI)) {
-              auto *NewMMO = MF.getMachineMemOperand(
-                  MMO->getPointerInfo(), MMO->getFlags(), MMO->getSize(),
-                  MFI.getObjectAlign(FI), MMO->getAAInfo(), MMO->getRanges(),
-                  MMO->getSyncScopeID(), MMO->getSuccessOrdering(),
-                  MMO->getFailureOrdering());
-              new_memops.push_back(NewMMO);
-              KeepOld = false;
-              continue;
-            }
-          }
-          new_memops.push_back(MMO);
-        }
-        if (!KeepOld)
-          MI.setMemRefs(MF, new_memops);
-      }
-    }
-  }
-
   // Set the physical aligned-stack base address register.
   Register AP = 0;
   if (const MachineInstr *AI = getAlignaInstr(MF))
     AP = AI->getOperand(0).getReg();
   auto &HMFI = *MF.getInfo<HexagonMachineFunctionInfo>();
-  HMFI.setStackAlignBasePhysReg(AP);
+  assert(!AP.isValid() || AP.isPhysical());
+  HMFI.setStackAlignBaseReg(AP);
 }
 
 /// Returns true if there are no caller-saved registers available in class RC.
@@ -1637,6 +1582,26 @@ bool HexagonFrameLowering::assignCalleeSavedSpillSlots(MachineFunction &MF,
   // (2) For each reserved register, remove that register and all of its
   // sub- and super-registers from SRegs.
   BitVector Reserved = TRI->getReservedRegs(MF);
+  // Unreserve the stack align register: it is reserved for this function
+  // only, it still needs to be saved/restored.
+  Register AP =
+      MF.getInfo<HexagonMachineFunctionInfo>()->getStackAlignBaseReg();
+  if (AP.isValid()) {
+    Reserved[AP] = false;
+    // Unreserve super-regs if no other subregisters are reserved.
+    for (MCSuperRegIterator SP(AP, TRI, false); SP.isValid(); ++SP) {
+      bool HasResSub = false;
+      for (MCSubRegIterator SB(*SP, TRI, false); SB.isValid(); ++SB) {
+        if (!Reserved[*SB])
+          continue;
+        HasResSub = true;
+        break;
+      }
+      if (!HasResSub)
+        Reserved[*SP] = false;
+    }
+  }
+
   for (int x = Reserved.find_first(); x >= 0; x = Reserved.find_next(x)) {
     Register R = x;
     for (MCSuperRegIterator SR(R, TRI, true); SR.isValid(); ++SR)
@@ -1936,21 +1901,16 @@ bool HexagonFrameLowering::expandStoreVec2(MachineBasicBlock &B,
   Register SrcHi = HRI.getSubReg(SrcR, Hexagon::vsub_hi);
   bool IsKill = MI->getOperand(2).isKill();
   int FI = MI->getOperand(0).getIndex();
-  bool NeedsAligna = needsAligna(MF);
 
   unsigned Size = HRI.getSpillSize(Hexagon::HvxVRRegClass);
   Align NeedAlign = HRI.getSpillAlign(Hexagon::HvxVRRegClass);
   Align HasAlign = MFI.getObjectAlign(FI);
   unsigned StoreOpc;
 
-  auto UseAligned = [&](Align NeedAlign, Align HasAlign) {
-    return !NeedsAligna && (NeedAlign <= HasAlign);
-  };
-
   // Store low part.
   if (LPR.contains(SrcLo)) {
-    StoreOpc = UseAligned(NeedAlign, HasAlign) ? Hexagon::V6_vS32b_ai
-                                               : Hexagon::V6_vS32Ub_ai;
+    StoreOpc = NeedAlign <= HasAlign ? Hexagon::V6_vS32b_ai
+                                     : Hexagon::V6_vS32Ub_ai;
     BuildMI(B, It, DL, HII.get(StoreOpc))
         .addFrameIndex(FI)
         .addImm(0)
@@ -1960,8 +1920,8 @@ bool HexagonFrameLowering::expandStoreVec2(MachineBasicBlock &B,
 
   // Store high part.
   if (LPR.contains(SrcHi)) {
-    StoreOpc = UseAligned(NeedAlign, HasAlign) ? Hexagon::V6_vS32b_ai
-                                               : Hexagon::V6_vS32Ub_ai;
+    StoreOpc = NeedAlign <= HasAlign ? Hexagon::V6_vS32b_ai
+                                     : Hexagon::V6_vS32Ub_ai;
     BuildMI(B, It, DL, HII.get(StoreOpc))
         .addFrameIndex(FI)
         .addImm(Size)
@@ -1988,28 +1948,23 @@ bool HexagonFrameLowering::expandLoadVec2(MachineBasicBlock &B,
   Register DstHi = HRI.getSubReg(DstR, Hexagon::vsub_hi);
   Register DstLo = HRI.getSubReg(DstR, Hexagon::vsub_lo);
   int FI = MI->getOperand(1).getIndex();
-  bool NeedsAligna = needsAligna(MF);
 
   unsigned Size = HRI.getSpillSize(Hexagon::HvxVRRegClass);
   Align NeedAlign = HRI.getSpillAlign(Hexagon::HvxVRRegClass);
   Align HasAlign = MFI.getObjectAlign(FI);
   unsigned LoadOpc;
 
-  auto UseAligned = [&](Align NeedAlign, Align HasAlign) {
-    return !NeedsAligna && (NeedAlign <= HasAlign);
-  };
-
   // Load low part.
-  LoadOpc = UseAligned(NeedAlign, HasAlign) ? Hexagon::V6_vL32b_ai
-                                            : Hexagon::V6_vL32Ub_ai;
+  LoadOpc = NeedAlign <= HasAlign ? Hexagon::V6_vL32b_ai
+                                  : Hexagon::V6_vL32Ub_ai;
   BuildMI(B, It, DL, HII.get(LoadOpc), DstLo)
       .addFrameIndex(FI)
       .addImm(0)
       .cloneMemRefs(*MI);
 
   // Load high part.
-  LoadOpc = UseAligned(NeedAlign, HasAlign) ? Hexagon::V6_vL32b_ai
-                                            : Hexagon::V6_vL32Ub_ai;
+  LoadOpc = NeedAlign <= HasAlign ? Hexagon::V6_vL32b_ai
+                                  : Hexagon::V6_vL32Ub_ai;
   BuildMI(B, It, DL, HII.get(LoadOpc), DstHi)
       .addFrameIndex(FI)
       .addImm(Size)
@@ -2028,7 +1983,6 @@ bool HexagonFrameLowering::expandStoreVec(MachineBasicBlock &B,
   if (!MI->getOperand(0).isFI())
     return false;
 
-  bool NeedsAligna = needsAligna(MF);
   auto &HRI = *MF.getSubtarget<HexagonSubtarget>().getRegisterInfo();
   DebugLoc DL = MI->getDebugLoc();
   Register SrcR = MI->getOperand(2).getReg();
@@ -2037,9 +1991,8 @@ bool HexagonFrameLowering::expandStoreVec(MachineBasicBlock &B,
 
   Align NeedAlign = HRI.getSpillAlign(Hexagon::HvxVRRegClass);
   Align HasAlign = MFI.getObjectAlign(FI);
-  bool UseAligned = !NeedsAligna && (NeedAlign <= HasAlign);
-  unsigned StoreOpc = UseAligned ? Hexagon::V6_vS32b_ai
-                                 : Hexagon::V6_vS32Ub_ai;
+  unsigned StoreOpc = NeedAlign <= HasAlign ? Hexagon::V6_vS32b_ai
+                                            : Hexagon::V6_vS32Ub_ai;
   BuildMI(B, It, DL, HII.get(StoreOpc))
       .addFrameIndex(FI)
       .addImm(0)
@@ -2059,17 +2012,15 @@ bool HexagonFrameLowering::expandLoadVec(MachineBasicBlock &B,
   if (!MI->getOperand(1).isFI())
     return false;
 
-  bool NeedsAligna = needsAligna(MF);
   auto &HRI = *MF.getSubtarget<HexagonSubtarget>().getRegisterInfo();
   DebugLoc DL = MI->getDebugLoc();
   Register DstR = MI->getOperand(0).getReg();
   int FI = MI->getOperand(1).getIndex();
 
   Align NeedAlign = HRI.getSpillAlign(Hexagon::HvxVRRegClass);
   Align HasAlign = MFI.getObjectAlign(FI);
-  bool UseAligned = !NeedsAligna && (NeedAlign <= HasAlign);
-  unsigned LoadOpc = UseAligned ? Hexagon::V6_vL32b_ai
-                                : Hexagon::V6_vL32Ub_ai;
+  unsigned LoadOpc = NeedAlign <= HasAlign ? Hexagon::V6_vL32b_ai
+                                           : Hexagon::V6_vL32Ub_ai;
   BuildMI(B, It, DL, HII.get(LoadOpc), DstR)
       .addFrameIndex(FI)
       .addImm(0)

llvm/lib/Target/Hexagon/HexagonISelDAGToDAG.cpp

@@ -749,7 +749,7 @@ void HexagonDAGToDAGISel::SelectFrameIndex(SDNode *N) {
     R = CurDAG->getMachineNode(Hexagon::PS_fi, DL, MVT::i32, FI, Zero);
   } else {
     auto &HMFI = *MF->getInfo<HexagonMachineFunctionInfo>();
-    unsigned AR = HMFI.getStackAlignBaseVReg();
+    Register AR = HMFI.getStackAlignBaseReg();
     SDValue CH = CurDAG->getEntryNode();
     SDValue Ops[] = { CurDAG->getCopyFromReg(CH, DL, AR, MVT::i32), FI, Zero };
     R = CurDAG->getMachineNode(Hexagon::PS_fia, DL, MVT::i32, Ops);
@@ -1285,11 +1285,22 @@ void HexagonDAGToDAGISel::emitFunctionEntryCode() {
 
   MachineFrameInfo &MFI = MF->getFrameInfo();
   MachineBasicBlock *EntryBB = &MF->front();
-  Register AR = FuncInfo->CreateReg(MVT::i32);
   Align EntryMaxA = MFI.getMaxAlign();
-  BuildMI(EntryBB, DebugLoc(), HII->get(Hexagon::PS_aligna), AR)
+
+  // Reserve the first non-volatile register.
+  Register AP = 0;
+  auto &HRI = *HST.getRegisterInfo();
+  BitVector Reserved = HRI.getReservedRegs(*MF);
+  for (const MCPhysReg *R = HRI.getCalleeSavedRegs(MF); *R; ++R) {
+    if (Reserved[*R])
+      continue;
+    AP = *R;
+    break;
+  }
+  assert(AP.isValid() && "Couldn't reserve stack align register");
+  BuildMI(EntryBB, DebugLoc(), HII->get(Hexagon::PS_aligna), AP)
       .addImm(EntryMaxA.value());
-  MF->getInfo<HexagonMachineFunctionInfo>()->setStackAlignBaseVReg(AR);
+  MF->getInfo<HexagonMachineFunctionInfo>()->setStackAlignBaseReg(AP);
 }
 
 void HexagonDAGToDAGISel::updateAligna() {

llvm/lib/Target/Hexagon/HexagonMachineFunctionInfo.h

@@ -27,8 +27,7 @@ class HexagonMachineFunctionInfo : public MachineFunctionInfo {
   // returning the value of the returned struct in a register. This field
   // holds the virtual register into which the sret argument is passed.
   unsigned SRetReturnReg = 0;
-  unsigned StackAlignBaseVReg = 0;    // Aligned-stack base register (virtual)
-  unsigned StackAlignBasePhysReg = 0; //                             (physical)
+  Register StackAlignBaseReg = 0;    // Aligned-stack base register
   int VarArgsFrameIndex;
   int RegSavedAreaStartFrameIndex;
   int FirstNamedArgFrameIndex;
@@ -82,11 +81,8 @@ class HexagonMachineFunctionInfo : public MachineFunctionInfo {
   bool hasEHReturn() const { return HasEHReturn; };
   void setHasEHReturn(bool H = true) { HasEHReturn = H; };
 
-  void setStackAlignBaseVReg(unsigned R) { StackAlignBaseVReg = R; }
-  unsigned getStackAlignBaseVReg() const { return StackAlignBaseVReg; }
-
-  void setStackAlignBasePhysReg(unsigned R) { StackAlignBasePhysReg = R; }
-  unsigned getStackAlignBasePhysReg() const { return StackAlignBasePhysReg; }
+  void setStackAlignBaseReg(Register R) { StackAlignBaseReg = R; }
+  Register getStackAlignBaseReg() const { return StackAlignBaseReg; }
 };
 
 } // end namespace llvm

llvm/lib/Target/Hexagon/HexagonRegisterInfo.cpp

@@ -194,6 +194,11 @@ BitVector HexagonRegisterInfo::getReservedRegs(const MachineFunction &MF)
   if (MF.getSubtarget<HexagonSubtarget>().hasReservedR19())
     Reserved.set(Hexagon::R19);
 
+  Register AP =
+      MF.getInfo<HexagonMachineFunctionInfo>()->getStackAlignBaseReg();
+  if (AP.isValid())
+    Reserved.set(AP);
+
   for (int x = Reserved.find_first(); x >= 0; x = Reserved.find_next(x))
     markSuperRegs(Reserved, x);
 

llvm/lib/Target/Hexagon/HexagonVExtract.cpp

@@ -105,7 +105,7 @@ bool HexagonVExtract::runOnMachineFunction(MachineFunction &MF) {
   MachineRegisterInfo &MRI = MF.getRegInfo();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   Register AR =
-      MF.getInfo<HexagonMachineFunctionInfo>()->getStackAlignBaseVReg();
+      MF.getInfo<HexagonMachineFunctionInfo>()->getStackAlignBaseReg();
   std::map<unsigned, SmallVector<MachineInstr *, 4>> VExtractMap;
   bool Changed = false;
 

llvm/test/CodeGen/Hexagon/spill-vector-alignment.mir

@@ -1,7 +1,7 @@
 # RUN: llc -march=hexagon -run-pass prologepilog %s -o - | FileCheck %s
 
-# Check that the spill of $q0 uses unaligned store instruction.
-# CHECK: V6_vS32Ub_ai $r30, -128, killed $v0
+# Check that the spill of $q0 no longer uses unaligned store instruction.
+# CHECK: V6_vS32b_ai $r16, -256, killed $v0
 
 ---
 name:            test
@@ -12,5 +12,6 @@ stack:
 body: |
   bb.0:
     liveins: $q0
+    $r16 = PS_aligna 128, implicit $r30
     PS_vstorerq_ai %stack.1, 0, $q0
 ...

llvm/test/CodeGen/Hexagon/stack-alloca2.ll

@@ -1,9 +1,9 @@
 ; RUN: llc -O0 -march=hexagon < %s | FileCheck %s
-; CHECK-DAG: r[[AP:[0-9]+]] = and(r30,#-32)
-; CHECK-DAG: r1 = add(r[[AP]],#-32)
-
-; CHECK-DAG: sub(r29,r[[SP:[0-9]+]])
-; CHECK-DAG: r29 = r[[SP]]
+; CHECK: r[[AP:[0-9]+]] = and(r30,#-32)
+; CHECK: sub(r29,r[[SP:[0-9]+]])
+; CHECK: r29 = r[[SP]]
+; CHECK: r1 = r[[AP]]
+; CHECK: r1 = add(r1,#-32)
 
 target triple = "hexagon-unknown-unknown"
 

llvm/test/CodeGen/Hexagon/v6-unaligned-spill.ll

@@ -1,10 +1,9 @@
 ; RUN: llc -march=hexagon < %s | FileCheck %s
 
-; Test that we generate an unaligned vector store for a spill when a function
-; has an alloca. Also, make sure the addressing mode for unaligned store does
-; is not a base+offset with a non-zero offset that is not a multiple of 128.
+; Test that we no longer generate an unaligned vector store for a spill when
+; a function has an alloca.
 
-; CHECK: vmemu({{.*}}) =
+; CHECK: vmem({{.*}}) =
 
 %s.0 = type { [5 x [4 x i8]], i32, i32, i32, i32 }