[X86] Add X86SafeStackBoundsCheckingCombiner pass

[mlemay-intel]

Summary

The X86SafeStackBoundsCheckingCombiner pre-emit pass analyzes MPX BNDCU
instructions inserted by the X86SafeStackBoundsChecking pass to combine
redundant checks and eliminate unneeded checks.

lib/Target/X86/CMakeLists.txt

set(sources
   X86InstrInfo.cpp
   X86EvexToVex.cpp
   X86MCInstLower.cpp
   X86MachineFunctionInfo.cpp
   X86OptimizeLEAs.cpp
   X86PadShortFunction.cpp
   X86RegisterInfo.cpp
   X86SafeStackBoundsChecking.cpp
+  X86SafeStackBoundsCheckingCombiner.cpp
   X86SelectionDAGInfo.cpp
   X86ShuffleDecodeConstantPool.cpp
   X86Subtarget.cpp
   X86TargetMachine.cpp
   X86TargetObjectFile.cpp
   X86TargetTransformInfo.cpp
   X86VZeroUpper.cpp
   X86WinAllocaExpander.cpp

lib/Target/X86/X86.h

 void initializeEvexToVexInstPassPass(PassRegistry &);
 
 /// This pass adds bounds checks to prevent stray writes from corrupting the
 /// safe stack.
 FunctionPass *createX86SafeStackBoundsCheckingPass(X86TargetMachine &TM);
 
 void initializeX86SafeStackBoundsCheckingPass(PassRegistry&);
 
+/// Return a pass that combines redundant bounds checks inserted by the
+/// X86SafeStackBoundsChecking pass.
+FunctionPass *createX86SafeStackBoundsCheckingCombinerPass();
+
+void initializeX86SafeStackBoundsCheckingCombinerPass(PassRegistry&);
 } // End llvm namespace
 
 #endif

lib/Target/X86/X86SafeStackBoundsCheckingCombiner.cpp

