[X86] Fix bug in -mlvi-cfi that may clobber a live register

[sconstab]

Summary

Fix for this bug: https://bugs.llvm.org/show_bug.cgi?id=47740

The fix uses the existing findDeadCallerSavedReg() function instead of a hacky heuristic to find a scratch register to clobber.

[sconstab]

Update required by https://reviews.llvm.org/D88924.

[craig.topper]

LGTM

[sconstab]

LGTM

llvm/lib/Target/X86/X86LoadValueInjectionRetHardening.cpp

bool X86LoadValueInjectionRetHardeningPass::runOnMachineFunction(
   // Don't skip functions with the "optnone" attr but participate in opt-bisect.
   const Function &F = MF.getFunction();
   if (!F.hasOptNone() && skipFunction(F))
     return false;
 
   ++NumFunctionsConsidered;
   const X86RegisterInfo *TRI = Subtarget->getRegisterInfo();
   const X86InstrInfo *TII = Subtarget->getInstrInfo();
-  unsigned ClobberReg = X86::NoRegister;
-  std::bitset<X86::NUM_TARGET_REGS> UnclobberableGR64s;
-  UnclobberableGR64s.set(X86::RSP); // can't clobber stack pointer
-  UnclobberableGR64s.set(X86::RIP); // can't clobber instruction pointer
-  UnclobberableGR64s.set(X86::RAX); // used for function return
-  UnclobberableGR64s.set(X86::RDX); // used for function return
-
-  // We can clobber any register allowed by the function's calling convention.
-  for (const MCPhysReg *PR = TRI->getCalleeSavedRegs(&MF); auto Reg = *PR; ++PR)
-    UnclobberableGR64s.set(Reg);
-  for (auto &Reg : X86::GR64RegClass) {
-    if (!UnclobberableGR64s.test(Reg)) {
-      ClobberReg = Reg;
-      break;
-    }
-  }
-
-  if (ClobberReg != X86::NoRegister) {
-    LLVM_DEBUG(dbgs() << "Selected register "
-                      << Subtarget->getRegisterInfo()->getRegAsmName(ClobberReg)
-                      << " to clobber\n");
-  } else {
-    LLVM_DEBUG(dbgs() << "Could not find a register to clobber\n");
-  }
 
   bool Modified = false;
   for (auto &MBB : MF) {
-    if (MBB.empty())
-      continue;
-
-    MachineInstr &MI = MBB.back();
-    if (MI.getOpcode() != X86::RETQ)
+    for (auto MBBI = MBB.begin(); MBBI != MBB.end(); ++MBBI) {
+      if (MBBI->getOpcode() != X86::RETQ)
         continue;
 
+      unsigned ClobberReg =
+          TRI->findDeadCallerSavedReg(MBB, MBBI, true /* Is64Bit */);
       if (ClobberReg != X86::NoRegister) {
-      MBB.erase_instr(&MI);
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::POP64r))
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::POP64r))
             .addReg(ClobberReg, RegState::Define)
             .setMIFlag(MachineInstr::FrameDestroy);
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::LFENCE));
-      BuildMI(MBB, MBB.end(), DebugLoc(), TII->get(X86::JMP64r))
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
+        BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::JMP64r))
             .addReg(ClobberReg);
+        MBB.erase(MBBI);
       } else {
-      // In case there is no available scratch register, we can still read from
-      // RSP to assert that RSP points to a valid page. The write to RSP is
-      // also helpful because it verifies that the stack's write permissions
-      // are intact.
-      MachineInstr *Fence = BuildMI(MBB, MI, DebugLoc(), TII->get(X86::LFENCE));
+        // In case there is no available scratch register, we can still read
+        // from RSP to assert that RSP points to a valid page. The write to RSP
+        // is also helpful because it verifies that the stack's write
+        // permissions are intact.
+        MachineInstr *Fence =
+            BuildMI(MBB, MBBI, DebugLoc(), TII->get(X86::LFENCE));
         addRegOffset(BuildMI(MBB, Fence, DebugLoc(), TII->get(X86::SHL64mi)),
                      X86::RSP, false, 0)
             .addImm(0)
             ->addRegisterDead(X86::EFLAGS, TRI);
       }
 
       ++NumFences;
       Modified = true;
+      break;
+    }
   }
 
   if (Modified)
     ++NumFunctionsMitigated;
   return Modified;
 }
 
 INITIALIZE_PASS(X86LoadValueInjectionRetHardeningPass, PASS_KEY,
                 "X86 LVI ret hardener", false, false)
 
 FunctionPass *llvm::createX86LoadValueInjectionRetHardeningPass() {
   return new X86LoadValueInjectionRetHardeningPass();
 }

llvm/test/CodeGen/X86/lvi-hardening-ret.ll

entry:
   %y.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   store i32 %y, i32* %y.addr, align 4
   %0 = load i32, i32* %x.addr, align 4
   %1 = load i32, i32* %y.addr, align 4
   %add = add nsw i32 %0, %1
   ret i32 %add
 ; CHECK-NOT:  retq
-; CHECK:      shlq $0, (%{{[^ ]*}})
+; CHECK:      popq %rcx
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: retq
+; CHECK-NEXT: jmpq *%rcx
 }
 
 ; Function Attrs: noinline nounwind optnone uwtable
 define dso_local preserve_mostcc void @preserve_most() #0 {
 ; CHECK-LABEL: preserve_most:
 entry:
   ret void
 ; CHECK-NOT:  retq
-; CHECK:      popq %r11
+; CHECK:      popq %rax
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: jmpq *%r11
+; CHECK-NEXT: jmpq *%rax
 }
 
 ; Function Attrs: noinline nounwind optnone uwtable
 define dso_local preserve_allcc void @preserve_all() #0 {
 ; CHECK-LABEL: preserve_all:
 entry:
   ret void
 ; CHECK-NOT:  retq
-; CHECK:      popq %r11
+; CHECK:      popq %rax
 ; CHECK-NEXT: lfence
-; CHECK-NEXT: jmpq *%r11
+; CHECK-NEXT: jmpq *%rax
+}
+
+define { i64, i128 } @ret_i64_i128() #0 {
+; CHECK-LABEL: ret_i64_i128:
+  ret { i64, i128 } { i64 1, i128 36893488147419103235 }
+; CHECK-NOT:  retq
+; CHECK:      popq %rsi
+; CHECK-NEXT: lfence
+; CHECK-NEXT: jmpq *%rsi
 }
 
 attributes #0 = { "target-features"="+lvi-cfi" }
 attributes #1 = { "no_caller_saved_registers" "target-features"="+lvi-cfi" }