[SLH] AArch64: correctly pick temporary register to mask SP

[kristof.beyls]

Summary

As part of speculation hardening, the stack pointer gets masked with the
taint register (X16) before a function call or before a function return.
Since there are no instructions that can directly mask writing to the
stack pointer, the stack pointer must first be transferred to another
register, where it can be masked, before that value is transferred back
to the stack pointer.
Before, that temporary register was always picked to be x17, since the
ABI allows clobbering x17 on any function call, resulting in the
following instruction pattern being inserted before function calls and
returns/tail calls:

mov x17, sp
and x17, x17, x16
mov sp, x17

However, x17 can be live in those locations, for example when the call
is an indirect call, using x17 as the target address (blr x17).

To fix this, this patch looks for an available register just before the
call or terminator instruction and uses that.

In the rare case when no register turns out to be available (this
situation is only encountered twice across the whole test-suite), just
insert a full speculation barrier at the start of the basic block where
this occurs.

[olista01]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

232 ↗

(On Diff #181773)

I think this requires a return instruction, not just any terminator, because terminators could have other values live across them. I'm also not sure that it's correct for calls, which could have registers live across them, without using them directly.

Maybe we could use the RegScavenger here, which looks like it can also handle some more difficult cases by spilling a register to the stack?

252 ↗

(On Diff #181773)

This can be hit by this code:

typedef int (*fptr_t)(void);
int foo(fptr_t f) {
  register fptr_t f2 asm("x30") = f;
  asm("" : "+r"(f2));
  return f2() + 1;
}

This uses inline assembly to force the function pointer into x30, but it is an allocatable register so I think it is possible for this to happen "naturally".

[chandlerc]

Adding Zola as she's been looking at ARM's SLH as well...

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

232 ↗

(On Diff #181773)

Yeah, you're right probably better to reuse RegScavenger than trying to reinvent the wheel here.
I am not fully sure if introducing new loads and stores in this pass which aims to protect loads and stores is a good idea, will try to think about that.
Even if we can't find a free register without introducing spill/fills, in that case, we could still use the "emergency" ISB/DSB barrier to block speculation entirely in the affected basic block and not need to mask values in that basic block.
I'll experiment with such an approach.

252 ↗

(On Diff #181773)

Thanks for that example. I didn't manage to come up with an example myself.
My understanding is that the llvm_unreachable is reached because the ABI rules on the indirect call are not modelled very accurately.
Some of the registers will be clobbered by the call, and could be used as temporary register, but it seems this isn't modelled on the BLR mir instruction.
Anyway, as I said above - probably I'll need to implement generation of the ISB/DSB full speculation barrier in such a rare case.

[zbrid]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

253 ↗

(On Diff #181773)

Since I don't know much about the ABI or the assumptions around calls, I had some quick questions about this for clarification. Can you be sure that this is unreachable because the way the calling convention is specified means there will always be at least one dead register before a call? For example, since there are several call clobbered registers specified in the calling convention even if x17 is used as an address for an indirect branch instruction (and is therefore live across the call), there's no way for all of the call clobbered registers to be live since there are no call instructions that could force all of them to be live?

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

253 ↗

(On Diff #181773)

Yes, indeed, the assumption is that the calling convention will result in at least one register being "killed" during the call. Or at least not being guaranteed to retain the same value.
The ABI actually specifies that x17 can never be live across a call, which is why this was picked initially. However, as can be seen in the new test cases in this patch, x17 could still be legally be used as containing the address an indirect call has to branch to - making the register still live just before the call.

[kristof.beyls]

Updated approach to use the register scavenger to find an available temporary register. In the rare case when no temporary register is available, fall back to inserting a full speculation barrier.

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

252 ↗

(On Diff #181773)

The example you provided is now (further simplified and as MIR) in one of the regression tests.

253 ↗

(On Diff #181773)

The new version of the patch handles no temporary register being available by inserting a full speculation barrier.
Across all of the test suite, this case was encountered just twice.

[olista01]

LGTM with one nit.

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

307 ↗

(On Diff #182498)

GPR64allRegClass includes SP and XZR, I don't think it actually matters here because neither of them are allocatable, but it would be clearer to use GPR64common.

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

307 ↗

(On Diff #182498)

Nice catch - thanks for noticing.
I've followed your suggestion in the committed code.

llvm/trunk/lib/Target/AArch64/AArch64SpeculationHardening.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Target/TargetMachine.h"
 #include <cassert>
 
 using namespace llvm;
 

 private:
   unsigned MisspeculatingTaintReg;
   unsigned MisspeculatingTaintReg32Bit;
   bool UseControlFlowSpeculationBarrier;
   BitVector RegsNeedingCSDBBeforeUse;
   BitVector RegsAlreadyMasked;
 
   bool functionUsesHardeningRegister(MachineFunction &MF) const;
-  bool instrumentControlFlow(MachineBasicBlock &MBB);
+  bool instrumentControlFlow(MachineBasicBlock &MBB,
+                             bool &UsesFullSpeculationBarrier);
   bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                                MachineBasicBlock *&FBB,
                                AArch64CC::CondCode &CondCode) const;
   void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
                           AArch64CC::CondCode &CondCode, DebugLoc DL) const;
-  void insertSPToRegTaintPropagation(MachineBasicBlock *MBB,
+  void insertSPToRegTaintPropagation(MachineBasicBlock &MBB,
                                      MachineBasicBlock::iterator MBBI) const;
-  void insertRegToSPTaintPropagation(MachineBasicBlock *MBB,
+  void insertRegToSPTaintPropagation(MachineBasicBlock &MBB,
                                      MachineBasicBlock::iterator MBBI,
                                      unsigned TmpReg) const;
+  void insertFullSpeculationBarrier(MachineBasicBlock &MBB,
+                                    MachineBasicBlock::iterator MBBI,
+                                    DebugLoc DL) const;
 
   bool slhLoads(MachineBasicBlock &MBB);
   bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MBBI,
                               MachineInstr &MI, unsigned Reg);
-  bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB);
+  bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB,
+                                        bool UsesFullSpeculationBarrier);
   bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
-                                  MachineBasicBlock::iterator MBBI);
+                                  MachineBasicBlock::iterator MBBI,
+                                  bool UsesFullSpeculationBarrier);
   bool insertCSDB(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                   DebugLoc DL);
 };
 
 } // end anonymous namespace
 
 char AArch64SpeculationHardening::ID = 0;
 
bool AArch64SpeculationHardening::endsWithCondControlFlow(
 
   assert(MBB.succ_size() == 2);
   // translate analyzeBranchCondCode to CondCode.
   assert(analyzeBranchCondCode.size() == 1 && "unknown Cond array format");
   CondCode = AArch64CC::CondCode(analyzeBranchCondCode[0].getImm());
   return true;
 }
 
+void AArch64SpeculationHardening::insertFullSpeculationBarrier(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    DebugLoc DL) const {
+  // A full control flow speculation barrier consists of (DSB SYS + ISB)
+  BuildMI(MBB, MBBI, DL, TII->get(AArch64::DSB)).addImm(0xf);
+  BuildMI(MBB, MBBI, DL, TII->get(AArch64::ISB)).addImm(0xf);
+}
+
 void AArch64SpeculationHardening::insertTrackingCode(
     MachineBasicBlock &SplitEdgeBB, AArch64CC::CondCode &CondCode,
     DebugLoc DL) const {
   if (UseControlFlowSpeculationBarrier) {
-    // insert full control flow speculation barrier (DSB SYS + ISB)
-    BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::ISB))
-        .addImm(0xf);
-    BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::DSB))
-        .addImm(0xf);
+    insertFullSpeculationBarrier(SplitEdgeBB, SplitEdgeBB.begin(), DL);
   } else {
     BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::CSELXr))
         .addDef(MisspeculatingTaintReg)
         .addUse(MisspeculatingTaintReg)
         .addUse(AArch64::XZR)
         .addImm(CondCode);
     SplitEdgeBB.addLiveIn(AArch64::NZCV);
   }
 }
 
 bool AArch64SpeculationHardening::instrumentControlFlow(
-    MachineBasicBlock &MBB) {
+    MachineBasicBlock &MBB, bool &UsesFullSpeculationBarrier) {
   LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);
 
   bool Modified = false;
   MachineBasicBlock *TBB = nullptr;
   MachineBasicBlock *FBB = nullptr;
   AArch64CC::CondCode CondCode;
 
   if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {
if (!endsWithCondControlFlow(MBB, TBB, FBB, CondCode)) {
     insertTrackingCode(*SplitEdgeFBB, InvCondCode, DL);
 
     LLVM_DEBUG(dbgs() << "SplitEdgeTBB: " << *SplitEdgeTBB << "\n");
     LLVM_DEBUG(dbgs() << "SplitEdgeFBB: " << *SplitEdgeFBB << "\n");
     Modified = true;
   }
 
   // Perform correct code generation around function calls and before returns.
-  {
-    SmallVector<MachineInstr *, 4> ReturnInstructions;
-    SmallVector<MachineInstr *, 4> CallInstructions;
+  // The below variables record the return/terminator instructions and the call
+  // instructions respectively; including which register is available as a
+  // temporary register just before the recorded instructions.
+  SmallVector<std::pair<MachineInstr *, unsigned>, 4> ReturnInstructions;
+  SmallVector<std::pair<MachineInstr *, unsigned>, 4> CallInstructions;
+  // if a temporary register is not available for at least one of the
+  // instructions for which we need to transfer taint to the stack pointer, we
+  // need to insert a full speculation barrier.
+  // TmpRegisterNotAvailableEverywhere tracks that condition.
+  bool TmpRegisterNotAvailableEverywhere = false;
+
+  RegScavenger RS;
+  RS.enterBasicBlock(MBB);
+
+  for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); I++) {
+    MachineInstr &MI = *I;
+    if (!MI.isReturn() && !MI.isCall())
+      continue;
 
-    for (MachineInstr &MI : MBB) {
-      if (MI.isReturn())
-        ReturnInstructions.push_back(&MI);
-      else if (MI.isCall())
-        CallInstructions.push_back(&MI);
-    }
+    // The RegScavenger represents registers available *after* the MI
+    // instruction pointed to by RS.getCurrentPosition().
+    // We need to have a register that is available *before* the MI is executed.
+    if (I != MBB.begin())
+      RS.forward(std::prev(I));
+    // FIXME: The below just finds *a* unused register. Maybe code could be
+    // optimized more if this looks for the register that isn't used for the
+    // longest time around this place, to enable more scheduling freedom. Not
+    // sure if that would actually result in a big performance difference
+    // though. Maybe RegisterScavenger::findSurvivorBackwards has some logic
+    // already to do this - but it's unclear if that could easily be used here.
+    unsigned TmpReg = RS.FindUnusedReg(&AArch64::GPR64commonRegClass);
+    LLVM_DEBUG(dbgs() << "RS finds "
+                      << ((TmpReg == 0) ? "no register " : "register ");
+               if (TmpReg != 0) dbgs() << printReg(TmpReg, TRI) << " ";
+               dbgs() << "to be available at MI " << MI);
+    if (TmpReg == 0)
+      TmpRegisterNotAvailableEverywhere = true;
+      ReturnInstructions.push_back({&MI, TmpReg});
+      CallInstructions.push_back({&MI, TmpReg});
 
-    Modified |=
-        (ReturnInstructions.size() > 0) || (CallInstructions.size() > 0);
+  if (TmpRegisterNotAvailableEverywhere) {
+    // When a temporary register is not available everywhere in this basic
+    // basic block where a propagate-taint-to-sp operation is needed, just
+    // emit a full speculation barrier at the start of this basic block, which
+    // renders the taint/speculation tracking in this basic block unnecessary.
+    insertFullSpeculationBarrier(MBB, MBB.begin(),
+                                 (MBB.begin())->getDebugLoc());
+    UsesFullSpeculationBarrier = true;
+    Modified = true;
+  } else {
+    for (auto MI_Reg : ReturnInstructions) {
+      assert(MI_Reg.second != 0);
+      LLVM_DEBUG(
+          dbgs()
+          << " About to insert Reg to SP taint propagation with temp register "
+          << printReg(MI_Reg.second, TRI)
+          << " on instruction: " << *MI_Reg.first);
+      insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
+      Modified = true;
+    }
 
-    for (MachineInstr *Return : ReturnInstructions)
-      insertRegToSPTaintPropagation(Return->getParent(), Return, AArch64::X17);
-    for (MachineInstr *Call : CallInstructions) {
+    for (auto MI_Reg : CallInstructions) {
+      assert(MI_Reg.second != 0);
+      LLVM_DEBUG(dbgs() << " About to insert Reg to SP and back taint "
+                           "propagation with temp register "
+                        << printReg(MI_Reg.second, TRI)
+                        << " around instruction: " << *MI_Reg.first);
       // Just after the call:
-      MachineBasicBlock::iterator i = Call;
-      i++;
-      insertSPToRegTaintPropagation(Call->getParent(), i);
+      insertSPToRegTaintPropagation(
+          MBB, std::next((MachineBasicBlock::iterator)MI_Reg.first));
       // Just before the call:
-      insertRegToSPTaintPropagation(Call->getParent(), Call, AArch64::X17);
+      insertRegToSPTaintPropagation(MBB, MI_Reg.first, MI_Reg.second);
+      Modified = true;
     }
   }
-
   return Modified;
 }
 
 void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
-    MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI) const {
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) const {
   // If full control flow speculation barriers are used, emit a control flow
   // barrier to block potential miss-speculation in flight coming in to this
   // function.
   if (UseControlFlowSpeculationBarrier) {
-    // insert full control flow speculation barrier (DSB SYS + ISB)
-    BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::DSB)).addImm(0xf);
-    BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ISB)).addImm(0xf);
+    insertFullSpeculationBarrier(MBB, MBBI, DebugLoc());
     return;
   }
 
   // CMP   SP, #0   === SUBS   xzr, SP, #0
