[BOLT] Fix AND evaluation bug in shrink wrapping

[rafauler]

Summary

Fix a bug where shrink-wrapping would use wrong stack offsets
because the stack was being aligned with an AND instruction, hence,
making its true offsets only available during runtime (we can't
statically determine where are the stack elements and we must give up
on this case).

[Amir]

I'm concerned that the fix is an overkill and doesn't address the root cause.
From what I understand, evaluateSimple can only evaluate a constant when given an instruction it understands and a pair of <reg, val> for use in evaluation. So the problem must be with the calling code – that the value of rsp is considered a constant while it's actually not (at least not in this case).

Even if it's not the case, I would like to keep the evaluation of and – it might be handy in other contexts – but restrict shrink-wrapping from using it.

[rafauler]

In D126110#3529103, @Amir wrote:

I'm concerned that the fix is an overkill and doesn't address the root cause.
From what I understand, evaluateSimple can only evaluate a constant when given an instruction it understands and a pair of <reg, val> for use in evaluation. So the problem must be with the calling code – that the value of rsp is considered a constant while it's actually not (at least not in this case).

Even if it's not the case, I would like to keep the evaluation of and – it might be handy in other contexts – but restrict shrink-wrapping from using it.

That makes a lot of sense! Here's what I'm thinking. I would like to have two versions of this function to support this requirement. When we are evaluating offsets (such as stack offsets), only a subset of operations are supported. That's because most of the time we are operating without full knowledge of the actual operands, but rather offsets that are added to an unknown base (x + Offset, and we are doing transformations to Offset to reason about a specific property). In these cases, using other operators that don't have associativity with addition of (x + Offset), where x is the stack base, will break the analysis. In the case of the bug fixed here, "(x + Offset) AND Constant" is not the same as "x + (Offset AND Constant)" (no associativity), hence we can't support AND. But "(x + Offset) + Displacement+ Index * Size", such as LEA, is the same as "x + (Offset + Displacement + Index * Size)", so it's fine to support LEA as long as we are only feeding Base as the input of the expression. ADD and SUB are supported as well.

But I don't think it's a good idea to add a boolean value such as "OnlyBaseOffsetAssociativeOperations=true" as a parameter to this function. Rather, I would prefer to break it into two functions, evaluateSimple, used for offsets, and evaluate(), which calls evaluateSimple and if it fails, try a bunch more operations that assume that the operands are fully known during evaluation time. But at the moment, I can't find any places in our codebase that would be users of evaluate() (the more general evaluation function), and thus I'm less inclined to add it as it would be currently dead code. Let me know if you have other ideas on how to move forward. Meanwhile I'll improve the comments on the usage of evaluateSimple to make it clear its intended usage.

[Amir]

In D126110#3532607, @rafauler wrote:

In D126110#3529103, @Amir wrote:

I'm concerned that the fix is an overkill and doesn't address the root cause.
From what I understand, evaluateSimple can only evaluate a constant when given an instruction it understands and a pair of <reg, val> for use in evaluation. So the problem must be with the calling code – that the value of rsp is considered a constant while it's actually not (at least not in this case).

Even if it's not the case, I would like to keep the evaluation of and – it might be handy in other contexts – but restrict shrink-wrapping from using it.

That makes a lot of sense! Here's what I'm thinking. I would like to have two versions of this function to support this requirement. When we are evaluating offsets (such as stack offsets), only a subset of operations are supported. That's because most of the time we are operating without full knowledge of the actual operands, but rather offsets that are added to an unknown base (x + Offset, and we are doing transformations to Offset to reason about a specific property). In these cases, using other operators that don't have associativity with addition of (x + Offset), where x is the stack base, will break the analysis. In the case of the bug fixed here, "(x + Offset) AND Constant" is not the same as "x + (Offset AND Constant)" (no associativity), hence we can't support AND. But "(x + Offset) + Displacement+ Index * Size", such as LEA, is the same as "x + (Offset + Displacement + Index * Size)", so it's fine to support LEA as long as we are only feeding Base as the input of the expression. ADD and SUB are supported as well.

