Fix incorrect choice of callee-saved registers save/restore points

[chill]

Summary

So, I have this testcase:

void f(int n, int x[]) {
  if (n < 0)
    return;

  int a[n];

  for (int i = 0; i < n; i++)
    a[i] = x[n - i - 1];

  for (int i = 0; i < n; i++)
    x[i] = a[i] + 1;
}

that, compiled with -O1/-Os for AArch64 and X86, generates machine code, which
fails to properly restore the stack pointer upon function return.

The testcase allocates a VLA, thus clang generates calls to llvm.stacksave and
llvm.stackrestore. The latter call is lowered to mov sp, x2, however this
move does not affect decisions made by the shrink wrapping pass, as the
instruction does not use or define a callee-saved register.

The end effect is that placement of register save/restore code is such that
along a certain path, the callee-saved registers and the stack pointer are
restored, but then the stack pointer is overwritten with an incorrect value.

This patches fixes the issue this by modifying ShrinkWrap::useOrDefCSROrFI to explicitly check
for the stack pointer register (using TLI.getStackPointerRegisterToSaveRestore) in non-call instructions.

[chill]

test/CodeGen/Generic/shrink-wrapping-vla.ll

88 ↗

(On Diff #142027)

I saw this one :/

[thegameg]

Thanks for this! If it's possible, could you also write a MIR test for this?

lib/CodeGen/ShrinkWrap.cpp

262 ↗

(On Diff #142027)

I think this is incorrect if the instruction calls a function with a different calling convention, right? In that case some CSRs would be clobbered while otherwise they were expected to be preserved.

IIUC, the issue here is that some instructions are marked as let Uses = [SP], like calls and returns, but are safe to skip during the analysis.

If I'm not missing anything, to handle all the cases I think it's better to verify all the register operands and reg masks, then only skip the instruction if it only uses SP and isCall or something similar. I'm not sure if isCall would handle all the cases, but from a quick glance over *InstrInfo.td, other instructions that have Uses = [SP] should definitely affect the placement of the save/restore blocks.

[junbuml]

lib/CodeGen/ShrinkWrap.cpp

265 ↗

(On Diff #142027)

Why don't we make it as a member variable?

[chill]

lib/CodeGen/ShrinkWrap.cpp

262 ↗

(On Diff #142027)

Thanks, indeed. For register masks, if I understand correctly they are used with call instructions and these should preserve SP, so no need to actually check it ?

265 ↗

(On Diff #142027)

Right, I made the stack pointer register number a member variable.

[thegameg]

Thanks. LGTM + one more comment.

lib/CodeGen/ShrinkWrap.cpp

271 ↗

(On Diff #142437)

Can you please add a comment explaining why we skip SP here?

[chill]

Added a comment; reformatted a little

[chill]

Thanks a lot for the comments and the review.

[thegameg]

lib/CodeGen/ShrinkWrap.cpp

274 ↗

(On Diff #142655)

Small typo: is -> as?

[thegameg]

Thanks!

[thegameg]

IIRC, there was a test case in the original patch but I can't see it in the landed commit. Any idea what happened?

[chill]

In D45524#1069899, @thegameg wrote:

IIRC, there was a test case in the original patch but I can't see it in the landed commit. Any idea what happened?

Oh, snap, I fail at svn ...

llvm/trunk/lib/CodeGen/ShrinkWrap.cpp

 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/RegisterClassInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
class ShrinkWrap : public MachineFunctionPass {
   uint64_t EntryFreq;
 
   /// Current opcode for frame setup.
   unsigned FrameSetupOpcode;
 
   /// Current opcode for frame destroy.
   unsigned FrameDestroyOpcode;
 
+  /// Stack pointer register, used by llvm.{savestack,restorestack}
+  unsigned SP;
+
   /// Entry block.
   const MachineBasicBlock *Entry;
 
   using SetOfRegs = SmallSetVector<unsigned, 16>;
 
   /// Registers that need to be saved for the current function.
   mutable SetOfRegs CurrentCSRs;
 
void init(MachineFunction &MF) {
     RCI.runOnMachineFunction(MF);
     MDT = &getAnalysis<MachineDominatorTree>();
     MPDT = &getAnalysis<MachinePostDominatorTree>();
     Save = nullptr;
     Restore = nullptr;
     MBFI = &getAnalysis<MachineBlockFrequencyInfo>();
     MLI = &getAnalysis<MachineLoopInfo>();
     EntryFreq = MBFI->getEntryFreq();
-    const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+    const TargetSubtargetInfo &Subtarget = MF.getSubtarget();
+    const TargetInstrInfo &TII = *Subtarget.getInstrInfo();
     FrameSetupOpcode = TII.getCallFrameSetupOpcode();
     FrameDestroyOpcode = TII.getCallFrameDestroyOpcode();
+    SP = Subtarget.getTargetLowering()->getStackPointerRegisterToSaveRestore();
     Entry = &MF.front();
     CurrentCSRs.clear();
     MachineFunc = &MF;
 
     ++NumFunc;
   }
 
   /// Check whether or not Save and Restore points are still interesting for
if (MO.isReg()) {
       // Ignore instructions like DBG_VALUE which don't read/def the register.
       if (!MO.isDef() && !MO.readsReg())
         continue;
       unsigned PhysReg = MO.getReg();
       if (!PhysReg)
         continue;
       assert(TargetRegisterInfo::isPhysicalRegister(PhysReg) &&
              "Unallocated register?!");
-      UseOrDefCSR = RCI.getLastCalleeSavedAlias(PhysReg);
+      // The stack pointer is not normally described as a callee-saved register
+      // in calling convention definitions, so we need to watch for it
+      // separately. An SP mentioned by a call instruction, we can ignore,
+      // though, as it's harmless and we do not want to effectively disable tail
+      // calls by forcing the restore point to post-dominate them.
+      UseOrDefCSR = (!MI.isCall() && PhysReg == SP) ||
+                    RCI.getLastCalleeSavedAlias(PhysReg);
     } else if (MO.isRegMask()) {
       // Check if this regmask clobbers any of the CSRs.
       for (unsigned Reg : getCurrentCSRs(RS)) {
         if (MO.clobbersPhysReg(Reg)) {
           UseOrDefCSR = true;
           break;
         }
       }