[AArch64][PAC] Refactor aarch64-ptrauth pass for better extensibility

[atrosinenko]

Refactor Pointer Authentication pass in preparation for adding more PAUTH_* pseudo instructions:

• dropped early return from runOnMachineFunction() as other PAUTH_* instructions need expansion even when pac-ret is disabled
• refactored runOnMachineFunction() to first collect all the instructions of interest without modifying anything and then performing changes in the later loops. There are two types of relevant instructions: PAUTH_* pseudos that should definitely be replaced by this pass and tail call instructions that may require attention if pac-ret is enabled
• made the loop iterating over all of the instructions handle instruction bundles by itself: even though this pass still does not support bundled TCRETURN* instructions (such as produced by KCFI) it does not crash anymore when no support is actually required

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

Refactor Pointer Authentication pass in preparation for adding more PAUTH_* pseudo instructions.

Fix handling of bundled TCRETURN* instructions (known to be generated by KCFI). As other PAUTH_* instructions may need expansion even when pac-ret is disabled, it is not generally possible to skip the whole function easily. While this pass still does not support pac-ret being enabled at the same time with KCFI, it should not crash if no checks are actually emitted prior to TCRETURN instruction.

---

Full diff: https://github.com/llvm/llvm-project/pull/70446.diff

1 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.cpp (+32-19)

diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 5d11f0d22574c90..0ae888cbaaa4612 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -297,52 +297,65 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
 
 bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
   const auto *MFnI = MF.getInfo<AArch64FunctionInfo>();
-  if (!MFnI->shouldSignReturnAddress(true))
-    return false;
 
   Subtarget = &MF.getSubtarget<AArch64Subtarget>();
   TII = Subtarget->getInstrInfo();
   TRI = Subtarget->getRegisterInfo();
 
-  SmallVector<MachineBasicBlock::iterator> DeletedInstrs;
-  SmallVector<MachineBasicBlock::iterator> TailCallInstrs;
+  SmallVector<MachineBasicBlock::instr_iterator> PAuthPseudoInstrs;
+  SmallVector<MachineBasicBlock::instr_iterator> TailCallInstrs;
 
   bool Modified = false;
   bool HasAuthenticationInstrs = false;
 
   for (auto &MBB : MF) {
-    for (auto &MI : MBB) {
-      auto It = MI.getIterator();
+    // Using instr_iterator to catch unsupported bundled TCRETURN* instructions
+    // instead of just skipping them.
+    for (auto &MI : MBB.instrs()) {
       switch (MI.getOpcode()) {
       default:
+        // Bundled TCRETURN* instructions (such as created by KCFI)
+        // are not supported yet, but no support is required if no
+        // PAUTH_EPILOGUE instructions exist in the same function.
+        if (MI.isBundle())
+          continue;
         if (AArch64InstrInfo::isTailCallReturnInst(MI))
-          TailCallInstrs.push_back(It);
+          TailCallInstrs.push_back(MI.getIterator());
         break;
       case AArch64::PAUTH_PROLOGUE:
-        signLR(MF, It);
-        DeletedInstrs.push_back(It);
-        Modified = true;
-        break;
       case AArch64::PAUTH_EPILOGUE:
-        authenticateLR(MF, It);
-        DeletedInstrs.push_back(It);
-        Modified = true;
-        HasAuthenticationInstrs = true;
+        assert(!MI.isBundled());
+        PAuthPseudoInstrs.push_back(MI.getIterator());
         break;
       }
     }
   }
 
+  for (auto It : PAuthPseudoInstrs) {
+    switch (It->getOpcode()) {
+    case AArch64::PAUTH_PROLOGUE:
+      signLR(MF, It);
+      break;
+    case AArch64::PAUTH_EPILOGUE:
+      authenticateLR(MF, It);
+      HasAuthenticationInstrs = true;
+      break;
+    default:
+      llvm_unreachable("Unhandled opcode");
+    }
+    It->eraseFromParent();
+    Modified = true;
+  }
+
   // FIXME Do we need to emit any PAuth-related epilogue code at all
   //       when SCS is enabled?
   if (HasAuthenticationInstrs &&
       !MFnI->needsShadowCallStackPrologueEpilogue(MF)) {
-    for (auto TailCall : TailCallInstrs)
+    for (auto TailCall : TailCallInstrs) {
+      assert(!TailCall->isBundled() && "Not yet supported");
       Modified |= checkAuthenticatedLR(TailCall);
+    }
   }
 
-  for (auto MI : DeletedInstrs)
-    MI->eraseFromParent();
-
   return Modified;
 }

[atrosinenko]

Rewritten the description to clarify the particular changes and added one more inline comment.

Ping.

[atrosinenko]

Ping

[pratlucas]

LGTM.

[atrosinenko]

Thank you! I will merge it after rebasing and retesting.