[X86] Add X86FixupSeparateStack pass

[mlemay-intel]

Summary

The X86FixupSeparateStack pass traverses the graph of basic blocks for
a function and tracks which physical registers derive from the stack
register (to point at stack data) at each instruction. It also
modifies each instruction that uses one of those registers as the base
register in a memory operand, unless the register is ESP or EBP. This
is to direct those memory accesses to the stack segment. It also
modifies each instruction that uses EBP as the base register in a
memory operand if EBP does not derive from the stack register at that
instruction. This is to direct those memory accesses to the data
segment.

This pass is enabled with the +separate-stack-seg option.

[eugenis]

I may be able to look at this next week, but it is a bit outside of my area of expertise. Maybe find another reviewer?

Also, this needs tests. A lot of tests.

[mlemay-intel]

In D17095#349401, @eugenis wrote:

I may be able to look at this next week, but it is a bit outside of my area of expertise. Maybe find another reviewer?

Thank you. I added you as a reviewer to all of these since they are all related to SafeStack, but I'll also look for another reviewer as you suggested.

Also, this needs tests. A lot of tests.

I haven't yet found any existing tests in LLVM for a pass like this that I can use as an example, but I'll keep looking. The challenge in writing a test for this pass that fits in the LLVM testing framework is that this pass modifies particular instructions depending on which registers get allocated as base registers for the instructions' memory operands, etc. When a test is written as IR, I don't know of a way to control which registers get selected.

[mlemay-intel]

lib/Target/X86/X86FixupSeparateStack.cpp

627

Perhaps this should be a warning instead of a debug message. Are there examples of passes that issue warnings for similar things?

[delena]

lib/Target/X86/X86FixupSeparateStack.cpp

594

Are you working with load? I see PUSH32r and MRMDestMem - DestMem looks like store

597

New line after )

600

Why masked instructions are out of scope?
DestMem can't be VEX_4V.

611

Add assert (MemoryOperand != -1)

639

if (!Op.isReg() || !Op.isDef())

continue;

662

Do you need to delete SuccReqs?

678

Do you check this when you add the pass?

682

Who deletes the object created by "new" ?

lib/Target/X86/X86TargetMachine.cpp

288

for 32-bit only?

[craig.topper]

lib/Target/X86/X86.td

254

Can this be done with a command line switch in X86TargetMachine.cpp like VZeroUpper? I don't see that this needs to be a feature bit.

lib/Target/X86/X86FixupSeparateStack.cpp

600

DestMem can be VEX_4V. We check for it in both MCCodeEmitter and the Disassembler. I believe at least vmaskmovpd uses it.

[mlemay-intel]

Sorry for my delayed response. Activity notifications have not been reaching my inbox.

lib/Target/X86/X86.td

254

Your suggestion should work fine at the LLVM level, but then how would you suggest passing the flag to Clang?

lib/Target/X86/X86FixupSeparateStack.cpp

594

Loads are handled as well on line 632.

597

Will do. Thanks.

600

I assume that VEX-encoded instructions are not used for storing the stack pointer. I think that handling them would increase the complexity of line 610 for computing the operand number of the base register (I looked at X86MCCodeEmitter::EmitVEXOpcodePrefix to understand how these instructions are represented in LLVM). Thus, I skip them to avoid that extra complexity. I'll add a comment to note this decision.

611

X86II::getMemoryOperandNo always returns 0 for instructions of this form. I'll add a comment noting this assumption.

639

Good catch!

662

Ownership of the pointer is assigned to a shared_ptr in processBasicBlock. However, I will revise processBasicBlock to accept a shared_ptr parameter so that the pointer ownership is more clear.

678

I think it may make sense to keep this check here so that an error message is emitted if the associated command line option is specified for an unsupported target.

682

This pointer is assigned to a shared_ptr in processBasicBlock.

[mlemay-intel]

Expand comments, streamline code, and clarify object ownership.

[mlemay-intel]

Fix bug that occurred when cycle was formed with
AddrRegReqs::addPredecessor. AddrRegReqs::useInMemOp could modify
DenseSet being iterated in addPredecessor.

Track registers pointing into SS segment across some spills and fills.

Detect some instructions emitted for variadic argument intrinsics that
store SS pointers, avoiding assertion failures for those instructions.

Add option to control level of verification that is applied to
instructions that store SS pointers.

[mlemay-intel]

More extensive support for tracking spills and fills of registers containing stack pointers across basic blocks.
Recognize additional instructions that are used to derive stack pointers.

[mlemay-intel]

Corrected/updated some comments. Added a check that stack pointer spill requirements are satisfied.

[mlemay-intel]

Rebase.
Revise StackPtrSpillReqs::addPredecessor to correctly handle the case when Demand is empty.
Revise comments and fix whitespace.

[mlemay-intel]

Fixed a bug in how segment switching instructions were inserted around string instructions.
Made comment indentation more consistent.
Added support for tracking the flow of stack pointer values through instructions with multiple register inputs (e.g. CMOV instructions).
Enhanced spill/fill tracking to support functions with eliminated frame pointers.
Selected more efficient data structures.
Revised AddrRegReqs::derive to reflect the fact that stack pointer values do not flow through the flags register.

