[llvm-branch-commits] [llvm] [AArch64][PAC] Prevent PAUTH_EPILOGUE from overwriting live registers (PR #220191)

[Anatoly Trosinenko via llvm-branch-commits]

================
@@ -11771,23 +11771,82 @@ unsigned llvm::getBLRCallOpcode(const MachineFunction &MF) {
 void AArch64InstrInfo::createPauthEpilogueInstr(MachineBasicBlock &MBB,
                                                 DebugLoc DL) const {
   MachineBasicBlock::iterator InsertPt = MBB.getFirstTerminator();
-  auto Builder = BuildMI(MBB, InsertPt, DL, get(AArch64::PAUTH_EPILOGUE))
-                     .setMIFlag(MachineInstr::FrameDestroy);
-
   MachineFunction &MF = *MBB.getParent();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
   const auto *AFI = MF.getInfo<AArch64FunctionInfo>();
   auto &AFL = *static_cast<const AArch64FrameLowering *>(
       MF.getSubtarget().getFrameLowering());
+
+  SmallVector<Register, 3> ImplicitDefs;
   if (AFL.getArgumentStackToRestore(MF, MBB)) {
-    Builder.addReg(AArch64::X17, RegState::ImplicitDefine);
-    Builder.addReg(AArch64::X16, RegState::ImplicitDefine);
+    ImplicitDefs.push_back(AArch64::X17);
+    ImplicitDefs.push_back(AArch64::X16);
     if (AFI->branchProtectionPAuthLR())
-      Builder.addReg(AArch64::X15, RegState::ImplicitDefine);
-    return;
+      ImplicitDefs.push_back(AArch64::X15);
+  } else if (AFI->branchProtectionPAuthLR() && !Subtarget.hasPAuthLR()) {
+    ImplicitDefs.push_back(AArch64::X16);
+  }
+
+  // If the scratch registers we plan using are alive at this point,
+  // try spilling them to other registers.
+
+  assert(MF.getProperties().hasTracksLiveness());
+  LivePhysRegs LiveRegs(TRI);
+  LiveRegs.addLiveOuts(MBB);
+  for (auto &MI : llvm::reverse(llvm::make_range(InsertPt, MBB.end())))
----------------
atrosinenko wrote:

> Is this intentionally skipping InsertPt itself

No, my intention is the opposite: to iterate over the "suffix" of MBB in reverse order and stop right after processing the `InsertPt` (i.e. at the insertion point of `PAUTH_EPILOGUE` right before `InsertPt`). Intuitively, I construct a range from `InsertPt` (including) to `MBB.end()` (excluding), and then iterate over the same set of elements but in reverse order. Just in case, I tried artificially inserting a few `RET`s before initializing `InsertPt = MBB.getFirstTerminator()` and modifying this loop like this:
```cpp
  dbgs() << "###\nMBB is\n";
  MBB.dump();
  dbgs() << "InsertPt is "; InsertPt->dump();
  dbgs() << "---\n";
  for (auto &MI : llvm::reverse(llvm::make_range(InsertPt, MBB.end()))) {
    MI.dump();
    LiveRegs.stepBackward(MI);
  }
  dbgs() << "===\n";
```
This seems to confirm my intuition:
```
###
MBB is
bb.1.if.then:
; predecessors: %bb.0
  successors: %bb.2(0x80000000); %bb.2(100.00%)
  liveins: $x16, $x17, $lr
  frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit $lr, implicit $sp
  early-clobber $sp = frame-setup STRXpre killed $lr, $sp(tied-def 0), -16 :: (store (s64) into %stack.0)
  frame-setup CFI_INSTRUCTION def_cfa_offset 16
  frame-setup CFI_INSTRUCTION offset $w30, -16
  INLINEASM &"mov x30, 12345" [sideeffect] [attdialect], $0:[clobber], implicit-def dead early-clobber $lr
  early-clobber $sp, $lr = frame-destroy LDRXpost $sp(tied-def 0), 16 :: (load (s64) from %stack.0)
  [ edit: PAUTH_EPILOGUE is to be inserted here ]
  RET $x1
  RET $x2
  RET $x3
InsertPt is   RET $x1
---
  RET $x3
  RET $x2
  RET $x1
===
```

https://github.com/llvm/llvm-project/pull/220191