diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp index b0451ca2edb6..a99dd356d4ff 100644 --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp @@ -597,6 +597,17 @@ static void adaptForLdStOpt(MachineBasicBlock &MBB, // } +static bool ShouldSignWithAKey(MachineFunction &MF) { + const Function &F = MF.getFunction(); + if (!F.hasFnAttribute("sign-return-address-key")) + return true; + + const StringRef Key = + F.getFnAttribute("sign-return-address-key").getValueAsString(); + assert(Key.equals_lower("a_key") || Key.equals_lower("b_key")); + return Key.equals_lower("a_key"); +} + void AArch64FrameLowering::emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const { MachineBasicBlock::iterator MBBI = MBB.begin(); @@ -620,7 +631,9 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF, DebugLoc DL; if (ShouldSignReturnAddress(MF)) { - BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACIASP)) + BuildMI( + MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::PACIASP : AArch64::PACIBSP)) .setMIFlag(MachineInstr::FrameSetup); } @@ -907,10 +920,14 @@ static void InsertReturnAddressAuth(MachineFunction &MF, // instructions, namely RETA{A,B}, that can be used instead. if (Subtarget.hasV8_3aOps() && MBBI != MBB.end() && MBBI->getOpcode() == AArch64::RET_ReallyLR) { - BuildMI(MBB, MBBI, DL, TII->get(AArch64::RETAA)).copyImplicitOps(*MBBI); + BuildMI(MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::RETAA : AArch64::RETAB)) + .copyImplicitOps(*MBBI); MBB.erase(MBBI); } else { - BuildMI(MBB, MBBI, DL, TII->get(AArch64::AUTIASP)) + BuildMI( + MBB, MBBI, DL, + TII->get(ShouldSignWithAKey(MF) ? AArch64::AUTIASP : AArch64::AUTIBSP)) .setMIFlag(MachineInstr::FrameDestroy); } } diff --git a/llvm/test/CodeGen/AArch64/sign-return-address.ll b/llvm/test/CodeGen/AArch64/sign-return-address.ll index a0c73058a301..c057c815acfd 100644 --- a/llvm/test/CodeGen/AArch64/sign-return-address.ll +++ b/llvm/test/CodeGen/AArch64/sign-return-address.ll @@ -84,3 +84,26 @@ define fastcc void @spill_lr_and_tail_call(i64 %x) "sign-return-address"="all" { tail call fastcc i64 @bar(i64 %x) ret void } + +; CHECK-LABEL: @leaf_sign_all_a_key +; CHECK: paciasp +; CHECK: autiasp +define i32 @leaf_sign_all_a_key(i32 %x) "sign-return-address"="all" "sign-return-address-key"="a_key" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_all_b_key +; CHECK: pacibsp +; CHECK: autibsp +define i32 @leaf_sign_all_b_key(i32 %x) "sign-return-address"="all" "sign-return-address-key"="b_key" { + ret i32 %x +} + +; CHECK-LABEL: @leaf_sign_all_v83_b_key +; CHECK: pacibsp +; CHECK-NOT: ret +; CHECK: retab +; CHECK-NOT: ret +define i32 @leaf_sign_all_v83_b_key(i32 %x) "sign-return-address"="all" "target-features"="+v8.3a" "sign-return-address-key"="b_key" { + ret i32 %x +}