[BOLT] Increase coverage of shrink wrapping [5/5]

[rafauler]

Summary

Add -experimental-shrink-wrapping flag to control when we
want to move callee-saved registers even when addresses of the stack
frame are captured and used in pointer arithmetic, making it more
challenging to do alias analysis to prove that we do not access
optimized stack positions. This alias analysis is not yet implemented,
hence, it is experimental. In practice, though, no compiler would emit
code to do pointer arithmetic to access a saved callee-saved register
unless there is a memory bug or we are failing to identify a
callee-saved reg, so I'm not sure how useful it would be to formally
prove that.

[Amir]

To recap:

1. This diff adds a new experimental shrink wrapping mode which ignores escaping stack pointer (ie. mov rsp -> reg) in moving callee-saved registers.
2. This is based on the assumption that stack frame locations for CSRs are not explicitly addressed by the function code.
3. For functions that have been optimized under this mode, the new property (FAA.doesEscapeStackAddress) is added to prevent further Frame optimizations.

bolt/include/bolt/Passes/FrameAnalysis.h

class FrameAnalysis {
 
   /// The set of functions we were able to perform the full analysis up to
   /// restoring frame indexes for all load/store instructions.
   DenseSet<const BinaryFunction *> AnalyzedFunctions;
 
   /// Set of functions that require the stack to be 16B aligned
   DenseSet<const BinaryFunction *> FunctionsRequireAlignment;
 
+  /// Set of functions that performs computations with stack addresses and
+  /// complicates our understanding of aliasing of stack spaces.
+  DenseSet<const BinaryFunction *> FunctionsWithStackArithmetic;
+
   /// Owns ArgAccesses for all instructions. References to elements are
   /// attached to instructions as indexes to this vector, in MCAnnotations.
   std::vector<ArgAccesses> ArgAccessesVector;
   /// Same for FrameIndexEntries.
   std::vector<FrameIndexEntry> FIEVector;
 
   /// Analysis stats counters
   uint64_t NumFunctionsNotOptimized{0};
bool hasFrameInfo(const BinaryFunction &Func) const {
     return AnalyzedFunctions.count(&Func);
   }
 
   /// Return true if \p Func cannot operate with a misaligned CFA
   bool requiresAlignment(const BinaryFunction &Func) const {
     return FunctionsRequireAlignment.count(&Func);
   }
 
+  /// Return true if \p Func does computation with the address of any stack
+  /// position, meaning we have limited alias analysis on this function.
+  bool hasStackArithmetic(const BinaryFunction &Func) const {
+    return FunctionsWithStackArithmetic.count(&Func);
+  }
+
   /// Functions for retrieving our specific MCAnnotation data from instructions
   ErrorOr<ArgAccesses &> getArgAccessesFor(const MCInst &Inst);
 
   ErrorOr<const ArgAccesses &> getArgAccessesFor(const MCInst &Inst) const;
 
   ErrorOr<const FrameIndexEntry &> getFIEFor(const MCInst &Inst) const;
 
   /// Remove all MCAnnotations attached by this pass

bolt/lib/Passes/FrameAnalysis.cpp

 
 static cl::opt<std::string> FrameOptFunctionNamesFile(
     "funcs-file-fop",
     cl::desc("file with list of functions to frame optimize"));
 