+//===----- X86SafeStackBoundsCheckingCombiner.cpp - Combine checks --------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a pass that analyzes MPX BNDCU instructions inserted by the
+// X86SafeStackBoundsChecking pass to combine redundant checks and eliminate
+// unneeded checks.
+//
+//===----------------------------------------------------------------------===//
+
+#include "X86.h"
+#include "X86InstrBuilder.h"
+#include "X86Subtarget.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include <forward_list>
+#include <unordered_set>
+
+using namespace llvm;
+
+#define SAFESTKBCCOMBINER_DESC "X86 Safe Stack Bounds Checking Combiner"
+#define SAFESTKBCCOMBINER_NAME "x86-safestack-bounds-checking-combiner"
+#define DEBUG_TYPE SAFESTKBCCOMBINER_NAME
+#define DEBUG_LABEL "[X86SafeStackBoundsCheckingCombiner] "
+
+namespace {
+
+class X86SafeStackBoundsCheckingCombiner : public MachineFunctionPass {
+public:
+  static char ID;
+
+  X86SafeStackBoundsCheckingCombiner();
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  llvm::StringRef getPassName() const override {
+    return SAFESTKBCCOMBINER_DESC;
+  }
+
+  bool processInstr(MachineInstr *MI);
+
+  const TargetInstrInfo *TII;
+  const X86RegisterInfo *TRI;
+  std::unordered_set<MachineInstr *> PrevBndChks;
+  std::unordered_set<MachineInstr *> ToErase;
+};
+
+char X86SafeStackBoundsCheckingCombiner::ID = 0;
+}
+
+INITIALIZE_PASS(X86SafeStackBoundsCheckingCombiner,
+                SAFESTKBCCOMBINER_NAME, SAFESTKBCCOMBINER_DESC, false, false)
+
+FunctionPass *llvm::createX86SafeStackBoundsCheckingCombinerPass() {
+  return new X86SafeStackBoundsCheckingCombiner();
+}
+
+X86SafeStackBoundsCheckingCombiner::X86SafeStackBoundsCheckingCombiner()
+  : MachineFunctionPass(ID) {
+
+  initializeX86SafeStackBoundsCheckingCombinerPass(*PassRegistry::getPassRegistry());
+}
+
+bool X86SafeStackBoundsCheckingCombiner::processInstr(MachineInstr *MI) {
+  X86AddressMode AM;
+  bool IsMem = false, IsBndc = false;
+
+  constexpr unsigned BNDCU_IASM_MEM_OP = 3;
+
+  std::forward_list<MachineInstr *> PrevToErase;
+  for (MachineInstr *PrevMI : PrevBndChks) {
+    for (MachineOperand &Op : PrevMI->uses()) {
+      if (!(Op.isReg() && MI->definesRegister(Op.getReg())))
+        continue;
+
+      DEBUG(dbgs() << DEBUG_LABEL << "Register " << Op.getReg() << " used by "; PrevMI->print(dbgs());
+            dbgs() << " is redefined by "; MI->print(dbgs()); dbgs() << "\n");
+
+      PrevToErase.push_front(PrevMI);
+      break;
+    }
+  }
+  for (MachineInstr *PrevMI : PrevToErase)
+    PrevBndChks.erase(PrevMI);
+
+  unsigned Opcode = MI->getOpcode();
+
+  const MCInstrDesc &Desc = TII->get(Opcode);
+  uint64_t TSFlags = Desc.TSFlags;
+
+  // Determine where the memory operand starts:
+  int MemoryOperand = -1;
+  if (MI->mayStore())
+    MemoryOperand = X86II::getMemoryOperandNo(TSFlags);
+  if (MemoryOperand == -1) {
+    if (!MI->isInlineAsm())
+      return false;
+
+    const char *IAsmStr = MI->getOperand(0).getSymbolName();
+    if (strcmp(IAsmStr, "bndcu $0, %bnd0") != 0)
+      return false;
+
+    IsMem = true;
+    IsBndc = true;
+
+    MemoryOperand = BNDCU_IASM_MEM_OP;
+  } else {
+    DEBUG(dbgs() << DEBUG_LABEL
+                 << "Found store: ";
+          MI->print(dbgs()); dbgs() << "\n");
+
+    MemoryOperand += X86II::getOperandBias(Desc);
+
+    IsMem = true;
+  }
+
+  if (!IsMem)
+    return false;
+
+  AM = getAddressFromInstr(MI, MemoryOperand);
+
+  if (!IsBndc)
+    return false;
+
+  DEBUG(dbgs() << DEBUG_LABEL
+               << "Identified bound check: ";
+        MI->print(dbgs()); dbgs() << "\n");
+
+  MachineOperand &SegRegOp = MI->getOperand(MemoryOperand + X86::AddrSegmentReg);
+  if (TRI->isPhysicalRegister(SegRegOp.getReg())) {
+    // This erases any bound checks with segment override prefixes that may be
+    // emitted by the X86SafeStackBoundsChecking pass, since that pass may not
+    // detect all such thread-local accesses.
+    DEBUG(dbgs() << DEBUG_LABEL << "Erasing bound check with segment override "
+                 << "prefix.\n");
+    ToErase.insert(MI);
+    return false;
+  }
+
+  assert(AM.BaseType == X86AddressMode::RegBase);
+
+  unsigned BaseReg = AM.Base.Reg;
+  unsigned IndexReg = AM.IndexReg;
+
+  bool BaseIsPhysReg = TRI->isPhysicalRegister(BaseReg);
+  bool IndexIsPhysReg = TRI->isPhysicalRegister(IndexReg);
+
+  if (!(BaseIsPhysReg || IndexIsPhysReg)) {
+    DEBUG(dbgs() << DEBUG_LABEL << "Memory operand without a base or index "
+                 << "register assumed to be safe.\n");
+    ToErase.insert(MI);
+    return false;
+  }
+
+  if (BaseReg == X86::RIP && !IndexIsPhysReg) {
+    DEBUG(dbgs() << DEBUG_LABEL << "Memory operand with simple RIP-relative "
+                 << "displacement assumed to be safe.\n");
+    ToErase.insert(MI);
+    return false;
+  }
+
+  auto hasSameBaseMemOp = [&](const X86AddressMode &ChkAM) {
+    return AM.Base.Reg == ChkAM.Base.Reg &&
+        AM.IndexReg == ChkAM.IndexReg &&
+        AM.Scale == ChkAM.Scale &&
+        AM.GV == ChkAM.GV &&
+        AM.GVOpFlags == ChkAM.GVOpFlags;
+  };
+
+  for (MachineInstr *PrevMI : PrevBndChks) {
+    X86AddressMode ChkAM = getAddressFromInstr(PrevMI, BNDCU_IASM_MEM_OP);
+    if (!hasSameBaseMemOp(ChkAM)) continue;
+
+    DEBUG(dbgs() << DEBUG_LABEL
+                 << "Eliding check by merging with earlier one: ";
+          PrevMI->print(dbgs()); dbgs() << "\n");
+
+    ToErase.insert(MI);
+
+    MachineOperand &DispOp =
+        PrevMI->getOperand(BNDCU_IASM_MEM_OP + X86::AddrDisp);
+    if (DispOp.isImm()) {
+      auto CurrOff = DispOp.getImm();
+      if (AM.Disp <= CurrOff)
+        return false;
+
+      DispOp.setImm(AM.Disp);
+    } else {
+      auto CurrOff = DispOp.getOffset();
+      if (AM.Disp <= CurrOff)
+        return false;
+
+      DispOp.setOffset(AM.Disp);
+    }
+
+    return true;
+  }
+
+  PrevBndChks.insert(MI);
+
+  return false;
+}
+
+bool X86SafeStackBoundsCheckingCombiner::runOnMachineFunction(MachineFunction &MF) {
+  const X86Subtarget *STI = &MF.getSubtarget<X86Subtarget>();
+  TII = STI->getInstrInfo();
+  TRI = STI->getRegisterInfo();
+
+  if (!(STI->useSeparateStackSeg() && STI->is64Bit()))
+    return false;
+
+  bool Changed = false;
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB)
+      Changed |= processInstr(&MI);
+
+    PrevBndChks.clear();
+  }
+
+  for (MachineInstr *MI : ToErase) {
+    MI->eraseFromParent();
+    Changed = true;
+  }
+  ToErase.clear();
+
+  return Changed;
+}

