[x86/SLH] Extract the core register hardening logic to a low-level

helper and restructure the post-load hardening to use this.

This isn't as trivial as I would have liked because the post-load
hardening used a trick that only works for it where it swapped in
a temporary register to the load rather than replacing anything.
However, there is a simple way to do this without that trick that allows
this to easily reuse a friendly API for hardening a value in a register.
That API will in turn be usable in subsequent patcehs.

This also techincally changes the position at which we insert the subreg
extraction for the predicate state, but that never resulted in an actual
instruction and so tests don't change at all.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@337825 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/lib/Target/X86/X86SpeculativeLoadHardening.cpp b/lib/Target/X86/X86SpeculativeLoadHardening.cpp
index 10e1132d87776fb48c62c5eb519286191c53a86b..65d53e9fd3af10c570e0aa28741656d6dd41569b 100644 (file)
--- a/lib/Target/X86/X86SpeculativeLoadHardening.cpp
+++ b/lib/Target/X86/X86SpeculativeLoadHardening.cpp
@@ -191,7 +191,10 @@ private:
   sinkPostLoadHardenedInst(MachineInstr &MI,
                            SmallPtrSetImpl<MachineInstr *> &HardenedInstrs);
   bool canHardenRegister(unsigned Reg);
-  void hardenPostLoad(MachineInstr &MI);
+  unsigned hardenValueInRegister(unsigned Reg, MachineBasicBlock &MBB,
+                                 MachineBasicBlock::iterator InsertPt,
+                                 DebugLoc Loc);
+  unsigned hardenPostLoad(MachineInstr &MI);
   void hardenReturnInstr(MachineInstr &MI);
 };
 
@@ -1430,12 +1433,11 @@ void X86SpeculativeLoadHardeningPass::hardenAllLoads(MachineFunction &MF) {
           }
         }
 
-        // The register def'ed by this instruction is trivially hardened so map
-        // it to itself.
-        AddrRegToHardenedReg[MI.getOperand(0).getReg()] =
-            MI.getOperand(0).getReg();
+        unsigned HardenedReg = hardenPostLoad(MI);
+
+        // Mark the resulting hardened register as such so we don't re-harden.
+        AddrRegToHardenedReg[HardenedReg] = HardenedReg;
 
-        hardenPostLoad(MI);
         continue;
       }
 
@@ -1866,53 +1868,51 @@ bool X86SpeculativeLoadHardeningPass::canHardenRegister(unsigned Reg) {
   return RC->hasSuperClassEq(GPRRegClasses[Log2_32(RegBytes)]);
 }
 
-/// Harden a load by hardening the loaded value in the defined register.
+/// Harden a value in a register.
 ///
