[llvm] [AArch64][PAC] Refactor aarch64-ptrauth pass for better extensibility (PR #70446)

[Anatoly Trosinenko via llvm-commits]

https://github.com/atrosinenko updated https://github.com/llvm/llvm-project/pull/70446

>From 7142937a7534869fc0d75f73a55c8034b5c15ecd Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Thu, 26 Oct 2023 17:56:38 +0300
Subject: [PATCH 1/2] [AArch64][PAC] Refactor aarch64-ptrauth pass for better
 extensibility

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.
---
 .../lib/Target/AArch64/AArch64PointerAuth.cpp | 51 ++++++++++++-------
 1 file changed, 32 insertions(+), 19 deletions(-)

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

>From 85236cf75ee468c18a336076ff93a68926b85b5c Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Fri, 3 Nov 2023 14:24:19 +0300
Subject: [PATCH 2/2] Add a clarifying comment

---
 llvm/lib/Target/AArch64/AArch64PointerAuth.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 0ae888cbaaa4612..7576d2a899d1afb 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -317,6 +317,8 @@ bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
         // 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.
+        // Skip the BUNDLE instruction itself (actual bundled instructions
+        // follow it in the instruction list).
         if (MI.isBundle())
           continue;
         if (AArch64InstrInfo::isTailCallReturnInst(MI))