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

[Jon Roelofs via llvm-branch-commits]

================
@@ -11771,23 +11771,75 @@ 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();
   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());
+  RegScavenger RS;
+  RS.enterBasicBlockEnd(MBB);
+  RS.backward(InsertPt);
+
+  // Find out which scratch registers have to be spilled to other GPRs,
+  // mark the rest as unavailable to be spilled-to.
+  SmallVector<std::pair<Register, Register>, 3> Spills;
+  for (Register ScratchReg : ImplicitDefs) {
+    // With Speculative Load Hardening, X16 is reported as reserved by MRI,
+    // but a fallback to DSB+ISB is actually supported as long as implicit-defs
+    // are set appropriately.
+    if (ScratchReg == AArch64::X16 &&
+        MF.getFunction().hasFnAttribute(Attribute::SpeculativeLoadHardening))
+      continue;
+
+    assert(!MF.getRegInfo().isReserved(ScratchReg));
+
+    if (RS.isRegUsed(ScratchReg))
+      Spills.emplace_back(ScratchReg, AArch64::NoRegister);
+    else
+      RS.setRegUsed(ScratchReg);
   }
 
-  if (AFI->branchProtectionPAuthLR() && !Subtarget.hasPAuthLR())
-    Builder.addReg(AArch64::X16, RegState::ImplicitDefine);
+  for (auto &[ScratchReg, SpillReg] : Spills) {
+    (void)ScratchReg;
+    SpillReg = RS.FindUnusedReg(&AArch64::GPR64RegClass);
+    if (!SpillReg)
+      reportFatalUsageError(
+          "Cannot insert PAUTH_EPILOGUE: ran out of registers");
+    RS.setRegUsed(SpillReg);
+  }
+
+  auto EmitMOV = [&](Register DstReg, Register SrcReg) {
+    BuildMI(MBB, InsertPt, DL, get(AArch64::ORRXrs), DstReg)
+        .addReg(AArch64::XZR)
+        .addReg(SrcReg)
+        .addImm(0)
+        .setMIFlag(MachineInstr::FrameDestroy);
+  };
+
+  for (auto [ScratchReg, SpillReg] : Spills)
+    EmitMOV(SpillReg, ScratchReg);
+
+  auto Builder = BuildMI(MBB, InsertPt, DL, get(AArch64::PAUTH_EPILOGUE))
+                     .setMIFlag(MachineInstr::FrameDestroy);
----------------
jroelofs wrote:

I think this conflicts a bit with some of the lowerings we're trying to do in `AArch64PointerAuth::authenticateLR`. In particular, when there are spills, we need to be more careful around deciding when the terminator is combinable, since we can't move the aut after the mov's that restore the spilled regs.

I was trying to address some related issues in https://github.com/llvm/llvm-project/pull/218783, but now I realize that approach is also flawed. We should treat the location of the `PAUTH_EPILOGUE` as ground-truth for where the aut may be inserted, and only allow combining with the terminator when the pseudo is immediately followed by the call/ret. I also think it would simplify things a bit in `authenticateLR` if the pseudo took an immediate for an offset from SP to recover the SP as of `PAUTH_PROLOGUE`.

I think what you're doing here is the right thing, FWIW, so I'll close & rework #218783.

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