[mlemay-intel]

Avoids erasing previous derivation data when a register is both used and defined in an instruction.
Ignores implicit and undef registers when tracking derivation.
Removes code and associated options for tracking variadic argument pointers.
Eliminates assert on storing stack pointers to memory.

[zvi]

In D17095#349469, @mlemay-intel wrote:

In D17095#349401, @eugenis wrote:

I may be able to look at this next week, but it is a bit outside of my area of expertise. Maybe find another reviewer?

Thank you. I added you as a reviewer to all of these since they are all related to SafeStack, but I'll also look for another reviewer as you suggested.

Also, this needs tests. A lot of tests.

I haven't yet found any existing tests in LLVM for a pass like this that I can use as an example, but I'll keep looking. The challenge in writing a test for this pass that fits in the LLVM testing framework is that this pass modifies particular instructions depending on which registers get allocated as base registers for the instructions' memory operands, etc. When a test is written as IR, I don't know of a way to control which registers get selected.

Now that tests were added to this patch in .ll form, would it be better to change them to machine IR (.mir) form?

[mlemay-intel]

In D17095#565252, @zvi wrote:

In D17095#349469, @mlemay-intel wrote:

In D17095#349401, @eugenis wrote:

I may be able to look at this next week, but it is a bit outside of my area of expertise. Maybe find another reviewer?

Thank you. I added you as a reviewer to all of these since they are all related to SafeStack, but I'll also look for another reviewer as you suggested.

Also, this needs tests. A lot of tests.

I haven't yet found any existing tests in LLVM for a pass like this that I can use as an example, but I'll keep looking. The challenge in writing a test for this pass that fits in the LLVM testing framework is that this pass modifies particular instructions depending on which registers get allocated as base registers for the instructions' memory operands, etc. When a test is written as IR, I don't know of a way to control which registers get selected.

Now that tests were added to this patch in .ll form, would it be better to change them to machine IR (.mir) form?

Yes, that's exactly the type of testing infrastructure I was looking for, but I didn't know it existed until you pointed it out. Thanks. I'll update this patch.

[mlemay-intel]

Add new pass to registry.
Convert tests to MIR form.
Update tests to reflect previous removal of assertion for certain stores of stack pointers.

[mlemay-intel]

Added -sep-stk-seg-flat-mem-func option to specify functions that use a flat memory model.
Fixed bug in StackPtrSpillReqs::supply to prevent coalescing of ranges that need to remain separate.

[mlemay-intel]

Revised comments.
Converted a condition into an assertion in AddrRegReqs::lookupRoots.

[mlemay-intel]

Mark EAX as not being derived from any possible stack pointer registers after calls.

lib/Target/X86/CMakeLists.txt

 tablegen(LLVM X86GenSubtargetInfo.inc -gen-subtarget)
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
   X86AsmPrinter.cpp
   X86CallFrameOptimization.cpp
   X86ExpandPseudo.cpp
   X86FastISel.cpp
+  X86FixupSeparateStack.cpp
   X86FloatingPoint.cpp
   X86FrameLowering.cpp
   X86ISelDAGToDAG.cpp
   X86ISelLowering.cpp
   X86InstrInfo.cpp
   X86MCInstLower.cpp
   X86MachineFunctionInfo.cpp
   X86PadShortFunction.cpp

lib/Target/X86/X86.h

 /// the MachineInstr to MC.
 FunctionPass *createX86ExpandPseudoPass();
 
 /// Return a Machine IR pass that selectively replaces
 /// certain byte and word instructions by equivalent 32 bit instructions,
 /// in order to eliminate partial register usage, false dependences on
 /// the upper portions of registers, and to save code size.
 FunctionPass *createX86FixupBWInsts();
+
+/// Return a pass that helps to support separate stack and data segments by
+/// emitting segment override prefixes where necessary.
+FunctionPass *createX86FixupSeparateStack();
 } // End llvm namespace
 
 #endif

lib/Target/X86/X86.td

 def FeatureSoftFloat
     : SubtargetFeature<"soft-float", "UseSoftFloat", "true",
                        "Use software floating point features.">;
 // On at least some AMD processors, there is no performance hazard to writing
 // only the lower parts of a YMM register without clearing the upper part.
 def FeatureFastPartialYMMWrite
     : SubtargetFeature<"fast-partial-ymm-write", "HasFastPartialYMMWrite",
                        "true", "Partial writes to YMM registers are fast">;
+def FeatureSeparateStackSeg
+    : SubtargetFeature<"separate-stack-seg", "UseSeparateStackSeg", "true",
+                       "Help to support using different stack and data segments"
+                       " by adding segment override prefixes (X86-32 only).">;
 
 //===----------------------------------------------------------------------===//
 // X86 processors supported.
 //===----------------------------------------------------------------------===//
 
 include "X86Schedule.td"
 
 def ProcIntelAtom : SubtargetFeature<"atom", "X86ProcFamily", "IntelAtom",
                     "Intel Atom processors">;