-  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
+  BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::SUBSXri))
       .addDef(AArch64::XZR)
       .addUse(AArch64::SP)
       .addImm(0)
       .addImm(0); // no shift
   // CSETM x16, NE  === CSINV  x16, xzr, xzr, EQ
-  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
+  BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::CSINVXr))
       .addDef(MisspeculatingTaintReg)
       .addUse(AArch64::XZR)
       .addUse(AArch64::XZR)
       .addImm(AArch64CC::EQ);
 }
 
 void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
-    MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI,
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
     unsigned TmpReg) const {
   // If full control flow speculation barriers are used, there will not be
   // miss-speculation when returning from this function, and therefore, also
   // no need to encode potential miss-speculation into the stack pointer.
   if (UseControlFlowSpeculationBarrier)
     return;
 
   // mov   Xtmp, SP  === ADD  Xtmp, SP, #0
-  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
+  BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
       .addDef(TmpReg)
       .addUse(AArch64::SP)
       .addImm(0)
       .addImm(0); // no shift
   // and   Xtmp, Xtmp, TaintReg === AND Xtmp, Xtmp, TaintReg, #0
-  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
+  BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ANDXrs))
       .addDef(TmpReg, RegState::Renamable)
       .addUse(TmpReg, RegState::Kill | RegState::Renamable)
       .addUse(MisspeculatingTaintReg, RegState::Kill)
       .addImm(0);
   // mov   SP, Xtmp === ADD SP, Xtmp, #0