-/// We can harden a non-leaking load into a register without touching the
-/// address by just hiding all of the loaded bits during misspeculation. We use
-/// an `or` instruction to do this because we set up our poison value as all
-/// ones. And the goal is just for the loaded bits to not be exposed to
-/// execution and coercing them to one is sufficient.
-void X86SpeculativeLoadHardeningPass::hardenPostLoad(MachineInstr &MI) {
-  MachineBasicBlock &MBB = *MI.getParent();
-  DebugLoc Loc = MI.getDebugLoc();
-
-  // For all of these, the def'ed register operand is operand zero.
-  auto &DefOp = MI.getOperand(0);
-  unsigned OldDefReg = DefOp.getReg();
-  assert(canHardenRegister(OldDefReg) &&
-         "Cannot harden this instruction's defined register!");
+/// This is the low-level logic to fully harden a value sitting in a register
+/// against leaking during speculative execution.
+///
+/// Unlike hardening an address that is used by a load, this routine is required
+/// to hide *all* incoming bits in the register.
+///
+/// `Reg` must be a virtual register. Currently, it is required to be a GPR no
+/// larger than the predicate state register. FIXME: We should support vector
+/// registers here by broadcasting the predicate state.
+///
+/// The new, hardened virtual register is returned. It will have the same
+/// register class as `Reg`.
+unsigned X86SpeculativeLoadHardeningPass::hardenValueInRegister(
+    unsigned Reg, MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPt,
+    DebugLoc Loc) {
+  assert(canHardenRegister(Reg) && "Cannot harden this register!");
+  assert(TRI->isVirtualRegister(Reg) && "Cannot harden a physical register!");
 
-  auto *DefRC = MRI->getRegClass(OldDefReg);
-  int DefRegBytes = TRI->getRegSizeInBits(*DefRC) / 8;
+  auto *RC = MRI->getRegClass(Reg);
+  int Bytes = TRI->getRegSizeInBits(*RC) / 8;
 
   unsigned StateReg = PS->SSA.GetValueAtEndOfBlock(&MBB);
 
   // FIXME: Need to teach this about 32-bit mode.
-  if (DefRegBytes != 8) {
+  if (Bytes != 8) {
     unsigned SubRegImms[] = {X86::sub_8bit, X86::sub_16bit, X86::sub_32bit};
-    unsigned SubRegImm = SubRegImms[Log2_32(DefRegBytes)];
-    unsigned NarrowStateReg = MRI->createVirtualRegister(DefRC);
-    BuildMI(MBB, MI.getIterator(), Loc, TII->get(TargetOpcode::COPY),
-            NarrowStateReg)
+    unsigned SubRegImm = SubRegImms[Log2_32(Bytes)];
+    unsigned NarrowStateReg = MRI->createVirtualRegister(RC);
+    BuildMI(MBB, InsertPt, Loc, TII->get(TargetOpcode::COPY), NarrowStateReg)
         .addReg(StateReg, 0, SubRegImm);
     StateReg = NarrowStateReg;
   }
 
-  auto InsertPt = std::next(MI.getIterator());
-
   unsigned FlagsReg = 0;
   if (isEFLAGSLive(MBB, InsertPt, *TRI))
     FlagsReg = saveEFLAGS(MBB, InsertPt, Loc);
 
-  unsigned NewDefReg = MRI->createVirtualRegister(DefRC);
-  DefOp.setReg(NewDefReg);
-
+  unsigned NewReg = MRI->createVirtualRegister(RC);
   unsigned OrOpCodes[] = {X86::OR8rr, X86::OR16rr, X86::OR32rr, X86::OR64rr};
-  unsigned OrOpCode = OrOpCodes[Log2_32(DefRegBytes)];
-  auto OrI = BuildMI(MBB, InsertPt, Loc, TII->get(OrOpCode), OldDefReg)
+  unsigned OrOpCode = OrOpCodes[Log2_32(Bytes)];
+  auto OrI = BuildMI(MBB, InsertPt, Loc, TII->get(OrOpCode), NewReg)
                  .addReg(StateReg)
-                 .addReg(NewDefReg);
+                 .addReg(Reg);
   OrI->addRegisterDead(X86::EFLAGS, TRI);
   ++NumInstsInserted;
   LLVM_DEBUG(dbgs() << "  Inserting or: "; OrI->dump(); dbgs() << "\n");
@@ -1920,7 +1920,44 @@ void X86SpeculativeLoadHardeningPass::hardenPostLoad(MachineInstr &MI) {
   if (FlagsReg)
     restoreEFLAGS(MBB, InsertPt, Loc, FlagsReg);
 
+  return NewReg;
+}
+
+/// Harden a load by hardening the loaded value in the defined register.
+///
+/// We can harden a non-leaking load into a register without touching the
+/// address by just hiding all of the loaded bits during misspeculation. We use
+/// an `or` instruction to do this because we set up our poison value as all
+/// ones. And the goal is just for the loaded bits to not be exposed to
+/// execution and coercing them to one is sufficient.
+///
+/// Returns the newly hardened register.
+unsigned X86SpeculativeLoadHardeningPass::hardenPostLoad(MachineInstr &MI) {
+  MachineBasicBlock &MBB = *MI.getParent();
+  DebugLoc Loc = MI.getDebugLoc();
+
+  auto &DefOp = MI.getOperand(0);
+  unsigned OldDefReg = DefOp.getReg();
+  auto *DefRC = MRI->getRegClass(OldDefReg);
+
+  // Because we want to completely replace the uses of this def'ed value with
+  // the hardened value, create a dedicated new register that will only be used
+  // to communicate the unhardened value to the hardening.
+  unsigned UnhardenedReg = MRI->createVirtualRegister(DefRC);
+  DefOp.setReg(UnhardenedReg);
+
+  // Now harden this register's value, getting a hardened reg that is safe to
+  // use. Note that we insert the instructions to compute this *after* the
+  // defining instruction, not before it.
+  unsigned HardenedReg = hardenValueInRegister(
+      UnhardenedReg, MBB, std::next(MI.getIterator()), Loc);
+
+  // Finally, replace the old register (which now only has the uses of the
+  // original def) with the hardened register.
+  MRI->replaceRegWith(/*FromReg*/ OldDefReg, /*ToReg*/ HardenedReg);
+
   ++NumPostLoadRegsHardened;
+  return HardenedReg;
 }
 
 /// Harden a return instruction.