But I don't think it's a good idea to add a boolean value such as "OnlyBaseOffsetAssociativeOperations=true" as a parameter to this function. Rather, I would prefer to break it into two functions, evaluateSimple, used for offsets, and evaluate(), which calls evaluateSimple and if it fails, try a bunch more operations that assume that the operands are fully known during evaluation time. But at the moment, I can't find any places in our codebase that would be users of evaluate() (the more general evaluation function), and thus I'm less inclined to add it as it would be currently dead code. Let me know if you have other ideas on how to move forward. Meanwhile I'll improve the comments on the usage of evaluateSimple to make it clear its intended usage.

Thank you, now I understand the problem and how this change addresses it.
I agree that removing the handling of AND and renaming the function to something like evaluateAssociative (simple is not self-explanatory) would be OK.
We can add a generic evaluate later on when the need arises.

[rafauler]

Rename evaluateSimple function and make it clear its intended usage.

[Amir]

Thanks!

[rafauler]

Fix nit in testcase and update assert message in ValidateInternalCalls

bolt/include/bolt/Core/MCPlusBuilder.h

public:
 
   /// Identify stack adjustment instructions -- those that change the stack
   /// pointer by adding or subtracting an immediate.
   virtual bool isStackAdjustment(const MCInst &Inst) const {
     llvm_unreachable("not implemented");
     return false;
   }
 
