[AArch64][MachineOutliner] Return address signing for outlined functions

[tellenbach]

Summary

During AArch64 frame lowering instructions to enable return address
signing are inserted into function if needed. Functions generated during
machine outlining don't run through target frame lowering and hence are
missing such instructions.

This patch introduces the following changes:

1. If not all functions that potentially participate in function outlining agree on their return address signing scope and their return address signing key, outlining is disabled for these functions.
2. If not all functions that potentially participate in function outlining agree on their support for v8.3A features, outlining is disabled for these functions.
3. If all candidate functions agree on the signing scope, signing key and and their support for v8.3 features, the outlined function behaves as if it had the same scope and key attributes and as if it would provide the same v8.3A support as the original functions.

[ostannard]

It looks like you're emitting PAC instructions but not AUT, so I'd expect this to fault when the outlined function returns. Or is some other code already returning using RETAA/RETAB? In which case, the test should check that.

If our callees have "sign-return-address"="non-leaf", I think we should inherit that in the outlined function, meaning that we only need PAC/AUT instructions if the outlined function spills LR.

Emitting no signing instructions when the callees have mismatched keys seems dangerous, maybe we should consider these candidates to be incompatible, and only outline from functions with matching return address signing options?

This is going to add one or two instructions (depending on whether we use RETAA/RETAB) to each outlined function which needs them, which we should consider when deciding if outlining is worth it, otherwise this might end up increasing code size.

[tellenbach]

In D69097#1712743, @ostannard wrote:

It looks like you're emitting PAC instructions but not AUT, so I'd expect this to fault when the outlined function returns. Or is some other code already returning using RETAA/RETAB? In which case, the test should check that.

Yep! AUT instructions are inserted during epilogue emission (again frame lowering). I'll fix that.

If our callees have "sign-return-address"="non-leaf", I think we should inherit that in the outlined function, meaning that we only need PAC/AUT instructions if the outlined function spills LR.

So, if a, b and c participate in outlining, a and b have sign-return-address=all and c has sign-return-address=non-leaf, the new outlined function should behave like having sign-return-address=non-leaf? Or should we only allow outlining if the attributes match?

Emitting no signing instructions when the callees have mismatched keys seems dangerous, maybe we should consider these candidates to be incompatible, and only outline from functions with matching return address signing options?

Sure, that would be certainly the safest possibility. This could include the key types as well as the overall outlining options (see comment above).

This is going to add one or two instructions (depending on whether we use RETAA/RETAB) to each outlined function which needs them, which we should consider when deciding if outlining is worth it, otherwise this might end up increasing code size.

That's a good point. IIRC the machine outliner already provides functionality for measuring "frame overhead".

[ostannard]

So, if a, b and c participate in outlining, a and b have sign-return-address=all and c has sign-return-address=non-leaf, the new outlined function should behave like having sign-return-address=non-leaf? Or should we only allow outlining if the attributes match?

I think for safety we should always pick the strongest attribute from the existing functions, so that case should get sign-return-address=all. Alternatively, we could consider functions with different return address signing attributes to be incompatible, and outline a and b separately from c. I'm not too worried about optimal code generation when mixing attributes, as I expect them to match in almost all real-world use cases (they are set by command-line options, which will probably be consistent across the whole program).

[tellenbach]

Update patch (see Oliver's comments)

[tellenbach]

Fix number of instructions needed for return address signing

[paquette]

I think this needs some more tests. In particular, I'd like to see some test cases that show what happens when

• We tail call
• We save LR to the stack
• We don't save LR to the stack (but emit a bl)
• We get an outlined thunk

Also another thing to consider here is whether or not we should ever outline, say, PAC/AUT instructions. It might be good to mark those as illegal to outline. That will also need a test. I think that writing a MIR test should make this fairly easy to check for.

edit: I had a couple other concerns there, but they were actually already addressed in the patch. :) I think most of my concerns here are to have explicit test cases using MIR.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5059–5060

I find "consensus" kind of confusing to use here. I'd expect it to return true when they agree, not when they disagree. Maybe something like "disagree" would be better to use here?

5626

I'd use IsLeafFunction versus isLeafFunction for the sake of consistency with the rest of the code.

5693–5700

Simplify?

else if (Scope.equals("non-leaf") && !isLeafFunction)
  ShouldSignReturnAddr = true;

5721–5723

Prefer

if (MBBAUT != MBB.end())
  DL = MBBAUT->getDebugLoc();

I think clang-format should probably get this?

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-diff-scope-same-key.ll

1 ↗