lib/Target/X86/X86FixupSeparateStack.cpp

+//===----- X86FixupSeparateStack.cpp - Fixup stack memory accesses --------===//
+//
+//                     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 emits segment override prefixes to help support
+// separate stack and data segments for X86-32.
+//
+//===----------------------------------------------------------------------===//
+
+#include <algorithm>
+#include <map>
+#include <memory>
+#include <set>
+#include <vector>
+
+#include "X86.h"
+#include "X86InstrInfo.h"
+#include "X86Subtarget.h"
+#include "X86MachineFunctionInfo.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetInstrInfo.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "x86-fixup-separate-stack"
+
+namespace {
+/// Records information about registers that directly or indirectly are used in
+/// memory operands within a basic block.  It is used to distinguish registers
+/// that point to stack memory from those that do not.  It only explicitly
+/// stores information about registers that point to the stack, which implies
+/// that all other physical registers do not point to the stack.
+class AddrRegs {
+private:
+  const X86RegisterInfo *TRI;
+
+public:
+  /// Registers that point to the stack
+  DenseSet<unsigned> Stack;
+
+  AddrRegs(const X86RegisterInfo *TRI_) : TRI(TRI_) {}
+
+  void insertStackReg(unsigned Reg);
+
+  /// \brief Return true iff Reg points to a stack location.
+  bool pointsIntoStack(unsigned Reg) const;
+};
+
+/// Requirements for how registers that are directly or indirectly used in
+/// memory operands are configured.
+///
+/// Each basic block is associated with one AddrRegReqs object, although that
+/// object may actually contain information about the successors of the basic
+/// block with which it is associated.
+class AddrRegReqs {
+private:
+  const X86RegisterInfo *TRI;
+
+  /// The collection of AddrRegReqs objects that is associated with a function
+  /// is structured as a directed graph that may contain cycles.  This set of
+  /// Predecessors points to the AddrRegReqs objects corresponding to all of the
+  /// basic blocks that immediately precede the one associated with this object
+  /// and that have already been visited or are currently being visited by this
+  /// pass.
+  std::set<AddrRegReqs *> Predecessors;
+  /// The original status of the registers at the time this object branched from
+  /// its first predecessor, according to the traversal order of this pass.
+  const AddrRegs OrigRegs;
+
+  /// Pairs of registers (A, B) in which A may derive a pointer to stack data
+  /// from B.  This is the case if an instruction form that the compiler selects
+  /// to compute a pointer to stack data uses B to define A.  See
+  /// X86FixupSeparateStack::mayComputeStackPtr for more details.
+  ///
+  /// This also indicates cancellations of implicit derivations.  Every physical
+  /// register is implicitly derived from itself unless an entry in Derivatives
+  /// indicates otherwise.  Specifically, it is as though Derivatives contains
+  /// (A, A) for all physical registers A unless there exists an entry (A, X) in
+  /// Derivatives for some physical register X or the special value
+  /// X86::NoRegister.  This special value indicates that A cannot possibly
+  /// point to stack data.
+  ///
+  /// Note that the presence of an entry (A, B) in Derivatives does NOT
+  /// automatically indicate that A points to stack data.  It only indicates
+  /// that if B pointed to stack data when this basic block was entered, then A
+  /// points to stack data at the current instruction.  Also keep in mind that
+  /// Derivatives gets updated as additional instructions in the basic block are
+  /// processed.
+  ///
+  /// The entry (ESP, ESP) is always implicitly (and never explicitly) present
+  /// in Derivatives.
+  std::vector<std::pair<unsigned, unsigned> > Derivatives;
+  /// Root registers that are used directly or indirectly in memory operands
+  DenseSet<unsigned> UsedInMemOp;
+
+  /// \brief Lookup the register from which To may derive a pointer to stack
+  /// data, if such a register exists.
+  ///
+  /// \returns The register Root if an implicit or explicit entry (To, Root) is
+  /// currently present in Derivatives.  Otherwise, returns X86::NoRegister.
+  /// Note that there may be an entry (To, X86::NoRegister) in Derivatives,
+  /// which would also result in X86::NoRegister being returned.
+  unsigned lookupRoot(unsigned To) const;
+
+public:
+  /// \brief Construct an object with no initial predecessor.  One or more
+  /// predecessors may still be added later with addPredecessor.
+  AddrRegReqs(const X86RegisterInfo *TRI_) : TRI(TRI_), OrigRegs(TRI_) {}
+  /// \brief Construct an object with one initial predecessor.  Initialize
+  /// OrigRegs to represent the register configuration at the end of
+  /// Predecessor, since this is only invoked after Predecessor has been
+  /// processed.
+  AddrRegReqs(AddrRegReqs *Predecessor);
+
+  /// \brief Extract the register configuration represented by the current state
+  /// of this object.  This is based on the contents of OrigRegs, and it is
+  /// adjusted to account for the contents of Derivatives.
+  AddrRegs extractRegs();
+
+  /// \brief Register a new predecessor for this object.  This involves two
+  /// steps:
+  /// 1. Invoke Predecessor->useInMemOp for each register in UsedInMemOp, since
+  ///    the basic block corresponding to this object uses those registers even
+  ///    though the instructions that are in the predecessor basic block may not
+  ///    have used those registers.
+  /// 2. Insert Predecessor into Predecessors so that any time useInMemOp is
+  ///    invoked on this object, it may also invoke Predecessor->useInMemOp as
+  ///    necessary.
+  void addPredecessor(AddrRegReqs *Predecessor);
+
+  /// \brief Record that To is derived from From in Derivatives.
+  void derive(unsigned From, unsigned To);
+
+  /// \brief Record that To is redefined by some operation that cannot be used
+  /// to derive a pointer into the stack.
+  void markNotDerived(unsigned To);
+
+  /// \brief Record that Reg is used as a base address register in a memory
+  /// operand and check that the requirement for this to be a stack or data
+  /// pointer is met.
+  ///
+  /// This routine is invoked recursively for each predecessor of this object,
+  /// and there may even be cycles.  However, the recursion will eventually
+  /// terminate.  The routine is only invoked recursively if an entry is added
+  /// to this->UsedInMemOp.  The only entries that can be added are those
+  /// computed from the contents of Derivatives and OrigRegs.  The resultant set
+  /// is a subset of the set of physical registers, which is finite.  Entries
+  /// are never removed from UsedInMemOp by this routine.  Thus, this routine
+  /// will eventually run out of entries to be added to UsedInMemOp.
+  ///
+  /// \param Reg Register that is used as a base address register in a memory
+  /// operand.
+  /// \param IsStackReg Set to true iff Reg is required to point to stack
+  /// memory.  Otherwise, Reg is required to point to data memory.
+  void useInMemOp(unsigned Reg, bool IsStackReg);
+
+  /// \brief Return true iff Reg points to a stack location at the current
+  /// instruction.
+  bool pointsIntoStack(unsigned Reg) const;
+};
+
+/// This pass performs a depth-first traversal of each function, visiting each
+/// basic block once.
+///
+/// To understand how this pass determines whether its transformation is
+/// correct, it may be helpful to think of each basic block as an electronic
+/// circuit with 12V and 5V inputs.  Those inputs take the form of register
+/// values fed into basic blocks.  Think of pointers to stack data as 12V
+/// inputs and other register values as 5V inputs.
+///
+/// Continuing this analogy, the problem that this pass seeks to solve is to
+/// determine for each circuit element whether it is always wired to 12V or 5V.
+/// This pass considers each instruction to be analogous to a circuit element.
+/// Of course, each basic block may be invoked after multiple predecessors.
+/// Think of this like a circuit that can have multiple copies of a particular
+/// circuit block comprising circuit elements such as resistors and capacitors,
+/// and all replicas of a particular element must use the same voltage inputs.
+///
+/// An assumption that helps to make the problem more tractable for this pass
+/// is that pointers to data on the stack are only computed using a couple of
+/// instruction forms, specifically MOV32rr and LEA32r.  This is analogous to a
+/// specific material being used to carry 12V in a circuit.  It is still
+/// possible that a 5V input signal could be connected to a 12V-capable trace,
+/// but a 12V input signal will never be connected to a 5V trace.  It is
+/// possible to determine which inputs may be connected to 12V input signals
+/// just by looking at a particular circuit block without needing to know how
+/// all replicas of it will be connected in the whole circuit.  Similarly, it is
+/// possible to determine which register inputs to a particular instruction may
+/// point to stack data by analyzing just the instructions preceding that one in
+/// its basic block.
+///
+/// Ultimately, it is still necessary to analyze the control flow between basic
+/// blocks to make sure that for each base address register used in each memory
+/// operand, the register consistently refers to either stack data or non-stack
+/// data for every control flow that contains the corresponding instruction.
+/// This is analogous to checking circuit blocks to make sure that for each
+/// 12V-capable trace leading to each circuit element, all replicas are
+/// connected to 12V input signals or all replicas are connected to 5V input
+/// signals.  For efficiency, this pass only traverses each basic block once.
+/// The connectivity check is performed by recording the predecessors of each
+/// basic block and tracing a particular register value back through the chain
+/// of derivations that produced it to check the property described above.  This
+/// is analogous to tracing a wire connected to a circuit element back to its
+/// source.
+///
+/// This pass should always be combined with the SafeStack instrumentation pass
+/// with the -safe-stack-subr-acc-as-unsafe option, since the SafeStack pass
+/// will then move any objects to the unsafe stack that could have pointers to
+/// them passed between functions.  This is necessary, since the
+/// X86FixupSeparateStack pass only performs intra-procedural analysis.  The
+/// runtime can place the unsafe stack outside of the stack segment, so that
+/// accesses to it do not require segment override prefixes.
+///
+/// Even with the SafeStack pass enabled, va_list objects containing pointers to
+/// the stack may still be passed between C/C++ functions.  That is why the
+/// Clang frontend recognizes when this pass is enabled and adds the appropriate
+/// segment override prefixes to __builtin_va_arg invocations.
+///
+class X86FixupSeparateStack : public MachineFunctionPass {
+public:
+  X86FixupSeparateStack() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  const char *getPassName() const override {
+    return "X86 Fixup Separate Stack";
+  }
+
+  const TargetInstrInfo *TII;
+  const X86RegisterInfo *TRI;
+  const X86Subtarget *STI;
+  static char ID;
+
+  /// Map from basic block to corresponding AddrRegReqs object
+  std::map<MachineBasicBlock *, std::shared_ptr<AddrRegReqs> > RegReqs;
+
+  /// \brief Perform fixups for a string instruction.
+  ///
+  /// \param I Iterator pointing to the string instruction in the basic block.
+  ///          This routine may add new instructions around the string
+  ///          instruction, in which case 'I' will point to the instruction that
+  ///          was added at the highest address.
+  /// \param StringDS true if the instruction accesses DS
+  /// \param StringES true if the instruction accesses ES
+  ///
+  /// \returns true if the basic block was changed
+  bool processStringInstr(MachineBasicBlock *BB, MachineBasicBlock::iterator &I,
+                          AddrRegReqs& Reqs, bool StringDS, bool StringES);
+
+  /// \brief Perform fixups for an instruction that accesses memory with an
+  /// explicit memory operand.
+  ///
+  /// \param MemoryOperand Offset to the memory operand within the instruction
+  ///
+  /// \returns true if the instruction was changed
+  bool processMemAccessInstr(MachineInstr &I, int MemoryOperand,
+                             AddrRegReqs& Reqs);
+
+  /// \brief Check whether an instruction is of a type that may be used to
+  /// compute a derived stack pointer.
+  ///
+  /// \returns X86::NoRegister if the instruction is not of a type that may be
+  /// used to computes a derived stack pointer.  Otherwise, returns the register
+  /// that is used as the source of the possible derivation.
+  unsigned mayComputeStackPtr(const MachineInstr &I) const;
+
+  bool processBasicBlock(MachineBasicBlock *BB,
+                         std::shared_ptr<AddrRegReqs> Reqs);
+};
+
+char X86FixupSeparateStack::ID = 0;
+}
+
+FunctionPass *llvm::createX86FixupSeparateStack() {
+  return new X86FixupSeparateStack();
+}
+
+void AddrRegs::insertStackReg(unsigned Reg) {
+  for (auto N : Stack)
+    if (TRI->regsOverlap(N, Reg))
+      return;
+
+  Stack.insert(Reg);
+}
+
+bool AddrRegs::pointsIntoStack(unsigned Reg) const {
+  if (!TRI->isPhysicalRegister(Reg))
+    return false;
+
+  if (TRI->regsOverlap(TRI->getStackRegister(), Reg))
+    // The stack register is assumed to always point to the stack
+    return true;
+
+  for (auto N : Stack)
+    if (TRI->regsOverlap(N, Reg))
+      return true;
+
+  return false;
+}
+
+AddrRegReqs::AddrRegReqs(AddrRegReqs *Predecessor) :
+  TRI(Predecessor->TRI), OrigRegs(Predecessor->extractRegs()) {
+  Predecessors.insert(Predecessor);
+}
+
+AddrRegs AddrRegReqs::extractRegs() {
+  AddrRegs Regs(TRI);
+
+  // Registers that have been redefined since the beginning of the basic block
+  DenseSet<unsigned> Redefined;
+
+  // Handle the explicit derivations:
+  for (auto &D : Derivatives) {
+    // Note that this derivation overrides the corresponding implicit one:
+    Redefined.insert(D.first);
+    if (OrigRegs.pointsIntoStack(D.second))
+      Regs.insertStackReg(D.first);
+  }
+
+  // Handle the implicit derivations:
+  for (auto N : OrigRegs.Stack)
+    if (Redefined.find(N) == Redefined.end())
+      Regs.insertStackReg(N);
+
+  return Regs;
+}
+
+void AddrRegReqs::addPredecessor(AddrRegReqs *Predecessor) {
+  if (!Predecessors.insert(Predecessor).second)
+    // This predecessor was already registered
+    return;
+
+  for (auto N : UsedInMemOp)
+    Predecessor->useInMemOp(N, OrigRegs.pointsIntoStack(N));
+}
+
+unsigned AddrRegReqs::lookupRoot(unsigned To) const {
+  if (!TRI->isPhysicalRegister(To))
+    return X86::NoRegister;
+
+  for (auto &Derivative : Derivatives)
+    if (TRI->regsOverlap(Derivative.first, To))
+      return Derivative.second;
+
+  return To;
+}
+
+void AddrRegReqs::derive(unsigned From, unsigned To) {
+  if (TRI->regsOverlap(To, TRI->getStackRegister()))
+    // The stack register is always assumed to point to the stack, regardless
+    // of what instructions are used to update it.
+    return;
+
+  // Lookup the root register from which From is derived, since the second
+  // identifier in each pair in Derivatives is actually for the register value
+  // at the beginning of the basic block from which the value for the register
+  // identified by the first number is derived.
+  unsigned RootFrom = lookupRoot(From);
+
+  for (auto &Derivative : Derivatives)
+    if (TRI->regsOverlap(Derivative.first, To)) {
+      // Update an existing entry
+      Derivative.second = RootFrom;
+      return;
+    }
+
+  // Replace the implicit derivation for To with the following explicit
+  // derivation.
+  Derivatives.push_back(std::make_pair(To, RootFrom));
+}
+
+void AddrRegReqs::markNotDerived(unsigned To) {
+  derive(X86::NoRegister, To);
+}
+
+void AddrRegReqs::useInMemOp(unsigned Reg, bool IsStackReg) {
+  bool FromStackReg = false;
+
+  unsigned RootFrom = lookupRoot(Reg);
+  if (RootFrom != X86::NoRegister)
+    FromStackReg = OrigRegs.pointsIntoStack(RootFrom);
+
+  // For registers that are set by instructions of the types used by the
+  // compiler to compute pointers to stack data, this pass checks that
+  // each such register either always points to stack data or never points
+  // to stack data at each instruction that uses the register in a memory
+  // operand.
+  assert(IsStackReg == FromStackReg &&
+         "One or more registers point to both stack and data locations at one "
+         "or more instructions, which is not supported by the "
+         "X86FixupSeparateStack pass.");
+
+  if (RootFrom == X86::NoRegister)
+    // This register cannot point to stack data, so it is unnecessary to perform
+    // consistency checks on its values.
+    return;
+
+  for (auto N : UsedInMemOp)
+    if (TRI->regsOverlap(RootFrom, N))
+      return;
+
+  UsedInMemOp.insert(RootFrom);
+
+  for (AddrRegReqs *P : Predecessors)
+    P->useInMemOp(RootFrom, IsStackReg);
+}
+
+/// \brief Return true iff Reg points to a stack location.
+bool AddrRegReqs::pointsIntoStack(unsigned Reg) const {
+  unsigned From = lookupRoot(Reg);
+  if (From == X86::NoRegister)
+    return false;
+
+  return OrigRegs.pointsIntoStack(From);
+}
+
+bool X86FixupSeparateStack::processStringInstr(MachineBasicBlock *BB,
+                                               MachineBasicBlock::iterator &I,
+                                               AddrRegReqs &Reqs,
+                                               bool StringDS, bool StringES) {
+  MachineFunction &MF = *BB->getParent();
+
+  // FIXME: Check for a segment override prefix on the string instruction and
+  // skip the following if one is detected.
+
+  bool StackDS = false;
+  bool StackES = false;
+
+  if (StringDS) {
+    StackDS = Reqs.pointsIntoStack(X86::ESI);
+    Reqs.useInMemOp(X86::ESI, StackDS);
+  }
+
+  if (StringES) {
+    StackES = Reqs.pointsIntoStack(X86::EDI);
+    Reqs.useInMemOp(X86::EDI, StackES);
+  }
+
+  if (!(StackDS || StackES))
+    return false;
+
+  // Emit instructions to save a segment register, overwrite it with the
+  // contents of SS, and restore its original value after the string instruction
+  // is complete.
+  auto tempOverwriteSegWithSS = [&](unsigned PushOp, unsigned PopOp) {
+    // Save the original value of the segment register to be overwritten.
+    MachineInstr *NewMI =
+      MF.CreateMachineInstr(TII->get(PushOp), I->getDebugLoc());
+    I = BB->insert(I, NewMI);
+    // Push the stack segment selector onto the stack.
+    NewMI = MF.CreateMachineInstr(TII->get(X86::PUSHSS32), I->getDebugLoc());
+    I = BB->insertAfter(I, NewMI);
+    // Overwrite the segment register with the stack segment selector.
+    NewMI = MF.CreateMachineInstr(TII->get(PopOp), I->getDebugLoc());
+    I = BB->insertAfter(I, NewMI);
+    // Advance to the string instruction.
+    I++;
+    // Restore the original value of the segment register.
+    NewMI = MF.CreateMachineInstr(TII->get(PopOp), I->getDebugLoc());
+    return BB->insertAfter(I, NewMI);
+  };
+
+  if (StackDS) {
+    tempOverwriteSegWithSS(X86::PUSHDS32, X86::POPDS32);
+  }
+  if (StackES) {
+    if (StackDS)
+      // Back up to the string instruction.
+      I--;
+    tempOverwriteSegWithSS(X86::PUSHES32, X86::POPES32);
+    if (StackDS)
+      // Advance past the POP instruction that restores DS.
+      I++;
+  }
+
+  return true;
+}
+
+bool X86FixupSeparateStack::processMemAccessInstr(MachineInstr &I,
+                                                  int MemoryOperand,
+                                                  AddrRegReqs &Reqs) {
+  MachineOperand &SegRegOp = I.getOperand(MemoryOperand + X86::AddrSegmentReg);
+
+  unsigned PrevSegReg = SegRegOp.getReg();
+
+  if (TRI->isPhysicalRegister(PrevSegReg))
+    // Do not replace the existing segment override prefix.
+    return false;
+
+  unsigned BaseReg = I.getOperand(MemoryOperand + X86::AddrBaseReg).getReg();
+
+  bool BaseIsPhysReg = TRI->isPhysicalRegister(BaseReg);
+
+  bool InStack = Reqs.pointsIntoStack(BaseReg);
+
+  if (BaseIsPhysReg)
+    Reqs.useInMemOp(BaseReg, InStack);
+
+  if (InStack &&
+       (TRI->regsOverlap(X86::ESP, BaseReg) ||
+        TRI->regsOverlap(X86::EBP, BaseReg)))
+      // Memory operand with a base register of ESP or EBP implicitly references
+      // SS.
+    return false;
+
+  if (!InStack &&
+      // Memory operand without a base register implicitly
+      // references DS.
+      (!(BaseIsPhysReg &&
+         // Memory operand with a base register other than ESP and EBP implicitly
+         // references DS.  This assumes that ESP is never used as a non-stack
+         // pointer.
+         TRI->regsOverlap(X86::EBP, BaseReg))))
+    return false;
+
+  SegRegOp.ChangeToRegister(InStack ? X86::SS : X86::DS, false);
+
+  return true;
+}
+
+unsigned X86FixupSeparateStack::mayComputeStackPtr(const MachineInstr &I) const {
+  switch (I.getOpcode()) {
+  // This assumes that only the following instructions are used by the compiler
+  // to compute a stack pointer derived directly or indirectly from ESP:
+  case X86::LEA32r:
+  case X86::MOV32rr:
+    break;
+  default:
+    return X86::NoRegister;
+  }
+
+  for (const MachineOperand &Op : I.operands()) {
+    if (!Op.isReg())
+      continue;
+
+    if (!Op.isUse())
+      continue;
+
+    // Return the register that is used in the address computation.
+    return Op.getReg();
+  }
+
+  assert(false &&
+         "Unexpected failure to identify register used to compute address");
+}
+
+bool X86FixupSeparateStack::processBasicBlock(MachineBasicBlock *BB,
+                                              std::shared_ptr<AddrRegReqs> Reqs) {
+  bool Changed = false;
+
+  RegReqs.emplace(BB, Reqs);
+
+  for (MachineBasicBlock::iterator I = BB->begin(); I != BB->end(); ++I) {
+    bool StringDS, StringES;
+
+    unsigned Opcode = I->getOpcode();
+
+    switch (Opcode) {
+    case X86::REP_MOVSB_32:
+    case X86::REP_MOVSW_32:
+    case X86::REP_MOVSD_32:
+      StringDS = true;
+      StringES = true;
+      break;
+    case X86::REP_STOSB_32:
+    case X86::REP_STOSW_32:
+    case X86::REP_STOSD_32:
+      StringDS = false;
+      StringES = true;
+      break;
+    default:
+      StringDS = false;
+      StringES = false;
+      break;
+    }
+
+    if (StringDS || StringES) {
+      Changed |= processStringInstr(BB, I, *Reqs, StringDS, StringES);
+
+      continue;
+    }
+
+    const MCInstrDesc &Desc = TII->get(Opcode);
+    uint64_t TSFlags = Desc.TSFlags;
+
+    int MemoryOperand = -1;
+    if (I->mayLoadOrStore()) {
+      // Determine where the memory operand starts, if present.
+      MemoryOperand = X86II::getMemoryOperandNo(TSFlags, Opcode);
+      if (MemoryOperand != -1)
+        MemoryOperand += X86II::getOperandBias(Desc);
+    }
+
+    bool HasEVEX_K = TSFlags & X86II::EVEX_K;
+    bool HasVEX_4V = TSFlags & X86II::VEX_4V;
+
+    // This assumes that the compiler only uses the following forms of
+    // instructions to store stack pointers (including derived ones) to memory.
+    // Certain VEX-encoded instructions are specifically excluded to reduce the
+    // complexity of the subsequent code that computes the operand number of the
+    // base register.  See X86MCCodeEmitter::EmitVEXOpcodePrefix for information
+    // about how these instructions are represented in LLVM.  Again, this
+    // assumes that such instructions are not used to store stack pointers.
+    MachineOperand *StackPtrStoreSrcOp = NULL;
+    if (Opcode == X86::PUSH32r) {
+      StackPtrStoreSrcOp = &(I->getOperand(0));
+    } else if (((TSFlags & X86II::FormMask) == X86II::MRMDestMem) &&
+               !(HasEVEX_K || HasVEX_4V)) {
+      // This assumes that X86II::getMemoryOperandNo always returns 0 for
+      // instructions of this form.  That is why there is no assertion preceding
+      // this statement to check that MemoryOperand != -1.
+      unsigned StackPtrStoreBaseReg =
+        I->getOperand(MemoryOperand + X86::AddrBaseReg).getReg();
+
+      // Only check stores to the stack, since the SafeStack instrumentation
+      // pass stores values derived from the stack pointer into other memory
+      // locations, and this pass assumes that those values will only be loaded
+      // back into the stack pointer later and do not need to be tracked.  This
+      // pass assumes that the SafeStack instrumentation pass will not load
+      // those values into other registers that are then used as base registers
+      // in memory operands.
+      if (Reqs->pointsIntoStack(StackPtrStoreBaseReg))
+        StackPtrStoreSrcOp =
+          &(I->getOperand(MemoryOperand + X86::AddrNumOperands));
+    }
+
+    if ((StackPtrStoreSrcOp != NULL) &&
+        Reqs->pointsIntoStack(StackPtrStoreSrcOp->getReg()))
+      DEBUG(dbgs() << "The X86FixupSeparateStack pass is unable to track stack "
+                      "pointers that are stored to memory.\n");
+
+    unsigned MayDeriveStackPtrFrom = X86::NoRegister;
+
+    if (MemoryOperand != -1) {
+      Changed |= processMemAccessInstr(*I, MemoryOperand, *Reqs);
+    } else {
+      MayDeriveStackPtrFrom = mayComputeStackPtr(*I);
+    }
+
+    for (MachineOperand &Op : I->operands()) {
+      if (!Op.isReg() || !Op.isDef())
+        continue;
+
+      unsigned Reg = Op.getReg();
+
+      if (!TRI->isPhysicalRegister(Reg))
+        continue;
+
+      if (MayDeriveStackPtrFrom != X86::NoRegister) {
+        Reqs->derive(MayDeriveStackPtrFrom, Reg);
+        continue;
+      }
+
+      Reqs->markNotDerived(Reg);
+    }
+  }
+
+  for (auto Succ : BB->successors()) {
+    auto RRI = RegReqs.find(Succ);
+    if (RRI == RegReqs.end()) {
+      std::shared_ptr<AddrRegReqs> SuccReqs(new AddrRegReqs(Reqs.get()));
+      Changed |= processBasicBlock(Succ, SuccReqs);
+    } else
+      RRI->second->addPredecessor(Reqs.get());
+  }
+
+  return Changed;
+}
+
+bool X86FixupSeparateStack::runOnMachineFunction(MachineFunction &MF) {
+  STI = &MF.getSubtarget<X86Subtarget>();
+  TII = STI->getInstrInfo();
+  TRI = STI->getRegisterInfo();
+
+  if (!STI->useSeparateStackSeg())
+    return false;
+
+  assert(!STI->is64Bit() && "Only X86-32 is supported.");
+
+  RegReqs.clear();
+
+  std::shared_ptr<AddrRegReqs> InitReqs(new AddrRegReqs(TRI));
+  return processBasicBlock(&MF.front(), InitReqs);
+}

