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 @@ -5037,8 +5037,76 @@ return 0u; } -outliner::OutlinedFunction -AArch64InstrInfo::getOutliningCandidateInfo( +static bool +outliningCandidatesSigningScopeConsensus(const outliner::Candidate &a, + const outliner::Candidate &b) { + const Function &Fa = a.getMF()->getFunction(); + const Function &Fb = b.getMF()->getFunction(); + + if (Fa.hasFnAttribute("sign-return-address")) { + StringRef ScopeA = + Fa.getFnAttribute("sign-return-address").getValueAsString(); + return ScopeA.equals("none"); + } + + if (Fb.hasFnAttribute("sign-return-address")) { + StringRef ScopeB = + Fb.getFnAttribute("sign-return-address").getValueAsString(); + return ScopeB.equals("none"); + } + if (!Fa.hasFnAttribute("sign-return-address") && + !Fb.hasFnAttribute("sign-return-address")) { + return true; + } + + if (Fa.hasFnAttribute("sign-return-address") && + Fb.hasFnAttribute("sign-return-address")) { + StringRef ScopeA = + Fa.getFnAttribute("sign-return-address").getValueAsString(); + StringRef ScopeB = + Fb.getFnAttribute("sign-return-address").getValueAsString(); + return ScopeA.equals(ScopeB); + } + + return false; +} + +static bool +outliningCandidatesSigningKeyConsensus(const outliner::Candidate &a, + const outliner::Candidate &b) { + const Function &Fa = a.getMF()->getFunction(); + const Function &Fb = b.getMF()->getFunction(); + + if (!Fa.hasFnAttribute("sign-return-address-key") && + !Fb.hasFnAttribute("sign-return-address-key")) { + return true; + } + + if (Fa.hasFnAttribute("sign-return-address-key") && + Fb.hasFnAttribute("sign-return-address-key")) { + StringRef KeyA = + Fa.getFnAttribute("sign-return-address-key").getValueAsString(); + StringRef KeyB = + Fb.getFnAttribute("sign-return-address-key").getValueAsString(); + return KeyA.equals(KeyB); + } + + if (Fa.hasFnAttribute("sign-return-address-key")) { + StringRef KeyA = + Fa.getFnAttribute("sign-return-address-key").getValueAsString(); + return KeyA.equals_lower("a_key"); + } + + if (Fb.hasFnAttribute("sign-return-address-key")) { + StringRef KeyB = + Fb.getFnAttribute("sign-return-address-key").getValueAsString(); + return KeyB.equals_lower("a_key"); + } + + return false; +} + +outliner::OutlinedFunction AArch64InstrInfo::getOutliningCandidateInfo( std::vector &RepeatedSequenceLocs) const { outliner::Candidate &FirstCand = RepeatedSequenceLocs[0]; unsigned SequenceSize = @@ -5046,6 +5114,30 @@ [this](unsigned Sum, const MachineInstr &MI) { return Sum + getInstSizeInBytes(MI); }); + unsigned NumBytesToCreateFrame = 0; + + // We only allow outlining for functions having exactly matching return + // address signing attributes, i.e., all share the same value for the + // attribute "sign-return-address" and all share the same type of key they + // are signed with. + if (std::adjacent_find( + RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(), + // Return true if a and b are non-equal w.r.t. return address signing + [](const outliner::Candidate &a, const outliner::Candidate &b) { + if (outliningCandidatesSigningScopeConsensus(a, b) && + outliningCandidatesSigningKeyConsensus(a, b)) { + return false; + } + return true; + }) != RepeatedSequenceLocs.end()) { + return outliner::OutlinedFunction(); + } + + // Since at this point all candidates agree on their return address signing + // picking just one is fine. + const Function &FCF = FirstCand.getMF()->getFunction(); + if (FCF.hasFnAttribute("sign-return-address")) + NumBytesToCreateFrame += 8; // Properties about candidate MBBs that hold for all of them. unsigned FlagsSetInAll = 0xF; @@ -5107,7 +5199,7 @@ }; unsigned FrameID = MachineOutlinerDefault; - unsigned NumBytesToCreateFrame = 4; + NumBytesToCreateFrame += 4; bool HasBTI = any_of(RepeatedSequenceLocs, [](outliner::Candidate &C) { return C.getMF()->getFunction().hasFnAttribute("branch-target-enforcement"); @@ -5543,16 +5635,19 @@ 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(); } + bool IsLeafFunction = true; + // Is there a call in the outlined range? - auto IsNonTailCall = [](MachineInstr &MI) { + auto IsNonTailCall = [](const MachineInstr &MI) { return MI.isCall() && !MI.isReturn(); }; + if (std::any_of(MBB.instr_begin(), MBB.instr_end(), IsNonTailCall)) { // Fix up the instructions in the range, since we're going to modify the // stack. @@ -5560,6 +5655,8 @@ "Can only fix up stack references once"); fixupPostOutline(MBB); + IsLeafFunction = false; + // LR has to be a live in so that we can save it. MBB.addLiveIn(AArch64::LR); @@ -5606,6 +5703,68 @@ Et = MBB.insert(Et, LDRXpost); } + // If a bunch of candidates reach this point they must agree on their return + // address signing. It is therefore enough to just consider the signing + // behaviour of one of them + const Function &CF = OF.Candidates.front().getMF()->getFunction(); + bool ShouldSignReturnAddr = false; + if (CF.hasFnAttribute("sign-return-address")) { + StringRef Scope = + CF.getFnAttribute("sign-return-address").getValueAsString(); + if (Scope.equals("all")) + ShouldSignReturnAddr = true; + else if (Scope.equals("non-leaf") && !IsLeafFunction) + ShouldSignReturnAddr = true; + } + + // a_key is the default + bool ShouldSignReturnAddrWithAKey = true; + if (CF.hasFnAttribute("sign-return-address-key")) { + const StringRef Key = + CF.getFnAttribute("sign-return-address-key").getValueAsString(); + // Key can either be a_key or b_key + assert((Key.equals_lower("a_key") || Key.equals_lower("b_key")) && + "Return address signing key must be either a_key or b_key"); + ShouldSignReturnAddrWithAKey = Key.equals_lower("a_key"); + } + + // Insert PAC/AUT instructions + if (ShouldSignReturnAddr) { + MachineBasicBlock::iterator MBBPAC = MBB.begin(); + MachineBasicBlock::iterator MBBAUT = MBB.getFirstTerminator(); + DebugLoc DL; + if (MBBAUT != MBB.end()) + DL = MBBAUT->getDebugLoc(); + if (ShouldSignReturnAddrWithAKey) { + BuildMI(MBB, MBBPAC, DebugLoc(), get(AArch64::PACIASP)) + .setMIFlag(MachineInstr::FrameSetup); + } else { + BuildMI(MBB, MBBPAC, DebugLoc(), get(AArch64::EMITBKEY)) + .setMIFlag(MachineInstr::FrameSetup); + BuildMI(MBB, MBBPAC, DebugLoc(), get(AArch64::PACIBSP)) + .setMIFlag(MachineInstr::FrameSetup); + } + unsigned CFIIndex = + MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr)); + BuildMI(MBB, MBBPAC, DebugLoc(), get(AArch64::CFI_INSTRUCTION)) + .addCFIIndex(CFIIndex) + .setMIFlags(MachineInstr::FrameSetup); + const AArch64Subtarget &Subtarget = MF.getSubtarget(); + if (Subtarget.hasV8_3aOps() && MBBAUT != MBB.end() && + MBBAUT->getOpcode() == AArch64::RET_ReallyLR) { + BuildMI( + MBB, MBBAUT, DL, + get(ShouldSignReturnAddrWithAKey ? AArch64::RETAA : AArch64::RETAB)) + .copyImplicitOps(*MBBAUT); + MBB.erase(MBBAUT); + } else { + BuildMI(MBB, MBBAUT, DL, + get(ShouldSignReturnAddrWithAKey ? AArch64::AUTIASP + : AArch64::AUTIBSP)) + .setMIFlag(MachineInstr::FrameDestroy); + } + } + // If this is a tail call outlined function, then there's already a return. if (OF.FrameConstructionID == MachineOutlinerTailCall || OF.FrameConstructionID == MachineOutlinerThunk) diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-diff-scope-same-key.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-diff-scope-same-key.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-diff-scope-same-key.ll @@ -0,0 +1,68 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple \ +; RUN: aarch64-arm-none-eabi %s -o - | FileCheck %s + +define void @a() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @b() "sign-return-address"="non-leaf" { +; CHECK-LABEL: b: // @b +; CHECK-NOT: paciasp +; CHECK-NOT: .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 +; CHECK-NOT: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @c() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +; CHECK-NOT: OUTLINED_FUNCTION_0: +; CHECK-NOT: // -- Begin function diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-non-leaf.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-non-leaf.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-non-leaf.ll @@ -0,0 +1,72 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple \ +; RUN: aarch64-arm-none-eabi %s -o - | FileCheck %s + +define i64 @a(i64 %x) "sign-return-address"="non-leaf" "sign-return-address-key"="b_key" { +; CHECK-LABEL: a: // @a +; CHECK: .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 + call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 + ret i64 %x +} + +define i64 @b(i64 %x) "sign-return-address"="non-leaf" "sign-return-address-key"="b_key" { +; CHECK-LABEL: b: // @b +; CHECK: .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 + call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 + ret i64 %x +} + +define i64 @c(i64 %x) "sign-return-address"="non-leaf" "sign-return-address-key"="b_key" { +; CHECK-LABEL: c: // @c +; CHECK: .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 + call void asm sideeffect "mov x30, $0", "r,~{lr}"(i64 %x) #1 + ret i64 %x +} + +; Outlined function is leaf-function => don't sign it +; CHECK-LABEL: OUTLINED_FUNCTION_0: +; CHECK-NOT: .cfi_b_key_frame +; CHECK-NOT: paci{{[a,b]}}sp +; CHECK-NOT: .cfi_negate_ra_state +; CHECK-NOT: auti{{[a,b]}}sp diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-diff-key.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-diff-key.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-diff-key.ll @@ -0,0 +1,69 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple \ +; RUN: aarch64-arm-none-eabi %s -o - | FileCheck %s + +define void @a() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @b() "sign-return-address"="all" "sign-return-address-key"="b_key" { +; CHECK-LABEL: b: // @b +; CHECK: .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 +; CHECK-NOT: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @c() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +; CHECK-NOT: OUTLINED_FUNCTION_0: +; CHECK-NOT: // -- Begin function diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-a.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-a.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-a.ll @@ -0,0 +1,71 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple \ +; RUN: aarch64-arm-none-eabi %s -o - | FileCheck %s + +define void @a() "sign-return-address"="all" "sign-return-address-key"="a_key" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @b() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +define void @c() "sign-return-address"="all" { +; CHECK-LABEL: 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 +; CHECK: autiasp + ret void +; CHECK: .cfi_endproc +} + +; CHECK-LABEL: OUTLINED_FUNCTION_0: +; CHECK: paciasp +; CHECK-NEXT: .cfi_negate_ra_state +; CHECK: autiasp +; CHECK-NEXT: ret diff --git a/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-b.ll b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-b.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-same-scope-same-key-b.ll @@ -0,0 +1,75 @@ +; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple \ +; RUN: aarch64-arm-none-eabi %s -o - | FileCheck %s + +define void @a() "sign-return-address"="all" "sign-return-address-key"="b_key" { +; CHECK-LABEL: a: // @a +; CHECK: .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 +; CHECK: autibsp + ret void +; CHECK: .cfi_endproc +} + +define void @b() "sign-return-address"="all" "sign-return-address-key"="b_key" { +; CHECK-LABEL: b: // @b +; CHECK: .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 +; CHECK: autibsp + ret void +; CHECK: .cfi_endproc +} + +define void @c() "sign-return-address"="all" "sign-return-address-key"="b_key" { +; CHECK-LABEL: c: // @c +; CHECK: .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 +; CHECK: autibsp + ret void +; CHECK: .cfi_endproc +} + +; CHECK-LABEL: OUTLINED_FUNCTION_0: +; CHECK: .cfi_b_key_frame +; CHECK-NEXT: pacibsp +; CHECK-NEXT: .cfi_negate_ra_state +; CHECK: autibsp +; CHECK-NEXT: ret