-  /// Use \p Input1 or Input2 as the current value for the input register and
-  /// put in \p Output the changes incurred by executing \p Inst. Return false
-  /// if it was not possible to perform the evaluation.
-  virtual bool evaluateSimple(const MCInst &Inst, int64_t &Output,
-                              std::pair<MCPhysReg, int64_t> Input1,
-                              std::pair<MCPhysReg, int64_t> Input2) const {
+  /// Use \p Input1 or Input2 as the current value for the input
+  /// register and put in \p Output the changes incurred by executing
+  /// \p Inst. Return false if it was not possible to perform the
+  /// evaluation. evaluateStackOffsetExpr is restricted to operations
+  /// that have associativity with addition. Its intended usage is for
+  /// evaluating stack offset changes. In these cases, expressions
+  /// appear in the form of (x + offset) OP constant, where x is an
+  /// unknown base (such as stack base) but offset and constant are
+  /// known. In these cases, \p Output represents the new stack offset
+  /// after executing \p Inst. Because we don't know x, we can't
+  /// evaluate operations such as multiply or AND/OR, e.g. (x +
+  /// offset) OP constant is not the same as x + (offset OP constant).
+  virtual bool
+  evaluateStackOffsetExpr(const MCInst &Inst, int64_t &Output,
     llvm_unreachable("not implemented");
     return false;
   }
 
   virtual bool isRegToRegMove(const MCInst &Inst, MCPhysReg &From,
                               MCPhysReg &To) const {
     llvm_unreachable("not implemented");
     return false;

bolt/include/bolt/Passes/StackPointerTracking.h

if (this->BC.MII->get(Point.getOpcode())
       else
         SP = std::make_pair(0, 0);
       std::pair<MCPhysReg, int64_t> FP;
       if (FPVal != EMPTY && FPVal != SUPERPOSITION)
         FP = std::make_pair(MIB->getFramePointer(), FPVal);
       else
         FP = std::make_pair(0, 0);
       int64_t Output;
-      if (!MIB->evaluateSimple(Point, Output, SP, FP)) {
+      if (!MIB->evaluateStackOffsetExpr(Point, Output, SP, FP)) {
         if (SPVal == EMPTY && FPVal == EMPTY)
           return SPVal;
         return SUPERPOSITION;
       }
 
       return static_cast<int>(Output);
     }
 
if (this->BC.MII->get(Point.getOpcode())
       else
         FP = std::make_pair(0, 0);
       std::pair<MCPhysReg, int64_t> SP;
       if (SPVal != EMPTY && SPVal != SUPERPOSITION)
         SP = std::make_pair(MIB->getStackPointer(), SPVal);
       else
         SP = std::make_pair(0, 0);
       int64_t Output;
-      if (!MIB->evaluateSimple(Point, Output, SP, FP)) {
+      if (!MIB->evaluateStackOffsetExpr(Point, Output, SP, FP)) {
         if (SPVal == EMPTY && FPVal == EMPTY)
           return FPVal;
         return SUPERPOSITION;
       }
 
       if (!HasFramePointer && MIB->escapesVariable(Point, false))
         HasFramePointer = true;
       return static_cast<int>(Output);

bolt/lib/Passes/AllocCombiner.cpp

 
 namespace llvm {
 namespace bolt {
 
 namespace {
 
 bool getStackAdjustmentSize(const BinaryContext &BC, const MCInst &Inst,
                             int64_t &Adjustment) {
-  return BC.MIB->evaluateSimple(Inst, Adjustment,
-                                std::make_pair(BC.MIB->getStackPointer(), 0LL),
-                                std::make_pair(0, 0LL));
+  return BC.MIB->evaluateStackOffsetExpr(
+      Inst, Adjustment, std::make_pair(BC.MIB->getStackPointer(), 0LL),
 }
 
 bool isIndifferentToSP(const MCInst &Inst, const BinaryContext &BC) {
   if (BC.MIB->isCFI(Inst))
     return true;
 
   const MCInstrDesc II = BC.MII->get(Inst.getOpcode());
   if (BC.MIB->isTerminator(Inst) ||

bolt/lib/Passes/ShrinkWrapping.cpp

void StackLayoutModifier::checkFramePointerInitialization(MCInst &Point) {
   std::pair<MCPhysReg, int64_t> SP;
 
   if (SPVal != SPT.EMPTY && SPVal != SPT.SUPERPOSITION)
     SP = std::make_pair(BC.MIB->getStackPointer(), SPVal);
   else
     SP = std::make_pair(0, 0);
 
   int64_t Output;
-  if (!BC.MIB->evaluateSimple(Point, Output, SP, FP))
+  if (!BC.MIB->evaluateStackOffsetExpr(Point, Output, SP, FP))
     return;
 
   // Not your regular frame pointer initialization... bail
   if (Output != SPVal)
     blacklistRegion(0, 0);
 }
 
 void StackLayoutModifier::checkStackPointerRestore(MCInst &Point) {
void StackLayoutModifier::checkStackPointerRestore(MCInst &Point) {
   std::pair<MCPhysReg, int64_t> SP;
 
   if (SPVal != SPT.EMPTY && SPVal != SPT.SUPERPOSITION)
     SP = std::make_pair(BC.MIB->getStackPointer(), SPVal);
   else
     SP = std::make_pair(0, 0);
 
   int64_t Output;
-  if (!BC.MIB->evaluateSimple(Point, Output, SP, FP))
+  if (!BC.MIB->evaluateStackOffsetExpr(Point, Output, SP, FP))
     return;
 
   // If the value is the same of FP, no need to adjust it
   if (Output == FPVal)
     return;
 
   // If an allocation happened through FP, bail
   if (Output <= SPVal) {

bolt/lib/Passes/StackAllocationAnalysis.cpp

if (BC.MII->get(Point.getOpcode())
     else
       SP = std::make_pair(0, 0);
     std::pair<MCPhysReg, int64_t> FP;
     if (FPOffset != SPT.EMPTY && FPOffset != SPT.SUPERPOSITION)
       FP = std::make_pair(MIB->getFramePointer(), FPOffset);
     else
       FP = std::make_pair(0, 0);
     int64_t Output;
-    if (!MIB->evaluateSimple(Point, Output, SP, FP))
+    if (!MIB->evaluateStackOffsetExpr(Point, Output, SP, FP))
       return Next;
 
     if (SPOffset < Output) {
       Next = doKill(Point, Next, Output - SPOffset);
       return Next;
     }
     if (SPOffset > Output) {
       Next.set(this->ExprToIdx[&Point]);
       return Next;
     }
   }
   return Next;
 }
 
 } // end namespace bolt
 } // end namespace llvm

bolt/lib/Passes/ValidateInternalCalls.cpp

for (MCInst &Inst : BB) {
         BitVector UsedRegs = BitVector(BC.MRI->getNumRegs(), false);
         BC.MIB->getTouchedRegs(Use, UsedRegs);
         if (!UsedRegs[Reg])
           continue;
         UseDetected = true;
         int64_t Output;
         std::pair<MCPhysReg, int64_t> Input1 = std::make_pair(Reg, 0);
         std::pair<MCPhysReg, int64_t> Input2 = std::make_pair(0, 0);
-        if (!BC.MIB->evaluateSimple(Use, Output, Input1, Input2)) {
-          LLVM_DEBUG(dbgs() << "Evaluate simple failed.\n");
+        if (!BC.MIB->evaluateStackOffsetExpr(Use, Output, Input1, Input2)) {
+          LLVM_DEBUG(dbgs() << "Evaluate stack offset expr failed.\n");
           return false;
         }
         if (Offset + Output < 0 ||
             Offset + Output > static_cast<int64_t>(Function.getSize())) {
           LLVM_DEBUG({
             dbgs() << "Detected out-of-range PIC reference in " << Function
                    << "\nReturn address load: ";
             BC.InstPrinter->printInst(TargetInst, 0, "", *BC.STI, dbgs());

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

bool isStackAdjustment(const MCInst &Inst) const override {
     for (int I = 0, E = MCII.getNumDefs(); I != E; ++I) {
       const MCOperand &Operand = Inst.getOperand(I);
       if (Operand.isReg() && Operand.getReg() == X86::RSP)
         return true;
     }
     return false;
   }
 
-  bool evaluateSimple(const MCInst &Inst, int64_t &Output,
+  bool
+  evaluateStackOffsetExpr(const MCInst &Inst, int64_t &Output,
                           std::pair<MCPhysReg, int64_t> Input1,
                           std::pair<MCPhysReg, int64_t> Input2) const override {
 
     auto getOperandVal = [&](MCPhysReg Reg) -> ErrorOr<int64_t> {
       if (Reg == Input1.first)
         return Input1.second;
       if (Reg == Input2.first)
         return Input2.second;
       return make_error_code(errc::result_out_of_range);
     };
 
     switch (Inst.getOpcode()) {
     default:
       return false;
 
-    case X86::AND64ri32:
-    case X86::AND64ri8:
-      if (!Inst.getOperand(2).isImm())
-        return false;
-      if (ErrorOr<int64_t> InputVal =
-              getOperandVal(Inst.getOperand(1).getReg()))
-        Output = *InputVal & Inst.getOperand(2).getImm();
-      else
-        return false;
-      break;
     case X86::SUB64ri32:
     case X86::SUB64ri8:
       if (!Inst.getOperand(2).isImm())
         return false;
       if (ErrorOr<int64_t> InputVal =
               getOperandVal(Inst.getOperand(1).getReg()))
         Output = *InputVal - Inst.getOperand(2).getImm();
       else

bolt/test/X86/shrinkwrapping-and-rsp.s

+# This checks that shrink wrapping does attempt at accessing stack elements
+# using RSP when the function is aligning RSP and changing offsets.
+
+# 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 -o %t.out -data %t.fdata \
+# RUN:     -frame-opt=all -simplify-conditional-tail-calls=false \
+# RUN:     -eliminate-unreachable=false | FileCheck %s
+
+# Here we have a function that aligns the stack at prologue. Stack pointer
+# analysis can't try to infer offset positions after AND because that depends
+# on the runtime value of the stack pointer of callee (whether it is misaligned
+# or not).
+  .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
+  and    $0xffffffffffffffe0,%rsp
+  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
+  # Block push-pop mode by using an instruction that requires the
+  # stack to be aligned to 128B. This will force the pass
+  # to try to index stack elements by using RSP +offset directly, but
+  # we do not know how to access individual elements of the stack thanks
+  # to the alignment.
+  movdqa  %xmm8, (%rsp)
+  addq  $0x20, %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
+
+# CHECK:   BOLT-INFO: Shrink wrapping moved 0 spills inserting load/stores and 0 spills inserting push/pops