lib/Target/X86/X86Subtarget.h

protected:
   bool HasPCOMMIT;
 
   /// Processor supports Cache Line Write Back instruction
   bool HasCLWB;
 
   /// Use software floating point for code generation.
   bool UseSoftFloat;
 
+  /// Use a separate stack segment
+  bool UseSeparateStackSeg;
+
   /// The minimum alignment known to hold of the stack frame on
   /// entry to the function and which must be maintained by every function.
   unsigned stackAlignment;
 
   /// Max. memset / memcpy size that is turned into rep/movs, rep/stos ops.
   ///
   unsigned MaxInlineSizeThreshold;
 
public:
   bool hasBWI() const { return HasBWI; }
   bool hasVLX() const { return HasVLX; }
   bool hasPKU() const { return HasPKU; }
   bool hasMPX() const { return HasMPX; }
 
   bool isAtom() const { return X86ProcFamily == IntelAtom; }
   bool isSLM() const { return X86ProcFamily == IntelSLM; }
   bool useSoftFloat() const { return UseSoftFloat; }
+  bool useSeparateStackSeg() const { return UseSeparateStackSeg; }
 
   /// Use mfence if we have SSE2 or we're on x86-64 (even if we asked for
   /// no-sse2). There isn't any reason to disable it if the target processor
   /// supports it.
   bool hasMFence() const { return hasSSE2() || is64Bit(); }
 
   const Triple &getTargetTriple() const { return TargetTriple; }
 

lib/Target/X86/X86Subtarget.cpp

void X86Subtarget::initializeEnvironment() {
   CallRegIndirect = false;
   LEAUsesAG = false;
   SlowLEA = false;
   SlowIncDec = false;
   stackAlignment = 4;
   // FIXME: this is a known good value for Yonah. How about others?
   MaxInlineSizeThreshold = 128;
   UseSoftFloat = false;
+  UseSeparateStackSeg = false;
 }
 
 X86Subtarget &X86Subtarget::initializeSubtargetDependencies(StringRef CPU,
                                                             StringRef FS) {
   initializeEnvironment();
   initSubtargetFeatures(CPU, FS);
   return *this;
 }

lib/Target/X86/X86TargetMachine.cpp

void X86PassConfig::addPreEmitPass() {
   if (UseVZeroUpper)
     addPass(createX86IssueVZeroUpperPass());
 
   if (getOptLevel() != CodeGenOpt::None) {
     addPass(createX86FixupBWInsts());
     addPass(createX86PadShortFunctions());
     addPass(createX86FixupLEAs());
   }
+
+  addPass(createX86FixupSeparateStack());
 }