-  BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
+  BuildMI(MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
       .addDef(AArch64::SP)
       .addUse(TmpReg, RegState::Kill)
       .addImm(0)
       .addImm(0); // no shift
 }
 
 bool AArch64SpeculationHardening::functionUsesHardeningRegister(
     MachineFunction &MF) const {
if (HardenAddressLoadedFrom)
       }
   }
   return Modified;
 }
 
 /// \brief If MBBI references a pseudo instruction that should be expanded
 /// here, do the expansion and return true. Otherwise return false.
 bool AArch64SpeculationHardening::expandSpeculationSafeValue(
-    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
+    bool UsesFullSpeculationBarrier) {
   MachineInstr &MI = *MBBI;
   unsigned Opcode = MI.getOpcode();
   bool Is64Bit = true;
 
   switch (Opcode) {
   default:
     break;
   case AArch64::SpeculationSafeValueW:
     Is64Bit = false;
     LLVM_FALLTHROUGH;
   case AArch64::SpeculationSafeValueX:
     // Just remove the SpeculationSafe pseudo's if control flow
     // miss-speculation isn't happening because we're already inserting barriers
     // to guarantee that.
-    if (!UseControlFlowSpeculationBarrier) {
+    if (!UseControlFlowSpeculationBarrier && !UsesFullSpeculationBarrier) {
       unsigned DstReg = MI.getOperand(0).getReg();
       unsigned SrcReg = MI.getOperand(1).getReg();
       // Mark this register and all its aliasing registers as needing to be
       // value speculation hardened before its next use, by using a CSDB
       // barrier instruction.
       for (MachineOperand Op : MI.defs())
         for (MCRegAliasIterator AI(Op.getReg(), TRI, true); AI.isValid(); ++AI)
           RegsNeedingCSDBBeforeUse.set(*AI);
assert(!UseControlFlowSpeculationBarrier && "No need to insert CSDBs when "
                                               "is already blocked");
   // insert data value speculation barrier (CSDB)
   BuildMI(MBB, MBBI, DL, TII->get(AArch64::HINT)).addImm(0x14);
   RegsNeedingCSDBBeforeUse.reset();
   return true;
 }
 
 bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudos(
-    MachineBasicBlock &MBB) {
+    MachineBasicBlock &MBB, bool UsesFullSpeculationBarrier) {
   bool Modified = false;
 
   RegsNeedingCSDBBeforeUse.reset();
 
   // The following loop iterates over all instructions in the basic block,
   // and performs 2 operations:
   // 1. Insert a CSDB at this location if needed.
   // 2. Expand the SpeculationSafeValuePseudo if the current instruction is
if (RegsNeedingCSDBBeforeUse.any() && (MI.isCall() || MI.isTerminator()))
       NeedToEmitBarrier = true;
     if (!NeedToEmitBarrier)
       for (MachineOperand Op : MI.uses())
         if (Op.isReg() && RegsNeedingCSDBBeforeUse[Op.getReg()]) {
           NeedToEmitBarrier = true;
           break;
         }
 
-    if (NeedToEmitBarrier)
+    if (NeedToEmitBarrier && !UsesFullSpeculationBarrier)
       Modified |= insertCSDB(MBB, MBBI, DL);
 
-    Modified |= expandSpeculationSafeValue(MBB, MBBI);
+    Modified |=
+        expandSpeculationSafeValue(MBB, MBBI, UsesFullSpeculationBarrier);
 
     MBBI = NMBBI;
   }
 
-  if (RegsNeedingCSDBBeforeUse.any())
+  if (RegsNeedingCSDBBeforeUse.any() && !UsesFullSpeculationBarrier)
     Modified |= insertCSDB(MBB, MBBI, DL);
 
   return Modified;
 }
 
 bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
   if (!MF.getFunction().hasFnAttribute(Attribute::SpeculativeLoadHardening))
     return false;
bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
   if (HardenLoads) {
     LLVM_DEBUG(
         dbgs() << "***** AArch64SpeculationHardening - automatic insertion of "
                   "SpeculationSafeValue intrinsics *****\n");
     for (auto &MBB : MF)
       Modified |= slhLoads(MBB);
   }
 
-  // 2.a Add instrumentation code to function entry and exits.
+  // 2. Add instrumentation code to function entry and exits.
   LLVM_DEBUG(
       dbgs()
       << "***** AArch64SpeculationHardening - track control flow *****\n");
 
   SmallVector<MachineBasicBlock *, 2> EntryBlocks;
   EntryBlocks.push_back(&MF.front());
   for (const LandingPadInfo &LPI : MF.getLandingPads())
     EntryBlocks.push_back(LPI.LandingPadBlock);
   for (auto Entry : EntryBlocks)
     insertSPToRegTaintPropagation(
-        Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
+        *Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
 
-  // 2.b Add instrumentation code to every basic block.
-  for (auto &MBB : MF)
-    Modified |= instrumentControlFlow(MBB);
-
-  LLVM_DEBUG(dbgs() << "***** AArch64SpeculationHardening - Lowering "
-                       "SpeculationSafeValue Pseudos *****\n");
-  // Step 3: Lower SpeculationSafeValue pseudo instructions.
-  for (auto &MBB : MF)
-    Modified |= lowerSpeculationSafeValuePseudos(MBB);
+  // 3. Add instrumentation code to every basic block.
+  for (auto &MBB : MF) {
+    bool UsesFullSpeculationBarrier = false;
+    Modified |= instrumentControlFlow(MBB, UsesFullSpeculationBarrier);
+    Modified |=
+        lowerSpeculationSafeValuePseudos(MBB, UsesFullSpeculationBarrier);
+  }
 
   return Modified;
 }
 
 /// \brief Returns an instance of the pseudo instruction expansion pass.
 FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
   return new AArch64SpeculationHardening();
 }

llvm/trunk/test/CodeGen/AArch64/speculation-hardening-loads.ll

 ; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --dump-input-on-failure
 
 define i128 @ldp_single_csdb(i128* %p) speculative_load_hardening {
 entry:
   %0 = load i128, i128* %p, align 16
   ret i128 %0
 ; CHECK-LABEL: ldp_single_csdb
 ; CHECK:      ldp   x8, x1, [x0]
 ; CHECK-NEXT: cmp sp, #0
 ; CHECK-NEXT: csetm x16, ne
 ; CHECK-NEXT: and   x8, x8, x16
 ; CHECK-NEXT: and   x1, x1, x16
 ; CHECK-NEXT: csdb
-; CHECK-NEXT: mov x17, sp
-; CHECK-NEXT: and x17, x17, x16
+; CHECK-NEXT: mov [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT: and [[TMPREG]], [[TMPREG]], x16
 ; CHECK-NEXT: mov x0, x8
-; CHECK-NEXT: mov sp, x17
+; CHECK-NEXT: mov sp, [[TMPREG]]
 ; CHECK-NEXT: ret
 }
 
 define double @ld_double(double* %p) speculative_load_hardening {
 entry:
   %0 = load double, double* %p, align 8
   ret double %0
 ; Checking that the address laoded from is masked for a floating point load.
 ; CHECK-LABEL: ld_double
 ; CHECK:      cmp sp, #0
 ; CHECK-NEXT: csetm x16, ne
 ; CHECK-NEXT: and   x0, x0, x16
 ; CHECK-NEXT: csdb
 ; CHECK-NEXT: ldr   d0, [x0]
-; CHECK-NEXT: mov x17, sp
-; CHECK-NEXT: and x17, x17, x16
-; CHECK-NEXT: mov sp, x17
+; CHECK-NEXT: mov [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT: and [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT: mov sp, [[TMPREG]]
 ; CHECK-NEXT: ret
 }
 
 define i32 @csdb_emitted_for_subreg_use(i64* %p, i32 %b) speculative_load_hardening {
 entry:
   %X = load i64, i64* %p, align 8
   %X_trunc = trunc i64 %X to i32
   %add = add i32 %b, %X_trunc
   %iszero = icmp eq i64 %X, 0
   %ret = select i1 %iszero, i32 %b, i32 %add
   ret i32 %ret
 ; Checking that the address laoded from is masked for a floating point load.
 ; CHECK-LABEL: csdb_emitted_for_subreg_use
 ; CHECK:      ldr x8, [x0]
 ; CHECK-NEXT: cmp sp, #0
 ; CHECK-NEXT: csetm x16, ne
 ; CHECK-NEXT: and x8, x8, x16
 ; csdb instruction must occur before the add instruction with w8 as operand.
 ; CHECK-NEXT: csdb
-; CHECK-NEXT: mov x17, sp
 ; CHECK-NEXT: add w9, w1, w8
 ; CHECK-NEXT: cmp x8, #0
-; CHECK-NEXT: and x17, x17, x16
 ; CHECK-NEXT: csel w0, w1, w9, eq
-; CHECK-NEXT: mov sp, x17
+; CHECK-NEXT: mov [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT: and [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT: mov sp, [[TMPREG]]
 ; CHECK-NEXT: ret
 }
 
 define i64 @csdb_emitted_for_superreg_use(i32* %p, i64 %b) speculative_load_hardening {
 entry:
   %X = load i32, i32* %p, align 4
   %X_ext = zext i32 %X to i64
   %add = add i64 %b, %X_ext
   %iszero = icmp eq i32 %X, 0
   %ret = select i1 %iszero, i64 %b, i64 %add
   ret i64 %ret
 ; Checking that the address laoded from is masked for a floating point load.
 ; CHECK-LABEL: csdb_emitted_for_superreg_use
 ; CHECK:      ldr w8, [x0]
 ; CHECK-NEXT: cmp sp, #0
 ; CHECK-NEXT: csetm x16, ne
 ; CHECK-NEXT: and w8, w8, w16
 ; csdb instruction must occur before the add instruction with x8 as operand.
 ; CHECK-NEXT: csdb
-; CHECK-NEXT: mov x17, sp
 ; CHECK-NEXT: add x9, x1, x8
 ; CHECK-NEXT: cmp w8, #0
-; CHECK-NEXT: and x17, x17, x16
 ; CHECK-NEXT: csel x0, x1, x9, eq
-; CHECK-NEXT: mov sp, x17
+; CHECK-NEXT: mov [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT: and [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT: mov sp, [[TMPREG]]
 ; CHECK-NEXT: ret
 }
 
 define i64 @no_masking_with_full_control_flow_barriers(i64 %a, i64 %b, i64* %p) speculative_load_hardening {
 ; CHECK-LABEL: no_masking_with_full_control_flow_barriers
 ; CHECK: dsb sy
 ; CHECK: isb
 entry:
entry:
   store <4 x i32> %shuffle, <4 x i32>* %dst, align 4
   ret void
 ; CHECK-LABEL: f_implicitdef_vector_load
 ; CHECK:       cmp     sp, #0
 ; CHECK-NEXT:  csetm   x16, ne
 ; CHECK-NEXT:  and     x1, x1, x16
 ; CHECK-NEXT:  csdb
 ; CHECK-NEXT:  ldr     d0, [x1]
-; CHECK-NEXT:  mov     x17, sp
-; CHECK-NEXT:  and     x17, x17, x16
 ; CHECK-NEXT:  mov     v0.d[1], v0.d[0]
 ; CHECK-NEXT:  str     q0, [x0]
-; CHECK-NEXT:  mov     sp, x17
+; CHECK-NEXT:  mov     [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT:  and     [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT:  mov     sp, [[TMPREG]]
 ; CHECK-NEXT:  ret
 }
 
 define <2 x double> @f_usedefvectorload(double* %a, double* %b) speculative_load_hardening {
 entry:
 ; CHECK-LABEL: f_usedefvectorload
 ; CHECK:       cmp     sp, #0
 ; CHECK-NEXT:  csetm   x16, ne
 ; CHECK-NEXT:  movi    v0.2d, #0000000000000000
 ; CHECK-NEXT:  and     x1, x1, x16
 ; CHECK-NEXT:  csdb
 ; CHECK-NEXT:  ld1     { v0.d }[0], [x1]
-; CHECK-NEXT:  mov     x17, sp
-; CHECK-NEXT:  and     x17, x17, x16
-; CHECK-NEXT:  mov     sp, x17
+; CHECK-NEXT:  mov     [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT:  and     [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT:  mov     sp, [[TMPREG]]
 ; CHECK-NEXT:  ret
   %0 = load double, double* %b, align 16
   %vld1_lane = insertelement <2 x double> <double undef, double 0.000000e+00>, double %0, i32 0
   ret <2 x double> %vld1_lane
 }
 
 define i32 @deadload() speculative_load_hardening {
 entry:
 ; CHECK-LABEL: deadload
 ; CHECK:       cmp     sp, #0
 ; CHECK-NEXT:  csetm   x16, ne
 ; CHECK-NEXT:  sub     sp, sp, #16
 ; CHECK-NEXT:  .cfi_def_cfa_offset 16
 ; CHECK-NEXT:  ldr     w8, [sp, #12]
 ; CHECK-NEXT:  add     sp, sp, #16
-; CHECK-NEXT:  mov     x17, sp
-; CHECK-NEXT:  and     x17, x17, x16
-; CHECK-NEXT:  mov     sp, x17
+; CHECK-NEXT:  mov     [[TMPREG:x[0-9]+]], sp
+; CHECK-NEXT:  and     [[TMPREG]], [[TMPREG]], x16
+; CHECK-NEXT:  mov     sp, [[TMPREG]]
 ; CHECK-NEXT:  ret
   %a = alloca i32, align 4
   %val = load volatile i32, i32* %a, align 4
   ret i32 undef
 }

llvm/trunk/test/CodeGen/AArch64/speculation-hardening.ll

-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,SLH,NOGISELSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,NOSLH,NOGISELNOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,SLH,GISELSLH --dump-input-on-failure
+; RUN sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,NOSLH,GISELNOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,SLH,NOGISELSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,NOSLH,NOGISELNOSLH --dump-input-on-failure
 
 define i32 @f(i8* nocapture readonly %p, i32 %i, i32 %N) local_unnamed_addr SLHATTR {
 ; CHECK-LABEL: f
 entry:
 ; SLH:  cmp sp, #0
 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
 
-; SLH:  mov x17, sp
-; SLH:  and x17, x17, x16
-; SLH:  mov sp, x17
-; NOSLH-NOT:  mov x17, sp
-; NOSLH-NOT:  and x17, x17, x16
-; NOSLH-NOT:  mov sp, x17
+; SLH:  mov [[TMPREG:x[0-9]+]], sp
+; SLH:  and [[TMPREG]], [[TMPREG]], x16
+; SLH:  mov sp, [[TMPREG]]
+; NOSLH-NOT:  mov [[TMPREG:x[0-9]+]], sp
+; NOSLH-NOT:  and [[TMPREG]], [[TMPREG]], x16
+; NOSLH-NOT:  mov sp, [[TMPREG]]
   %call = tail call i32 @tail_callee(i32 %i)
 ; SLH:  cmp sp, #0
 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
   %cmp = icmp slt i32 %call, %N
   br i1 %cmp, label %if.then, label %return
 ; GlobalISel lowers the branch to a b.ne sometimes instead of b.ge as expected..
 ; CHECK: b.[[COND:(ge)|(lt)|(ne)]]
 
 if.then:                                          ; preds = %entry
 ; NOSLH-NOT: csel x16, x16, xzr, {{(lt)|(ge)|(eq)}}
 ; SLH-DAG: csel x16, x16, xzr, {{(lt)|(ge)|(eq)}}
   %idxprom = sext i32 %i to i64
   %arrayidx = getelementptr inbounds i8, i8* %p, i64 %idxprom
   %0 = load i8, i8* %arrayidx, align 1
 ; CHECK-DAG:      ldrb [[LOADED:w[0-9]+]],
   %conv = zext i8 %0 to i32
   br label %return
 
 ; SLH-DAG: csel x16, x16, xzr, [[COND]]
 ; NOSLH-NOT: csel x16, x16, xzr, [[COND]]
 return:                                           ; preds = %entry, %if.then
   %retval.0 = phi i32 [ %conv, %if.then ], [ 0, %entry ]
-; SLH:  mov x17, sp
-; SLH:  and x17, x17, x16
-; SLH:  mov sp, x17
-; NOSLH-NOT:  mov x17, sp
-; NOSLH-NOT:  and x17, x17, x16
-; NOSLH-NOT:  mov sp, x17
+; SLH:  mov [[TMPREG:x[0-9]+]], sp
+; SLH:  and [[TMPREG]], [[TMPREG]], x16
+; SLH:  mov sp, [[TMPREG]]
+; NOSLH-NOT:  mov [[TMPREG:x[0-9]+]], sp
+; NOSLH-NOT:  and [[TMPREG]], [[TMPREG]], x16
+; NOSLH-NOT:  mov sp, [[TMPREG]]
   ret i32 %retval.0
 }
 
 ; Make sure that for a tail call, taint doesn't get put into SP twice.
 define i32 @tail_caller(i32 %a) local_unnamed_addr SLHATTR {
 ; CHECK-LABEL: tail_caller:
-; SLH:     mov     x17, sp
-; SLH:     and     x17, x17, x16
-; SLH:     mov     sp, x17
-; NOSLH-NOT:     mov     x17, sp
-; NOSLH-NOT:     and     x17, x17, x16
-; NOSLH-NOT:     mov     sp, x17
+; NOGISELSLH:     mov [[TMPREG:x[0-9]+]], sp
+; NOGISELSLH:     and [[TMPREG]], [[TMPREG]], x16
+; NOGISELSLH:     mov sp, [[TMPREG]]
+; NOGISELNOSLH-NOT:     mov [[TMPREG:x[0-9]+]], sp
+; NOGISELNOSLH-NOT:     and [[TMPREG]], [[TMPREG]], x16
+; NOGISELNOSLH-NOT:     mov sp, [[TMPREG]]
+; GISELSLH:     mov [[TMPREG:x[0-9]+]], sp
+; GISELSLH:     and [[TMPREG]], [[TMPREG]], x16
+; GISELSLH:     mov sp, [[TMPREG]]
+; GISELNOSLH-NOT:     mov [[TMPREG:x[0-9]+]], sp
+; GISELNOSLH-NOT:     and [[TMPREG]], [[TMPREG]], x16
+; GISELNOSLH-NOT:     mov sp, [[TMPREG]]
 ;  GlobalISel doesn't optimize tail calls (yet?), so only check that
 ;  cross-call taint register setup code is missing if a tail call was
 ;  actually produced.
-; SLH:     {{(bl tail_callee[[:space:]] cmp sp, #0)|(b tail_callee)}}
-; SLH-NOT: cmp sp, #0
+; NOGISELSLH:     b tail_callee
+; GISELSLH:       bl tail_callee
+; GISELSLH:       cmp sp, #0
+; SLH-NOT:        cmp sp, #0
   %call = tail call i32 @tail_callee(i32 %a)
   ret i32 %call
 }
 
 declare i32 @tail_callee(i32) local_unnamed_addr
 
 ; Verify that no cb(n)z/tb(n)z instructions are produced when implementing
 ; SLH

llvm/trunk/test/CodeGen/AArch64/speculation-hardening.mir

 # RUN: llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu \
 # RUN:     -start-before aarch64-speculation-hardening -o - %s \
+# RUN:     -debug-only=aarch64-speculation-hardening \
 # RUN:   | FileCheck %s --dump-input-on-failure
 
 # Check that the speculation hardening pass generates code as expected for
 # basic blocks ending with a variety of branch patterns:
 # - (1) no branches (fallthrough)
 # - (2) one unconditional branch
 # - (3) one conditional branch + fall-through
 # - (4) one conditional branch + one unconditional branch
--- |
    ret void
   }
   define void @condbranch_uncondbranch(i32 %a, i32 %b) speculative_load_hardening {
    ret void
   }
   define void @indirectbranch(i32 %a, i32 %b) speculative_load_hardening {
    ret void
   }
+  ; Also check that a non-default temporary register gets picked correctly to
+  ; transfer the SP to to and it with the taint register when the default
+  ; temporary isn't available.
+  define void @indirect_call_x17(i32 %a, i32 %b) speculative_load_hardening {
+   ret void
+  }
+  @g = common dso_local local_unnamed_addr global i64 (...)* null, align 8
+  define void @indirect_tailcall_x17(i32 %a, i32 %b) speculative_load_hardening {
+   ret void
+  }
+  define void @indirect_call_lr(i32 %a, i32 %b) speculative_load_hardening {
+   ret void
+  }
+  define void @RS_cannot_find_available_regs() speculative_load_hardening {
+   ret void
+  }
 ...
 ---
 name:            nobranch_fallthrough
 tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: nobranch_fallthrough
   bb.0:
     successors: %bb.1
body:             |
    liveins: $x0
   ; CHECK-NOT: csel
    RET undef $lr, implicit $x0
   bb.2:
    liveins: $x0
   ; CHECK-NOT: csel
    RET undef $lr, implicit $x0
 ...
+---
+name:            indirect_call_x17
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x17
+    ; CHECK-LABEL: indirect_call_x17
+    ; CHECK:       mov x0, sp
+    ; CHECK:       and x0, x0, x16
+    ; CHECK:       mov sp, x0
+    ; CHECK:       blr x17
+    BLR killed renamable $x17, implicit-def dead $lr, implicit $sp
+    RET undef $lr, implicit undef $w0
+...
+---
+name:           indirect_tailcall_x17
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: indirect_tailcall_x17
+    ; CHECK:       mov x1, sp
+    ; CHECK:       and x1, x1, x16
+    ; CHECK:       mov sp, x1
+    ; CHECK:       br x17
+    $x8 = ADRP target-flags(aarch64-page) @g
+    $x17 = LDRXui killed $x8, target-flags(aarch64-pageoff, aarch64-nc) @g
+    TCRETURNri killed $x17, 0, implicit $sp, implicit $x0
+...
+---
+name:           indirect_call_lr
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: indirect_call_lr
+    ; CHECK:            mov x1, sp
+    ; CHECK-NEXT:       and x1, x1, x16
+    ; CHECK-NEXT:       mov sp, x1
+    ; CHECK-NEXT:       blr x30
+    liveins: $x0, $lr
+    BLR killed renamable $lr, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $w0
+    $w0 = nsw ADDWri killed $w0, 1, 0
+    RET undef $lr, implicit $w0
+...
+---
+name:           RS_cannot_find_available_regs
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; In the rare case when no free temporary register is available for the
+    ; propagate taint-to-sp operation, just put in a full speculation barrier
+    ; (isb+dsb sy) at the start of the basic block. And don't put masks on
+    ; instructions for the rest of the basic block, since speculation in that
+    ; basic block was already done, so no need to do masking.
+    ; CHECK-LABEL: RS_cannot_find_available_regs
+    ; CHECK:       dsb sy
+    ; CHECK-NEXT:  isb
+    ; CHECK-NEXT:  ldr x0, [x0]
+    ; The following 2 instructions come from propagating the taint encoded in
+    ; sp at function entry to x16. It turns out the taint info in x16 is not
+    ; used in this function, so those instructions could be optimized away. An
+    ; optimization for later if it turns out this situation occurs often enough.
+    ; CHECK-NEXT:  cmp sp, #0
+    ; CHECK-NEXT:  csetm x16, ne
+    ; CHECK-NEXT:  ret
+    liveins: $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8, $x9, $x10, $x11, $x12, $x13, $x14, $x15, $x17, $x18, $x19, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp, $lr
+     $x0 = LDRXui killed $x0, 0
+     RET undef $lr, implicit $x0
+...