Index: llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp =================================================================== --- llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp +++ llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp @@ -157,6 +157,7 @@ void insertRegToSPTaintPropagation(MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI, unsigned TmpReg) const; + unsigned findDeadReg(MachineInstr *MI, unsigned PreferredReg) const; bool slhLoads(MachineBasicBlock &MBB); bool makeGPRSpeculationSafe(MachineBasicBlock &MBB, @@ -226,6 +227,32 @@ } } +unsigned AArch64SpeculationHardening::findDeadReg(MachineInstr *MI, + unsigned PreferredReg) const { + // This function will only be called on either terminator instructions or on + // call instructions. The implementation makes use of that knowledge. + assert(MI->isTerminator() || MI->isCall()); + auto RegIsDead = [&](unsigned Reg) { + if (MI->isTerminator()) + return !MI->readsRegister(Reg, TRI); + else + return MI->definesRegister(Reg, TRI) && !MI->readsRegister(Reg, TRI); + }; + if (RegIsDead(PreferredReg)) + return PreferredReg; + for (unsigned Reg : AArch64::GPR64allRegClass) { + if (RegIsDead(Reg)) { + LLVM_DEBUG(dbgs() << " - Found alternative " << printReg(Reg, TRI) + << " to be available as a tmp reg, to transfer taint " + "to SP before terminator." + << *MI); + return Reg; + } + } + llvm_unreachable("The nature of call instructions and the ABI must result in " + "finding at least one register above."); +} + bool AArch64SpeculationHardening::instrumentControlFlow( MachineBasicBlock &MBB) { LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB); @@ -278,14 +305,18 @@ (ReturnInstructions.size() > 0) || (CallInstructions.size() > 0); for (MachineInstr *Return : ReturnInstructions) - insertRegToSPTaintPropagation(Return->getParent(), Return, AArch64::X17); + insertRegToSPTaintPropagation(Return->getParent(), Return, + findDeadReg(Return, AArch64::X17)); for (MachineInstr *Call : CallInstructions) { // Just after the call: MachineBasicBlock::iterator i = Call; i++; insertSPToRegTaintPropagation(Call->getParent(), i); // Just before the call: - insertRegToSPTaintPropagation(Call->getParent(), Call, AArch64::X17); + // Prefer LR as the temporary register to use, as it should always be + // available, as the call instruction writes to it. + insertRegToSPTaintPropagation(Call->getParent(), Call, + findDeadReg(Call, AArch64::LR)); } } Index: llvm/test/CodeGen/AArch64/speculation-hardening.ll =================================================================== --- llvm/test/CodeGen/AArch64/speculation-hardening.ll +++ llvm/test/CodeGen/AArch64/speculation-hardening.ll @@ -1,9 +1,9 @@ -; 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 @@ -13,12 +13,12 @@ ; 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 x30, sp +; SLH: and x30, x30, x16 +; SLH: mov sp, x30 +; NOSLH-NOT: mov x30, sp +; NOSLH-NOT: and x30, x30, x16 +; NOSLH-NOT: mov sp, x30 %call = tail call i32 @tail_callee(i32 %i) ; SLH: cmp sp, #0 ; SLH: csetm x16, ne @@ -55,17 +55,25 @@ ; 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 x17, sp +; NOGISELSLH: and x17, x17, x16 +; NOGISELSLH: mov sp, x17 +; NOGISELNOSLH-NOT: mov x17, sp +; NOGISELNOSLH-NOT: and x17, x17, x16 +; NOGISELNOSLH-NOT: mov sp, x17 +; GISELSLH: mov x30, sp +; GISELSLH: and x30, x30, x16 +; GISELSLH: mov sp, x30 +; GISELNOSLH-NOT: mov x30, sp +; GISELNOSLH-NOT: and x30, x30, x16 +; GISELNOSLH-NOT: mov sp, x30 ; 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 } Index: llvm/test/CodeGen/AArch64/speculation-hardening.mir =================================================================== --- llvm/test/CodeGen/AArch64/speculation-hardening.mir +++ llvm/test/CodeGen/AArch64/speculation-hardening.mir @@ -25,6 +25,16 @@ 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 + } ... --- name: nobranch_fallthrough @@ -115,3 +125,32 @@ ; 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 x30, sp + ; CHECK: and x30, x30, x16 + ; CHECK: mov sp, x30 + ; 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 +...