[AArch64] Fix a machine verifier ICE when using __builtin_return_address with return address sign and authentication

[chill]

Summary

When using PAuth instructions in the NOP-space, the compiler can emit an
AUTIxSP instruction in the function epilogue.

In a certain test case, using __builtin_return_address, the compiler causes
the machine verifier to exit with an error after turning

%0 = COPY LR
x0 = COPY killed %0

into

x0 = COPY killed LR

and the backend later emitting

x0 = COPY killed LR
AUTIASP implicit def LR, implicit LR

This is obviously incorrect since LR is live beyond the COPY instruction.

The issue is that lifetime of the return address is not tracked properly. The
backend relies on the callee-saved registers mechanism to ensure LR is
restored before used for return. For the return instruction that's handled with
a little hack to appease the verifier.

This patch employs the same hack for AUTIxSP, marking the implicit LR
operand as undef.

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

static void InsertReturnAddressAuth(MachineFunction &MF,
   // instructions, namely RETA{A,B}, that can be used instead.
   if (Subtarget.hasV8_3aOps() && MBBI != MBB.end() &&
       MBBI->getOpcode() == AArch64::RET_ReallyLR) {
     BuildMI(MBB, MBBI, DL,
             TII->get(ShouldSignWithAKey(MF) ? AArch64::RETAA : AArch64::RETAB))
         .copyImplicitOps(*MBBI);
     MBB.erase(MBBI);
   } else {
-    BuildMI(
-        MBB, MBBI, DL,
-        TII->get(ShouldSignWithAKey(MF) ? AArch64::AUTIASP : AArch64::AUTIBSP))
+    MachineInstrBuilder B =
+        BuildMI(MBB, MBBI, DL,
+                TII->get(ShouldSignWithAKey(MF) ? AArch64::AUTIASP
+                                                : AArch64::AUTIBSP))
             .setMIFlag(MachineInstr::FrameDestroy);
+    // Similarly to way we handle LR use in RET_ReallyLR (see comment in
+    // AArch64ExpandPseudeInstrs.cpp), rely on callee-saved register handling to
+    // ensure LR is restored before we emit the AUTIxSP, but set the undef flag
+    // here to appease the MachineVerifier liveness checks.
+    for (MachineOperand &MO : B.getInstr()->operands()) {
+      if (!MO.isReg() || MO.isDef())
+        continue;
+      if (MO.getReg() != AArch64::LR)
+        continue;
+      MO.setIsUndef();
+    }
   }
 }
 
 static bool isFuncletReturnInstr(const MachineInstr &MI) {
   switch (MI.getOpcode()) {
   default:
     return false;
   case AArch64::CATCHRET:

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

if (ShouldSignReturnAddr) {
     if (Subtarget.hasV8_3aOps() && MBBAUT != MBB.end() &&
         MBBAUT->getOpcode() == AArch64::RET) {
       BuildMI(MBB, MBBAUT, DL,
               TII->get(ShouldSignReturnAddrWithAKey ? AArch64::RETAA
                                                     : AArch64::RETAB))
           .copyImplicitOps(*MBBAUT);
       MBB.erase(MBBAUT);
     } else {
+      MachineInstrBuilder B =
           BuildMI(MBB, MBBAUT, DL,
                   TII->get(ShouldSignReturnAddrWithAKey ? AArch64::AUTIASP
                                                         : AArch64::AUTIBSP))
               .setMIFlag(MachineInstr::FrameDestroy);
+      // Similarly to way we handle LR use in RET_ReallyLR (see comment in
+      // AArch64ExpandPseudeInstrs.cpp), rely on callee-saved register handling
+      // to ensure LR is restored before we emit the AUTIxSP, but set the undef
+      // flag here to appease the MachineVerifier liveness checks.
+      for (MachineOperand &MO : B.getInstr()->operands()) {
+        if (!MO.isReg() || MO.isDef())
+          continue;
+        if (MO.getReg() != AArch64::LR)
+          continue;
+        MO.setIsUndef();
+      }
     }
   }
 }
 
 void AArch64InstrInfo::buildOutlinedFrame(
     MachineBasicBlock &MBB, MachineFunction &MF,
     const outliner::OutlinedFunction &OF) const {
 

llvm/test/CodeGen/AArch64/builtin-return-address-autisap.ll

+;; RUN: llc  -mtriple aarch64 %s -o - | FileCheck %s
+define i8* @f() #0 {
+entry:
+  %0 = tail call i8* @llvm.returnaddress(i32 0)
+  ret i8* %0
+}
+
+declare i8* @llvm.returnaddress(i32 immarg) #1
+
+attributes #0 = { nounwind readnone  "sign-return-address"="all" "sign-return-address-key"="a_key" }
+attributes #1 = { nounwind readnone }
+
+;; CHECK-LABEL: f:
+;; CHECK:      mov x0, x30
+;; CHECK-NEXT: hint #29
+;; CHECK-NEXT: ret

llvm/test/CodeGen/AArch64/machine-outliner-retaddr-sign-sp-mod.mir

liveins: $lr
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui killed renamable $x9, killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     $sp = frame-destroy ADDXri $sp, 16, 0
-    frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+    frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
     RET undef $lr
 
 # CHECK-LABEL:    name:            legal0
 # CHECK:          body:             |
 # CHECK-NEXT:         bb.0 (%ir-block.0):
 # CHECK-NEXT:           liveins: $lr
 # CHECK:                frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
 # CHECK-NEXT:           frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK:                BL @[[OUTLINED_FUNC:OUTLINED_FUNCTION_[0-9]+]]
-# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 # CHECK-NEXT:           RET undef $lr
 
 ...
 ---
 name:            legal1
 tracksRegLiveness: true
 body:             |
   bb.0 (%ir-block.0):
   liveins: $lr
     frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
     frame-setup CFI_INSTRUCTION negate_ra_sign_state
     $sp = frame-setup SUBXri $sp, 16, 0
     renamable $x8 = ADRP target-flags(aarch64-page) @v
     $x9 = ADDXri $sp, 12, 0
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui killed renamable $x9, killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     $sp = frame-destroy ADDXri $sp, 16, 0
-    frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+    frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
     RET undef $lr
 
 # CHECK-LABEL:    name:            legal1
 # CHECK:          body:             |
 # CHECK-NEXT:         bb.0 (%ir-block.0):
 # CHECK-NEXT:           liveins: $lr
 # CHECK:                frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
 # CHECK-NEXT:           frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK:                BL @[[OUTLINED_FUNC]]
-# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 # CHECK-NEXT:           RET undef $lr
 
 ...
 ---
 name:            illegal0
 tracksRegLiveness: true
 body:             |
   bb.0 (%ir-block.0):
   liveins: $lr
     frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
     frame-setup CFI_INSTRUCTION negate_ra_sign_state
     $sp = frame-setup SUBXri $sp, 16, 0
     renamable $x8 = ADRP target-flags(aarch64-page) @v
     $x9 = ADDXri $sp, 12, 0
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui killed renamable $x9, killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     $sp = frame-destroy ADDXri $sp, 12, 0
-    frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+    frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
     RET undef $lr
 
 ...
 ---
 name:            illegal1
 tracksRegLiveness: true
 body:             |
   bb.0 (%ir-block.0):
   liveins: $lr
     frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
     frame-setup CFI_INSTRUCTION negate_ra_sign_state
     $sp = frame-setup SUBXri $sp, 16, 0
     renamable $x8 = ADRP target-flags(aarch64-page) @v
     $x9 = ADDXri $sp, 12, 0
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui renamable $x9, renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     STRXui killed renamable $x9, killed renamable $x8, target-flags(aarch64-pageoff, aarch64-nc) @v :: (volatile store 8 into @v)
     $sp = frame-destroy ADDXri $sp, 12, 0
-    frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+    frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
     RET undef $lr
 
 # CHECK-LABEL:    name:            illegal0
 # CHECK:          body:             |
 # CHECK-NEXT:         bb.0 (%ir-block.0):
 # CHECK-NEXT:           liveins: $lr
 # CHECK:                frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
 # CHECK-NEXT:           frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK-NOT:            BL @OUTLINED_FUNCTION_{{.*}}
-# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 # CHECK-NEXT:           RET undef $lr
 
 # CHECK-LABEL:    name:            illegal1
 # CHECK:          body:             |
 # CHECK-NEXT:         bb.0 (%ir-block.0):
 # CHECK-NEXT:           liveins: $lr
 # CHECK:                frame-setup PACIASP implicit-def $lr, implicit killed $lr, implicit $sp
 # CHECK-NEXT:           frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK-NOT:            BL @OUTLINED_FUNCTION_{{.*}}
-# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit killed $lr, implicit $sp
+# CHECK:                frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 # CHECK-NEXT:           RET undef $lr
 
 # Outlined function that contains only legal sp modifications
 # CHECK:          name:            [[OUTLINED_FUNC]]
 # CHECK:          body:             |
 # CHECK-NEXT:       bb.0:
 # CHECK-NEXT: liveins: $lr
 # CHECK-NEXT: {{^  $}}
 # CHECK-NEXT:         frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp
 # CHECK-NEXT:         frame-setup CFI_INSTRUCTION negate_ra_sign_state
 # CHECK-NEXT:         $sp = frame-setup SUBXri $sp, 16, 0
 # CHECK:              $sp = frame-destroy ADDXri $sp, 16, 0
-# CHECK-NEXT:         frame-destroy AUTIASP implicit-def $lr, implicit $lr, implicit $sp
+# CHECK-NEXT:         frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 # CHECK-NEXT:         RET $lr

llvm/test/CodeGen/MIR/AArch64/return-address-signing.mir

 #CHECK: foo
 name:            foo
 alignment:       4
 tracksRegLiveness: true
 frameInfo:
   maxCallFrameSize: 0
 #CHECK:    frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp
 #CHECK:    frame-setup CFI_INSTRUCTION negate_ra_sign_state
-#CHECK:    frame-destroy AUTIASP implicit-def $lr, implicit $lr, implicit $sp
+#CHECK:    frame-destroy AUTIASP implicit-def $lr, implicit undef $lr, implicit $sp
 body:             |
   bb.0.entry:
     $w0 = MOVi32imm 2
     RET_ReallyLR implicit killed $w0
 
 ...
 ---
 #CHECK: bar
 name:            bar
 alignment:       4
 tracksRegLiveness: true
 frameInfo:
   maxCallFrameSize: 0
 #CHECK:    frame-setup EMITBKEY
 #CHECK:    frame-setup PACIBSP implicit-def $lr, implicit $lr, implicit $sp
 #CHECK:    frame-setup CFI_INSTRUCTION negate_ra_sign_state
-#CHECK:    frame-destroy AUTIBSP implicit-def $lr, implicit $lr, implicit $sp
+#CHECK:    frame-destroy AUTIBSP implicit-def $lr, implicit undef $lr, implicit $sp
 body:             |
   bb.0.entry:
     $w0 = MOVi32imm 2
     RET_ReallyLR implicit killed $w0
 
 ...