diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp @@ -5528,9 +5528,83 @@ } } +static bool ShouldSignReturnAddress(const MachineFunction &MF) { + // The function should be signed in the following situations: + // - sign-return-address=all + // - sign-return-address=non-leaf and the functions spills the LR + + const Function &F = MF.getFunction(); + if (!F.hasFnAttribute("sign-return-address")) + return false; + + StringRef Scope = F.getFnAttribute("sign-return-address").getValueAsString(); + if (Scope.equals("none")) + return false; + + if (Scope.equals("all")) + return true; + + assert(Scope.equals("non-leaf") && "Expected all, none or non-leaf"); + + for (const auto &Info : MF.getFrameInfo().getCalleeSavedInfo()) + if (Info.getReg() == AArch64::LR) + return true; + + return false; +} + +static bool ShouldSignWithAKey(const 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 AArch64InstrInfo::buildOutlinedFrame( MachineBasicBlock &MBB, MachineFunction &MF, const outliner::OutlinedFunction &OF) const { + + // If for the machine functions of all candiates ShouldSignReturnAddress + // returns true *and* if the machine functions of all candidates agree on the + // return value of ShouldSignWithAKey, we let the outlined function behave + // like the candidates in this point + if (all_of(OF.Candidates.begin(), OF.Candidates.end(), + [](const outliner::Candidate &c) { + return ShouldSignReturnAddress(*c.getMF()); + })) { + MachineBasicBlock::iterator MBBI = MBB.begin(); + if (all_of(OF.Candidates.begin(), OF.Candidates.end(), + [](const outliner::Candidate &c) { + return ShouldSignWithAKey(*c.getMF()); + })) { + BuildMI(MBB, MBBI, DebugLoc(), get(AArch64::PACIASP)) + .setMIFlag(MachineInstr::FrameSetup); + unsigned CFIIndex = + MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr)); + BuildMI(MBB, MBBI, DebugLoc(), get(AArch64::CFI_INSTRUCTION)) + .addCFIIndex(CFIIndex) + .setMIFlags(MachineInstr::FrameSetup); + } else if (all_of(OF.Candidates.begin(), OF.Candidates.end(), + [](const outliner::Candidate &c) { + return !ShouldSignWithAKey(*c.getMF()); + })) { + + BuildMI(MBB, MBBI, DebugLoc(), get(AArch64::EMITBKEY)) + .setMIFlag(MachineInstr::FrameSetup); + BuildMI(MBB, MBBI, DebugLoc(), get(AArch64::PACIBSP)) + .setMIFlag(MachineInstr::FrameSetup); + unsigned CFIIndex = + MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr)); + BuildMI(MBB, MBBI, DebugLoc(), get(AArch64::CFI_INSTRUCTION)) + .addCFIIndex(CFIIndex) + .setMIFlags(MachineInstr::FrameSetup); + } + } + // For thunk outlining, rewrite the last instruction from a call to a // tail-call. if (OF.FrameConstructionID == MachineOutlinerThunk) { @@ -5543,8 +5617,8 @@ TailOpcode = AArch64::TCRETURNriALL; } MachineInstr *TC = BuildMI(MF, DebugLoc(), get(TailOpcode)) - .add(Call->getOperand(0)) - .addImm(0); + .add(Call->getOperand(0)) + .addImm(0); MBB.insert(MBB.end(), TC); Call->eraseFromParent(); } diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-akey.ll b/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-akey.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-akey.ll @@ -0,0 +1,77 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner \ +; RUN: -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s + +define void @a() #0 { +; CHECK: a: // @a +; CHECK: paciasp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +define void @b() #0 { +; CHECK: b: // @b +; CHECK: paciasp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +define void @c() #0 { +; CHECK: c: // @c +; CHECK: paciasp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +attributes #0 = { "sign-return-address"="all" "sign-return-address-key"="a_key" } + +; OUTLINED_FUNCTION_0 +; CHECK: OUTLINED_FUNCTION_0: // @OUTLINED_FUNCTION_0 +; CHECK: // %bb.0: +; a_key +; CHECK-NEXT: paciasp +; CHECK-NEXT: .cfi_negate_ra_state +; not b_key +; CHECK-NOT: .cfi_b_key_frame +; CHECK-NOT: pacibsp +; CHECK-NOT: .cfi_negate_ra_state +; CHECK: .cfi_endproc + diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-bkey.ll b/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-bkey.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-bkey.ll @@ -0,0 +1,86 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner \ +; RUN: -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s + +define void @a() #0 { +; CHECK: a: // @a +; CHECK-NEXT: .cfi_startproc +; CHECK-NEXT: // %bb.0: +; CHECK-NEXT: .cfi_b_key_frame +; CHECK-NEXT: pacibsp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +define void @b() #0 { +; CHECK: b: // @b +; CHECK-NEXT: .cfi_startproc +; CHECK-NEXT: // %bb.0: +; CHECK-NEXT: .cfi_b_key_frame +; CHECK-NEXT: pacibsp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +define void @c() #0 { +; CHECK: c: // @c +; CHECK-NEXT: .cfi_startproc +; CHECK-NEXT: // %bb.0: +; CHECK-NEXT: .cfi_b_key_frame +; CHECK-NEXT: pacibsp +; CHECK-NEXT: .cfi_negate_ra_state + %1 = alloca i32, align 4 + %2 = alloca i32, align 4 + %3 = alloca i32, align 4 + %4 = alloca i32, align 4 + %5 = alloca i32, align 4 + %6 = alloca i32, align 4 + store i32 1, i32* %1, align 4 + store i32 2, i32* %2, align 4 + store i32 3, i32* %3, align 4 + store i32 4, i32* %4, align 4 + store i32 5, i32* %5, align 4 + store i32 6, i32* %6, align 4 + ret void +; CHECK: .cfi_endproc +} + +attributes #0 = { "sign-return-address"="all" "sign-return-address-key"="b_key" } + +; OUTLINED_FUNCTION_0 +; CHECK: OUTLINED_FUNCTION_0: // @OUTLINED_FUNCTION_0 +; CHECK: // %bb.0: +; not a_key +; CHECK-NOT: paciasp +; CHECK-NOT: .cfi_negate_ra_state +; b_key +; CHECK-NEXT: .cfi_b_key_frame +; CHECK-NEXT: pacibsp +; CHECK-NEXT: .cfi_negate_ra_state +; CHECK: .cfi_endproc +