 static cl::opt<bool> TimeFA("time-fa", cl::desc("time frame analysis steps"),
                             cl::ReallyHidden, cl::cat(BoltOptCategory));
 
+static cl::opt<bool>
+    ExperimentalSW("experimental-shrink-wrapping",
+                   cl::desc("process functions with stack pointer arithmetic"),
+                   cl::ReallyHidden, cl::ZeroOrMore, cl::cat(BoltOptCategory));
+
 bool shouldFrameOptimize(const llvm::bolt::BinaryFunction &Function) {
   if (Function.hasUnknownControlFlow())
     return false;
 
   if (!FrameOptFunctionNamesFile.empty()) {
     assert(!FrameOptFunctionNamesFile.empty() && "unexpected empty file name");
     std::ifstream FuncsFile(FrameOptFunctionNamesFile, std::ios::in);
     std::string FuncName;
class FrameAccessAnalysis {
   int FPOffset{0};
   int64_t CfaOffset{-8};
   uint16_t CfaReg{7};
   std::stack<std::pair<int64_t, uint16_t>> CFIStack;
   /// Our pointer to access SPT info
   const MCInst *Prev{nullptr};
   /// Info about the last frame access
   bool IsValidAccess{false};
+  bool EscapesStackAddress{false};
   FrameIndexEntry FIE;
 
   bool decodeFrameAccess(const MCInst &Inst) {
     int32_t SrcImm = 0;
     MCPhysReg Reg = 0;
     int64_t StackOffset = 0;
     bool IsIndexed = false;
     if (!BC.MIB->isStackAccess(
             Inst, FIE.IsLoad, FIE.IsStore, FIE.IsStoreFromReg, Reg, SrcImm,
             FIE.StackPtrReg, StackOffset, FIE.Size, FIE.IsSimple, IsIndexed)) {
       return true;
     }
 
-    if (IsIndexed || FIE.Size == 0) {
+    if (IsIndexed || (!FIE.Size && (FIE.IsLoad || FIE.IsStore))) {
       LLVM_DEBUG(dbgs() << "Giving up on indexed memory access/unknown size\n");
       LLVM_DEBUG(dbgs() << "Blame insn: ");
+      LLVM_DEBUG(BC.printInstruction(outs(), Inst, 0, &BF, true, false, false));
       LLVM_DEBUG(Inst.dump());
       return false;
     }
 
-    assert(FIE.Size != 0);
+    assert(FIE.Size != 0 || (!FIE.IsLoad && !FIE.IsStore));
 
     FIE.RegOrImm = SrcImm;
     if (FIE.IsLoad || FIE.IsStoreFromReg)
       FIE.RegOrImm = Reg;
 
     if (FIE.StackPtrReg == BC.MIB->getStackPointer() && SPOffset != SPT.EMPTY &&
         SPOffset != SPT.SUPERPOSITION) {
       LLVM_DEBUG(

 public:
   FrameAccessAnalysis(BinaryFunction &BF, StackPointerTracking &SPT)
       : SPT(SPT), BC(BF.getBinaryContext()), BF(BF) {}
 
   void enterNewBB() { Prev = nullptr; }
   const FrameIndexEntry &getFIE() const { return FIE; }
   int getSPOffset() const { return SPOffset; }
   bool isValidAccess() const { return IsValidAccess; }
+  bool doesEscapeStackAddress() const { return EscapesStackAddress; }
 
   bool doNext(const BinaryBasicBlock &BB, const MCInst &Inst) {
     IsValidAccess = false;
+    EscapesStackAddress = false;
     std::tie(SPOffset, FPOffset) =
         Prev ? *SPT.getStateAt(*Prev) : *SPT.getStateAt(BB);
     Prev = &Inst;
     // Use CFI information to keep track of which register is being used to
     // access the frame
     if (BC.MIB->isCFI(Inst)) {
       const MCCFIInstruction *CFI = BF.getCFIFor(Inst);
       switch (CFI->getOperation()) {
if (BC.MIB->isCFI(Inst)) {
         break;
       default:
         break;
       }
       return true;
     }
 
     if (BC.MIB->escapesVariable(Inst, SPT.HasFramePointer)) {
+      EscapesStackAddress = true;
+      if (!opts::ExperimentalSW) {
         LLVM_DEBUG(
             dbgs() << "Leaked stack address, giving up on this function.\n");
         LLVM_DEBUG(dbgs() << "Blame insn: ");
         LLVM_DEBUG(Inst.dump());
         return false;
       }
+    }
 
     return decodeFrameAccess(Inst);
   }
 };
 
 } // end anonymous namespace
 
 void FrameAnalysis::addArgAccessesFor(MCInst &Inst, ArgAccesses &&AA) {
bool FrameAnalysis::computeArgsAccessed(BinaryFunction &BF) {
   bool UpdatedArgsTouched = false;
   bool NoInfo = false;
   FrameAccessAnalysis FAA(BF, getSPT(BF));
 
   for (BinaryBasicBlock *BB : BF.layout()) {
     FAA.enterNewBB();
 
     for (MCInst &Inst : *BB) {
-      if (!FAA.doNext(*BB, Inst)) {
+      if (!FAA.doNext(*BB, Inst) || FAA.doesEscapeStackAddress()) {
         ArgsTouchedMap[&BF].emplace(std::make_pair(-1, 0));
         NoInfo = true;
         break;
       }
 
       // Check for calls -- attach stack accessing info to them regarding their
       // target
       if (updateArgsTouchedFor(BF, Inst, FAA.getSPOffset()))
for (MCInst &Inst : *BB) {
       if (!FAA.doNext(*BB, Inst))
         return false;
       LLVM_DEBUG({
         dbgs() << "\t\tNow at ";
         Inst.dump();
         dbgs() << "\t\t\tSP offset is " << FAA.getSPOffset() << "\n";
       });
 
+      if (FAA.doesEscapeStackAddress()) {
+        if (!FunctionsWithStackArithmetic.count(&BF))
+          FunctionsWithStackArithmetic.insert(&BF);
+        continue;
+      }
+
       if (!FAA.isValidAccess())
         continue;
 
       const FrameIndexEntry &FIE = FAA.getFIE();
 
       addFIEFor(Inst, FIE);
       LLVM_DEBUG({
         dbgs() << "Frame index annotation " << FIE << " added to:\n";

bolt/lib/Passes/FrameOptimizer.cpp

if (opts::FrameOptimization == FOP_HOT) {
                  << I.second.getKnownExecutionCount()
                  << " ) exceeds our hotness threshold ( "
                  << BC.getHotThreshold() << " )\n");
     }
 
     {
       NamedRegionTimer T1("removeloads", "remove loads", "FOP", "FOP breakdown",
                           opts::TimeOpts);
+      if (!FA->hasStackArithmetic(I.second))
         removeUnnecessaryLoads(*RA, *FA, I.second);
     }
 
     if (opts::RemoveStores) {
       NamedRegionTimer T1("removestores", "remove stores", "FOP",
                           "FOP breakdown", opts::TimeOpts);
+      if (!FA->hasStackArithmetic(I.second))
         removeUnusedStores(*FA, I.second);
     }
     // Don't even start shrink wrapping if no profiling info is available
     if (I.second.getKnownExecutionCount() == 0)
       continue;
   }
 
   {
     NamedRegionTimer T1("shrinkwrapping", "shrink wrapping", "FOP",

bolt/lib/Passes/ShrinkWrapping.cpp

bool ShrinkWrapping::validatePushPopsMode(unsigned CSR, MCInst *BestPosSave,
   if (FA.requiresAlignment(BF)) {
     LLVM_DEBUG({
       dbgs() << "Reg " << CSR
              << " is not using push/pops due to function "
                 "alignment requirements.\n";
     });
     return false;
   }
+  if (FA.hasStackArithmetic(BF)) {
+    LLVM_DEBUG({
+      dbgs() << "Reg " << CSR
+             << " is not using push/pops due to function "
+                "taking the address of a stack position.\n";
+    });
+    return false;
+  }
   for (MCInst *Save : CSA.getSavesByReg(CSR)) {
     if (!SLM.canCollapseRegion(Save)) {
       LLVM_DEBUG(dbgs() << "Reg " << CSR << " cannot collapse region.\n");
       return false;
     }
   }
   // Abort if one of the restores for this CSR is not a POP.
   for (MCInst *Load : CSA.getRestoresByReg(CSR)) {

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

bool isStackAccess(const MCInst &Inst, bool &IsLoad, bool &IsStore,
     if (MemOpNo == -1 || MCII.isCall())
       return false;
 
     switch (Inst.getOpcode()) {
     default: {
       bool IsLoad = MCII.mayLoad();
       bool IsStore = MCII.mayStore();
       // Is it LEA? (deals with memory but is not loading nor storing)
-      if (!IsLoad && !IsStore)
-        return false;
+      if (!IsLoad && !IsStore) {
+        I = {0, IsLoad, IsStore, false, false};
+        break;
+      }
       uint8_t Sz = getMemDataSize(Inst, MemOpNo);
       I = {Sz, IsLoad, IsStore, false, false};
       break;
     }
     // Report simple stack accesses
     case X86::MOV8rm: I = {1, true, false, false, true}; break;
     case X86::MOV16rm: I = {2, true, false, false, true}; break;
     case X86::MOV32rm: I = {4, true, false, false, true}; break;

bolt/test/X86/frame-opt-lea.s

+# This checks that frame optimizer does not try to optimize away caller-saved
+# regs when we do not have complete aliasing info (when there is an LEA
+# instruction and the function does arithmetic with stack addresses).
+
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt %t.exe -relocs -o %t.out -data %t.fdata \
+# RUN:     -frame-opt=all -simplify-conditional-tail-calls=false \
+# RUN:     -lite=0 -eliminate-unreachable=false | FileCheck %s
+# RUN: llvm-objdump -d %t.out --print-imm-hex | \
+# RUN:   FileCheck --check-prefix CHECK-OBJDUMP %s
+
+  .globl foo
+  .type foo, %function
+foo:
+  .cfi_startproc
+  movq $0, (%rsi)
+  ret
+  .cfi_endproc
+  .size foo, .-foo
+
+
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 _start 0 0 1
+  push  %rbp
+  mov   %rsp, %rbp
+  subq  $0x20, %rsp
+  je  b
+c:
+  addq  $0x20, %rsp
+  pop %rbp
+  ret
+b:
+  movq %rdi, %r13
+  movq %r13, -0x08(%rbp)
+  leaq -0x08(%rbp), %rsi
+  callq foo
+  movq -0x08(%rbp), %r13
+  jmp c
+  .cfi_endproc
+  .size _start, .-_start
+
+
+# CHECK:   BOLT-INFO: FOP deleted 0 load(s) (dyn count: 0) and 0 store(s)
+
+# CHECK-OBJDUMP:     <_start>:
+# CHECK-OBJDUMP:         movq    %rdi, %r13
+# CHECK-OBJDUMP-NEXT:    movq    %r13, -0x8(%rbp)
+# CHECK-OBJDUMP-NEXT:    leaq
+# CHECK-OBJDUMP-NEXT:    callq
+# CHECK-OBJDUMP-NEXT:    movq    -0x8(%rbp), %r13

bolt/test/X86/shrinkwrapping-lea.s

+# This checks that shrink wrapping correctly drops moving push/pops when
+# there is an LEA instruction.
+
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: ld.lld %t.o -o %t.exe -q
+# RUN: llvm-bolt %t.exe -relocs -o %t.out -data %t.fdata \
+# RUN:     -frame-opt=all -simplify-conditional-tail-calls=false \
+# RUN:     -experimental-shrink-wrapping \
+# RUN:     -eliminate-unreachable=false | FileCheck %s
+# RUN: llvm-objdump -d %t.out --print-imm-hex | \
+# RUN:   FileCheck --check-prefix CHECK-OBJDUMP %s
+
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 _start 0 0 1
+  push  %rbp
+  mov   %rsp, %rbp
+  push  %rbx
+  push  %r14
+  subq  $0x20, %rsp
+  je  b
+c:
+  addq  $0x20, %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+b:
+  je  f
+  jmp *JT(,%rdi,8)
+d:
+  mov %r14, %rdi
+  mov %rbx, %rdi
+  leaq -0x20(%rbp), %r14
+  movq -0x20(%rbp), %rdi
+f:
+  addq  $0x20, %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+  .cfi_endproc
+  .size _start, .-_start
+  .data
+JT:
+  .quad c
+  .quad d
+  .quad f
+
+
+# CHECK:   BOLT-INFO: Shrink wrapping moved 2 spills inserting load/stores and 0 spills inserting push/pops
+
+# Checks that offsets of instructions accessing the stack were not changed
+# CHECK-OBJDUMP:     <_start>:
+# CHECK-OBJDUMP:         movq    %rbx, %rdi
+# CHECK-OBJDUMP-NEXT:    leaq    -0x20(%rbp), %r14
+# CHECK-OBJDUMP:         movq    -0x20(%rbp), %rdi

bolt/test/X86/shrinkwrapping-mov.s

+# This checks that shrink wrapping correctly drops moving push/pops when
+# there is a MOV instruction loading the value of the stack pointer in
+# order to do pointer arithmetic with a stack address.
+
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: ld.lld %t.o -o %t.exe -q
+# RUN: llvm-bolt %t.exe -relocs -o %t.out -data %t.fdata \
+# RUN:     -frame-opt=all -simplify-conditional-tail-calls=false \
+# RUN:     -experimental-shrink-wrapping \
+# RUN:     -eliminate-unreachable=false | FileCheck %s
+
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 _start 0 0 1
+  push  %rbp
+  mov   %rsp, %rbp
+  push  %rbx
+  push  %r14
+  subq  $0x20, %rsp
+  je  b
+c:
+  addq  $0x20, %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+b:
+  je  f
+  jmp *JT(,%rdi,8)
+d:
+  mov %r14, %rdi
+  mov %rbx, %rdi
+  mov %rbp, %rdi
+  sub 0x20, %rdi
+f:
+  addq  $0x20, %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+  .cfi_endproc
+  .size _start, .-_start
+  .data
+JT:
+  .quad c
+  .quad d
+  .quad f
+
+
+# CHECK:   BOLT-INFO: Shrink wrapping moved 2 spills inserting load/stores and 0 spills inserting push/pops

bolt/test/X86/shrinkwrapping-popf.s

# FDATA: 0 [unknown] 0 1 _start 0 0 6
     je a
 b:  jne _start
 # FDATA: 1 _start #b# 1 _start #c# 0 3
 
 c:
     pushf
     push  %rbx
     push  %rbp
+    addq $0x50, -0x30(%rbp)
     pop   %rbp
     pop   %rbx
     popf
 
 # This basic block is treated as having 0 execution count.
 a:
     ud2
     .cfi_endproc

bolt/test/X86/shrinkwrapping-restore-position.s

+# This checks that shrink wrapping uses the red zone defined in the X86 ABI by
+# placing restores that access elements already deallocated by the stack.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt -relocs %t.exe -o %t.out -data %t.fdata \
+# RUN:     -frame-opt=all -simplify-conditional-tail-calls=false \
+# RUN:     -experimental-shrink-wrapping \
+# RUN:     -eliminate-unreachable=false | FileCheck %s
+# RUN: llvm-objdump -d %t.out --print-imm-hex | \
+# RUN:   FileCheck --check-prefix CHECK-OBJDUMP %s
+
+
+# Here we create a CFG where the restore position matches the previous (deleted)
+# restore position. Shrink wrapping then will put a stack access to an element
+# that was deallocated at the previously deleted POP, which falls in the red
+# zone and should be safe for X86 Linux ABI.
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+# FDATA: 0 [unknown] 0 1 _start 0 0 1
+  push  %rbp
+  mov   %rsp, %rbp
+  push  %rbx
+  push  %r14
+  subq  $0x20, %rsp
+b:  je  hot_path
+# FDATA: 1 _start #b# 1 _start #hot_path# 0 1
+cold_path:
+  mov %r14, %rdi
+  mov %rbx, %rdi
+  movq rel(%rip), %rdi  # Add this to create a relocation and run bolt w/ relocs
+  leaq -0x20(%rbp), %r14
+  movq -0x20(%rbp), %rdi
+  leaq -0x10(%rbp), %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+hot_path:
+  addq  $0x20, %rsp
+  pop %r14
+  pop %rbx
+  pop %rbp
+  ret
+  .cfi_endproc
+  .size _start, .-_start
+
+  .data
+rel:  .quad cold_path
+
+# CHECK:   BOLT-INFO: Shrink wrapping moved 2 spills inserting load/stores and 0 spills inserting push/pops
+
+# CHECK-OBJDUMP:     <_start>:
+# CHECK-OBJDUMP:         leaq    (%rbp), %rsp
+# CHECK-OBJDUMP-NEXT:    popq    %rbp
+# CHECK-OBJDUMP-NEXT:    movq    -0x10(%rsp), %rbx
+# CHECK-OBJDUMP-NEXT:    movq    -0x18(%rsp), %r14
+# CHECK-OBJDUMP-NEXT:    retq