(On Diff #225873)

Can you use MIR tests here? That would make it easier to test each frame variant with signing.

[paquette]

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-non-leaf.ll

73 ↗

(On Diff #225873)

add newline

[tellenbach]

Changes to address some of Jessica's comments

[tellenbach]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5059–5060

What you expect is exactly what is happening: outliningCandidatesSigningScopeConsensus(a, b) returns true if a and b *agree* on their return address signing. If they do, the lambda for std::adjacent_find returns false and no differing Candidates are found.

This is basically a pairwise search through the iterator range for unequal candidates (w.r.t. to their return address signing).

However, the logic could be reversed:

if (OutliningCandidatesDisagree)
  return true;

[tellenbach]

In D69097#1716855, @paquette wrote:

I think this needs some more tests. In particular, I'd like to see some test cases that show what happens when

• We tail call
• We save LR to the stack
• We don't save LR to the stack (but emit a bl)
• We get an outlined thunk

I agree, thanks.

Also another thing to consider here is whether or not we should ever outline, say, PAC/AUT instructions. It might be good to mark those as illegal to outline. That will also need a test. I think that writing a MIR test should make this fairly easy to check for.

I think this should never happen anyways: After/before any PAC/AUT instruction a cfi instruction is inserted which is a position (isPosition() returns true). Positions don't get outlined on AArch64 (AArch64InstrInfo.cpp:5565) and the one instruction for the actual signing that remains, is too small to get outlined.

However, making this explicit is probably the better way to do it.

edit: I had a couple other concerns there, but they were actually already addressed in the patch. :) I think most of my concerns here are to have explicit test cases using MIR.

Thanks for your comments!

[tellenbach]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5072

This is actually an assumption: If the candidates agree on "sign-return-address"="non-leaf", the outlined function could potentially get signed. (and the outlined function could potentially increase by two instructions). However, it is not certainly true that it will be no leaf function (and hence needs no signing), This gets determined later when the decision to outline or not is already made.

[tellenbach]

Remove braces after if-stmt

[tellenbach]

• Make test more generic
• If a function has no "sign-return-address" attribute, its signing behaviour is equivalent to a function with the "sign-return-address"="none" attribute

[tellenbach]

• Fix determination of signing scope equality
• Add comments

[tellenbach]

Add .mir test for MachineOutlinerRegSave

[tellenbach]

Add test for outlined thunk

[ostannard]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5065

This comment doesn't make sense, I think you swapped A and B half way through.

5072

I think this would be useful as a comment.

Could we also check if we have the v8.3A instructions here, as using RETAA will reduce the size bloat by 4 bytes?

5080

I think this can be llvm_unreachable, the above code covers all possible combinations.

5121

llvm_unreachable.

[tellenbach]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5072

The problem here is the same: Here we only know whether RETAA/RETAB instructions are available, not if we really can use them. E.g. in the case of a tail calling outlined function we rather emit a branch than a return and therefore can't replace the RET by RETAA/RETAB. The signing falls back to AUTIASP/AUTIBSP in this case.

However, the decision to tail-call or not is made later, when the decision to outline has already been made.

Assuming 2 instructions is therefore the only safe way I see.

[tellenbach]

• Improve support for v8.3a signing instructions
• Add tests
• Address review comments
• Moved signing logic into separate function

[ostannard]

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-v8-3.ll

61 ↗

(On Diff #226904)

I don't think this is correct: @c is being compiled for Armv8.0-A, but we've introduced a call to a function with the retab instruction, so it will fail if run on an v8.0-A machine. I think we need to check that all of the outlining candidate functions are being compiled for v8.3-A, or carefully construct the new function's subtarget to be a subset of the candidate functions.

[tellenbach]

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-v8-3.ll

61 ↗

(On Diff #226904)

I see your point. Maybe this test is just too artificial and in reality functions in one module will all be compiled for the same subtarget? I think the machine outliner is not the correct place to check if this is true.

And if we would like to check it here this would be a more general issue and should probably addressed in a separate patch.

[ostannard]

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-v8-3.ll

61 ↗

(On Diff #226904)

Subtargets can vary per-function, this might happen because the user is doing runtime checks, and dispatching to optimised versions using later architecture features when possible. Fixing this as a separate patch probably makes sense, but it will nee to be committed before this one otherwise we'll start generating invalid code.

I guess this hasn't been a problem for the outliner until now, because doesn't introduce any instructions not present in v8.0-A.

[tellenbach]

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-v8-3.ll

61 ↗

(On Diff #226904)

Okay, I'm going to disable outlining for v8.3A and non-v8.3A subtargets (if the outlined function would be signed). If this needs to be extended to other non-matching subtargets needs more investigation. I definitely don't want to overshoot here and disable outlining for too many functions.

[tellenbach]

• Disable outlining for functions that disagree on their v8.3A feature support
• Add a test for the above bullet point

[ostannard]

LGTM, thanks!

[tellenbach]

Thanks for the great review!

[ostannard]

I've reverted this (rGa3f474542) because it is causing failures when an instruction which modifies SP gets outlined. here's a reproducer:

int *volatile v;

void foo() {
  int a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
}

void bar() {
  int a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
  v = &a;
}

Which gets compiled to:

$ /work/llvm/build/bin/clang --target=aarch64--none-eabi -march=armv8.3-a -c test2.c -o - -S -Oz -mbranch-protection=pac-ret+leaf 
        .text
        .file   "test2.c"
        .globl  foo                     // -- Begin function foo
        .p2align        2
        .type   foo,@function
foo:                                    // @foo
// %bb.0:                               // %entry
        paciasp
        sub     sp, sp, #16             // =16
        mov     x0, x30
        bl      OUTLINED_FUNCTION_0
        mov     x30, x0
        retaa
.Lfunc_end0:
        .size   foo, .Lfunc_end0-foo
                                        // -- End function
        .globl  bar                     // -- Begin function bar
        .p2align        2
        .type   bar,@function
bar:                                    // @bar
// %bb.0:                               // %entry
        paciasp
        sub     sp, sp, #16             // =16
        mov     x0, x30
        bl      OUTLINED_FUNCTION_0
        mov     x30, x0
        retaa
.Lfunc_end1:
        .size   bar, .Lfunc_end1-bar
                                        // -- End function
        .p2align        2               // -- Begin function OUTLINED_FUNCTION_0
        .type   OUTLINED_FUNCTION_0,@function
OUTLINED_FUNCTION_0:                    // @OUTLINED_FUNCTION_0
        .cfi_sections .debug_frame
        .cfi_startproc
// %bb.0:
        paciasp
        .cfi_negate_ra_state
        adrp    x8, v
        add     x9, sp, #12             // =12
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        str     x9, [x8, :lo12:v]
        add     sp, sp, #16             // =16
        retaa
.Lfunc_end2:
        .size   OUTLINED_FUNCTION_0, .Lfunc_end2-OUTLINED_FUNCTION_0
        .cfi_endproc
                                        // -- End function
        .type   v,@object               // @v
        .comm   v,8,8

        .ident  "clang version 10.0.0 (https://github.com/llvm/llvm-project.git 7849862f46933306454342b0e8ee05e4e6806646)"
        .section        ".note.GNU-stack","",@progbits
        .addrsig
        .addrsig_sym v

The problem is the instruction add sp, sp, #16 in OUTLINED_FUNCTION_0. This causes SP to have a different value for the paciasp and retaa instructions, so the signature does not match, and the return causes a fault.

I think we need some additional checks to avoid outlining anything which modifies SP when doing return address signing (though it would be valid to outline a balanced sub/add pair).

[tellenbach]

@ostannard wow, good catch! Thanks for reverting this.

$ clang --target=aarch64--none-eabi -c test2.c -o - -mllvm -stop-before=machine-outliner
name:            bar
body:             |
  bb.0.entry:
    $sp = frame-setup SUBXri $sp, 16, 0
    frame-setup CFI_INSTRUCTION def_cfa_offset <mcsymbol .Ltmp1>16
    $x8 = ADRP target-flags(aarch64-page) @v
    renamable $x8 = ADDXri $x8, target-flags(aarch64-pageoff, aarch64-nc) @v, 0
    $x9 = ADDXri $sp, 12, 0
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui renamable $x9, renamable $x8, 0 :: (volatile store 8 into @v)
    STRXui killed renamable $x9, killed renamable $x8, 0 :: (volatile store 8 into @v)
    $sp = frame-destroy ADDXri $sp, 16, 0
    RET undef $lr

Since CFI instructions don't get outlined things get messed up here. If possible I'd like to avoid to disable outlining for all functions that alter sp. So I'll probably try to find matching sp changing instructions.

[tellenbach]

I've opened the following new revision that includes a fix for the bug in this patch and added @ostannard and @paquette again as reviewers: https://reviews.llvm.org/D70635

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

AArch64InstrInfo::getOutliningCandidateInfo(
   unsigned FlagsSetInAll = 0xF;
 
   // Compute liveness information for each candidate, and set FlagsSetInAll.
   const TargetRegisterInfo &TRI = getRegisterInfo();
   std::for_each(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
                 [&FlagsSetInAll](outliner::Candidate &C) {
                   FlagsSetInAll &= C.Flags;
                 });
 
   // According to the AArch64 Procedure Call Standard, the following are
   // undefined on entry/exit from a function call:
   //
   // * Registers x16, x17, (and thus w16, w17)
   // * Condition codes (and thus the NZCV register)
   //
   // Because if this, we can't outline any sequence of instructions where
   // one
   // of these registers is live into/across it. Thus, we need to delete
   // those
   // candidates.
   auto CantGuaranteeValueAcrossCall = [&TRI](outliner::Candidate &C) {
     // If the unsafe registers in this block are all dead, then we don't need
     // to compute liveness here.
     if (C.Flags & UnsafeRegsDead)
       return false;
     C.initLRU(TRI);
     LiveRegUnits LRU = C.LRU;
     return (!LRU.available(AArch64::W16) || !LRU.available(AArch64::W17) ||
             !LRU.available(AArch64::NZCV));
   };
 
   // Are there any candidates where those registers are live?
   if (!(FlagsSetInAll & UnsafeRegsDead)) {
     // Erase every candidate that violates the restrictions above. (It could be
     // true that we have viable candidates, so it's not worth bailing out in
     // the case that, say, 1 out of 20 candidates violate the restructions.)
     RepeatedSequenceLocs.erase(std::remove_if(RepeatedSequenceLocs.begin(),
                                               RepeatedSequenceLocs.end(),
bool HasBTI = any_of(RepeatedSequenceLocs, [](outliner::Candidate &C) {
     return C.getMF()->getFunction().hasFnAttribute("branch-target-enforcement");
   });
 
   // Returns true if an instructions is safe to fix up, false otherwise.
   auto IsSafeToFixup = [this, &TRI](MachineInstr &MI) {
     if (MI.isCall())
       return true;
 
     if (!MI.modifiesRegister(AArch64::SP, &TRI) &&
         !MI.readsRegister(AArch64::SP, &TRI))
       return true;
 
     // Any modification of SP will break our code to save/restore LR.
     // FIXME: We could handle some instructions which add a constant
     // offset to SP, with a bit more work.
     if (MI.modifiesRegister(AArch64::SP, &TRI))
       return false;
for (MachineInstr &MI : MBB) {
     // We've pushed the return address to the stack, so add 16 to the offset.
     // This is safe, since we already checked if it would overflow when we
     // checked if this instruction was legal to outline.
     int64_t NewImm = (Offset + 16) / Scale;
     StackOffsetOperand.setImm(NewImm);
   }
 }
 
+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) {
     MachineInstr *Call = &*--MBB.instr_end();
     unsigned TailOpcode;
     if (Call->getOpcode() == AArch64::BL) {
       TailOpcode = AArch64::TCRETURNdi;
     } else {
       assert(Call->getOpcode() == AArch64::BLR);
       TailOpcode = AArch64::TCRETURNriALL;
     }
     MachineInstr *TC = BuildMI(MF, DebugLoc(), get(TailOpcode))
-                            .add(Call->getOperand(0))
-                            .addImm(0);
     MBB.insert(MBB.end(), TC);
     Call->eraseFromParent();
   }
 
   // Is there a call in the outlined range?
   auto IsNonTailCall = [](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.
     assert(OF.FrameConstructionID != MachineOutlinerDefault &&
            "Can only fix up stack references once");
if (OF.FrameConstructionID == MachineOutlinerTailCall ||
       OF.FrameConstructionID == MachineOutlinerThunk)
     return;
 
   // It's not a tail call, so we have to insert the return ourselves.
   MachineInstr *ret = BuildMI(MF, DebugLoc(), get(AArch64::RET))
                           .addReg(AArch64::LR, RegState::Undef);
   MBB.insert(MBB.end(), ret);
 
   // Did we have to modify the stack by saving the link register?
   if (OF.FrameConstructionID != MachineOutlinerDefault)
     return;
 
   // We modified the stack.
   // Walk over the basic block and fix up all the stack accesses.
   fixupPostOutline(MBB);
 }
 
 MachineBasicBlock::iterator AArch64InstrInfo::insertOutlinedCall(
     Module &M, MachineBasicBlock &MBB, MachineBasicBlock::iterator &It,
     MachineFunction &MF, const outliner::Candidate &C) const {
 
   // Are we tail calling?
   if (C.CallConstructionID == MachineOutlinerTailCall) {
     // If yes, then we can just branch to the label.
     It = MBB.insert(It, BuildMI(MF, DebugLoc(), get(AArch64::TCRETURNdi))
                             .addGlobalAddress(M.getNamedValue(MF.getName()))
                             .addImm(0));
     return It;
   }
 
   // Are we saving the link register?
   if (C.CallConstructionID == MachineOutlinerNoLRSave ||
       C.CallConstructionID == MachineOutlinerThunk) {
     // No, so just insert the call.
     It = MBB.insert(It, BuildMI(MF, DebugLoc(), get(AArch64::BL))
                             .addGlobalAddress(M.getNamedValue(MF.getName())));
     return It;
   }
 
   // We want to return the spot where we inserted the call.
   MachineBasicBlock::iterator CallPt;
 
   // Instructions for saving and restoring LR around the call instruction we're
   // going to insert.
   MachineInstr *Save;
   MachineInstr *Restore;
   // Can we save to a register?

llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-akey.ll

+; 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
+

llvm/test/CodeGen/AArch64/machine-outliner-sign-return-address-bkey.ll

+; 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
+