[llvm] r337895 - [x86/SLH] Sink the return hardening into the main block-walk + hardening

[Chandler Carruth via llvm-commits]

Author: chandlerc
Date: Wed Jul 25 02:18:48 2018
New Revision: 337895

URL: http://llvm.org/viewvc/llvm-project?rev=337895&view=rev
Log:
[x86/SLH] Sink the return hardening into the main block-walk + hardening
code.

This consolidates all our hardening calls, and simplifies the code
a bit. It seems much more clear to handle all of these together.

No functionality changed here.

Modified:
    llvm/trunk/lib/Target/X86/X86SpeculativeLoadHardening.cpp

Modified: llvm/trunk/lib/Target/X86/X86SpeculativeLoadHardening.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/Target/X86/X86SpeculativeLoadHardening.cpp?rev=337895&r1=337894&r2=337895&view=diff
==============================================================================
--- llvm/trunk/lib/Target/X86/X86SpeculativeLoadHardening.cpp (original)
+++ llvm/trunk/lib/Target/X86/X86SpeculativeLoadHardening.cpp Wed Jul 25 02:18:48 2018
@@ -541,24 +541,6 @@ bool X86SpeculativeLoadHardeningPass::ru
       PS->SSA.RewriteUse(Op);
     }
 
-  // If we are hardening interprocedurally, find each returning block and
-  // protect the caller from being returned to through misspeculation.
-  if (HardenInterprocedurally)
-    for (MachineBasicBlock &MBB : MF) {
-      if (MBB.empty())
-        continue;
-
-      MachineInstr &MI = MBB.back();
-
-      // We only care about returns that are not also calls. For calls, that
-      // happen to also be returns (tail calls) we will have already handled
-      // them as calls.
-      if (!MI.isReturn() || MI.isCall())
-        continue;
-
-      hardenReturnInstr(MI);
-    }
-
   LLVM_DEBUG(dbgs() << "Final speculative load hardened function:\n"; MF.dump();
              dbgs() << "\n"; MF.verify(this));
   return true;
@@ -1588,16 +1570,24 @@ void X86SpeculativeLoadHardeningPass::tr
           hardenIndirectCallOrJumpInstr(MI, AddrRegToHardenedReg);
       }
 
-      // After we finish processing the instruction and doing any hardening
-      // necessary for it, we need to handle transferring the predicate state
-      // into a call and recovering it after the call returns (if it returns).
-      if (!MI.isCall())
+      // After we finish hardening loads we handle interprocedural hardening if
+      // enabled and relevant for this instruction.
+      if (!HardenInterprocedurally)
+        continue;
+      if (!MI.isCall() && !MI.isReturn())
         continue;
 
-      // If we're not hardening interprocedurally, we can just skip calls.
-      if (!HardenInterprocedurally)
+      // If this is a direct return (IE, not a tail call) just directly harden
+      // it.
+      if (MI.isReturn() && !MI.isCall()) {
+        hardenReturnInstr(MI);
         continue;
+      }
 
+      // Otherwise we have a call. We need to handle transferring the predicate
+      // state into a call and recovering it after the call returns unless this
+      // is a tail call.
+      assert(MI.isCall() && "Should only reach here for calls!");
       auto InsertPt = MI.getIterator();
       DebugLoc Loc = MI.getDebugLoc();
 
@@ -1607,11 +1597,12 @@ void X86SpeculativeLoadHardeningPass::tr
       unsigned StateReg = PS->SSA.GetValueAtEndOfBlock(&MBB);
       mergePredStateIntoSP(MBB, InsertPt, Loc, StateReg);
 
-      // If this call is also a return (because it is a tail call) we're done.
+      // If this call is also a return, it is a tail call and we don't need
+      // anything else to handle it so just continue.
       if (MI.isReturn())
         continue;
 
-      // Otherwise we need to step past the call and recover the predicate
+      // We need to step past the call and recover the predicate
       // state from SP after the return, and make this new state available.
       ++InsertPt;
       unsigned NewStateReg = extractPredStateFromSP(MBB, InsertPt, Loc);