lib/Target/X86/X86TargetMachine.cpp

extern "C" void LLVMInitializeX86Target() {
   RegisterTargetMachine<X86TargetMachine> Y(getTheX86_64Target());
 
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   initializeGlobalISel(PR);
   initializeWinEHStatePassPass(PR);
   initializeFixupBWInstPassPass(PR);
   initializeEvexToVexInstPassPass(PR);
   initializeX86SafeStackBoundsCheckingPass(PR);
+  initializeX86SafeStackBoundsCheckingCombinerPass(PR);
 }
 
 static std::unique_ptr<TargetLoweringObjectFile> createTLOF(const Triple &TT) {
   if (TT.isOSBinFormatMachO()) {
     if (TT.getArch() == Triple::x86_64)
       return make_unique<X86_64MachoTargetObjectFile>();
     return make_unique<TargetLoweringObjectFileMachO>();
   }
if (UseVZeroUpper)
     addPass(createX86IssueVZeroUpperPass());
 
   if (getOptLevel() != CodeGenOpt::None) {
     addPass(createX86FixupBWInsts());
     addPass(createX86PadShortFunctions());
     addPass(createX86FixupLEAs());
     addPass(createX86EvexToVexInsts());
   }
+
+  addPass(createX86SafeStackBoundsCheckingCombinerPass());
 }

test/CodeGen/X86/safestack-bounds-checking-combiner.mir

+# RUN: llc -mtriple=x86_64-linux-gnu -mattr=+mpx,+separate-stack-seg -start-before x86-safestack-bounds-checking-combiner -stop-after x86-safestack-bounds-checking-combiner -o - %s | FileCheck %s
+
+--- |
+
+  @__safestack_unsafe_stack_ptr = external thread_local(initialexec) global i8*
+
+  ; Function Attrs: nounwind safestack uwtable
+  define void @overflow_gep_store() #0 {
+  entry:
+    %unsafe_stack_ptr = load i8*, i8** @__safestack_unsafe_stack_ptr
+    %unsafe_stack_static_top = getelementptr i8, i8* %unsafe_stack_ptr, i32 -16
+    store i8* %unsafe_stack_static_top, i8** @__safestack_unsafe_stack_ptr
+    %0 = getelementptr i8, i8* %unsafe_stack_ptr, i32 -8
+    %a.unsafe = bitcast i8* %0 to i32*
+    %1 = bitcast i32* %a.unsafe to i8*
+    %2 = getelementptr i8, i8* %1, i32 4
+    %3 = getelementptr i8, i8* %2, i64 1
+    call void asm "bndcu $0, %bnd0", "*m"(i8* %3)
+    store i8 0, i8* %2
+    %4 = getelementptr i8, i8* %unsafe_stack_ptr, i32 -4
+    %b.unsafe = bitcast i8* %4 to i32*
+    %5 = bitcast i32* %b.unsafe to i8*
+    %6 = getelementptr i8, i8* %5, i32 4
+    %7 = getelementptr i8, i8* %6, i64 1
+    call void asm "bndcu $0, %bnd0", "*m"(i8* %7)
+    store i8 0, i8* %6
+    store i8* %unsafe_stack_ptr, i8** @__safestack_unsafe_stack_ptr
+    ret void
+  }
+
+  attributes #0 = { nounwind safestack uwtable "target-features"="+mpx,+separate-stack-seg" }
+
+...
+---
+name:            overflow_gep_store
+alignment:       4
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+tracksRegLiveness: true
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    0
+  adjustsStack:    false
+  hasCalls:        false
+  maxCallFrameSize: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+body:             |
+  bb.0.entry:
+    %rax = MOV64rm %rip, 1, _, target-flags(x86-gottpoff) @__safestack_unsafe_stack_ptr, _ :: (load 8 from got)
+    %rcx = MOV64rm %rax, 1, _, 0, %fs :: (dereferenceable load 8 from @__safestack_unsafe_stack_ptr)
+    %rdx = LEA64r %rcx, 1, _, -16, _
+    MOV64mr %rax, 1, _, 0, %fs, killed %rdx :: (store 8 into @__safestack_unsafe_stack_ptr)
+    INLINEASM $"bndcu $0, %bnd0", 8, 196654, %rcx, 1, _, -3, _
+    ; CHECK: INLINEASM $"bndcu $0, %bnd0", 8, 196654, %rcx, 1, _, 1, _
+    ; CHECK-NEXT: MOV8mi
+    ; CHECK-NEXT: MOV8mi
+    MOV8mi %rcx, 1, _, -4, _, 0 :: (store 1 into %ir.2)
+    INLINEASM $"bndcu $0, %bnd0", 8, 196654, %rcx, 1, _, 1, _
+    MOV8mi %rcx, 1, _, 0, _, 0 :: (store 1 into %ir.6)
+    MOV64mr killed %rax, 1, _, 0, %fs, killed %rcx :: (store 8 into @__safestack_unsafe_stack_ptr)
+    RETQ
+
+...