Index: lib/Target/AArch64/AArch64FrameLowering.cpp =================================================================== --- lib/Target/AArch64/AArch64FrameLowering.cpp +++ lib/Target/AArch64/AArch64FrameLowering.cpp @@ -572,6 +572,17 @@ // } +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(); @@ -595,7 +606,9 @@ 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); unsigned CFIIndex = @@ -888,10 +901,14 @@ // 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); unsigned CFIIndex = Index: test/CodeGen/AArch64/sign-return-address.ll =================================================================== --- test/CodeGen/AArch64/sign-return-address.ll +++ test/CodeGen/AArch64/sign-return-address.ll @@ -84,3 +84,26 @@ 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 +}