[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

5116–5117

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?

5634

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

5706–5713

Simplify?

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

5734–5736

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

2

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

add newline

[tellenbach]

Changes to address some of Jessica's comments

[tellenbach]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5116–5117

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

5129

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

5064

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

5079

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

5120

llvm_unreachable.

5129

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?

[tellenbach]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

5129

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

if (!ARI->isReservedReg(*MF, Reg) &&
         C.LRU.available(Reg) && C.UsedInSequence.available(Reg))
       return Reg;
   }
 
   // No suitable register. Return 0.
   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") &&
+      !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<outliner::Candidate> &RepeatedSequenceLocs) const {
   outliner::Candidate &FirstCand = RepeatedSequenceLocs[0];
   unsigned SequenceSize =
       std::accumulate(FirstCand.front(), std::next(FirstCand.back()), 0,
                       [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;
 
   // Compute liveness information for each candidate, and set FlagsSetInAll.
   const TargetRegisterInfo &TRI = getRegisterInfo();
   std::for_each(RepeatedSequenceLocs.begin(), RepeatedSequenceLocs.end(),
                 [&FlagsSetInAll](outliner::Candidate &C) {
outliner::OutlinedFunction AArch64InstrInfo::getOutliningCandidateInfo(
   // Helper lambda which sets call information for every candidate.
   auto SetCandidateCallInfo =
       [&RepeatedSequenceLocs](unsigned CallID, unsigned NumBytesForCall) {
         for (outliner::Candidate &C : RepeatedSequenceLocs)
           C.setCallInfo(CallID, NumBytesForCall);
       };
 
   unsigned FrameID = MachineOutlinerDefault;
-  unsigned NumBytesToCreateFrame = 4;
+  NumBytesToCreateFrame += 4;
 
   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())
for (MachineInstr &MI : MBB) {
     int64_t NewImm = (Offset + 16) / Scale;
     StackOffsetOperand.setImm(NewImm);
   }
 }
 
 void AArch64InstrInfo::buildOutlinedFrame(
     MachineBasicBlock &MBB, MachineFunction &MF,
     const outliner::OutlinedFunction &OF) const {
+
   // 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();
   }
 
+  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.
     assert(OF.FrameConstructionID != MachineOutlinerDefault &&
            "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);
 
     MachineBasicBlock::iterator It = MBB.begin();
     MachineBasicBlock::iterator Et = MBB.end();
 
     if (OF.FrameConstructionID == MachineOutlinerTailCall ||
         OF.FrameConstructionID == MachineOutlinerThunk)
if (std::any_of(MBB.instr_begin(), MBB.instr_end(), IsNonTailCall)) {
     MachineInstr *LDRXpost = BuildMI(MF, DebugLoc(), get(AArch64::LDRXpost))
                                  .addReg(AArch64::SP, RegState::Define)
                                  .addReg(AArch64::LR, RegState::Define)
                                  .addReg(AArch64::SP)
                                  .addImm(16);
     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<AArch64Subtarget>();
+    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)
     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);

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

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

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

+; 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:            pacibsp
+; CHECK-NOT:            .cfi_negate_ra_state
+; CHECK-NOT:            autibsp

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

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

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

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

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

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