[AArch64][PAC] Expand blend(reg, imm) operation in aarch64-pauth pass

[atrosinenko]

In preparation for implementing code generation for more @llvm.ptrauth.* intrinsics, move the expansion of blend(register, small integer) variant of @llvm.ptrauth.blend to the AArch64PointerAuth pass, where most other PAuth-related code generation takes place.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Anatoly Trosinenko (atrosinenko)

Changes

In preparation for implementing code generation for more @llvm.ptrauth.* intrinsics, move the expansion of blend(register, small integer) variant of @llvm.ptrauth.blend to the AArch64PointerAuth pass, where most other PAuth-related code generation takes place.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64InstrInfo.td (+4-3)
• (modified) llvm/lib/Target/AArch64/AArch64PointerAuth.cpp (+37)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 44b0337fe7879..bca9dccf69adc 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -1570,6 +1570,9 @@ def PAUTH_PROLOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
 def PAUTH_EPILOGUE : Pseudo<(outs), (ins), []>, Sched<[]>;
 }
 
+def PAUTH_BLEND : Pseudo<(outs GPR64:$disc),
+                         (ins GPR64:$addr_disc, i32imm:$int_disc), []>, Sched<[]>;
+
 // These pointer authentication instructions require armv8.3a
 let Predicates = [HasPAuth] in {
 
@@ -9188,12 +9191,10 @@ let Predicates = [HasMOPS, HasMTE], Defs = [NZCV], Size = 12, mayLoad = 0, maySt
 //-----------------------------------------------------------------------------
 // v8.3 Pointer Authentication late patterns
 
-let Predicates = [HasPAuth] in {
 def : Pat<(int_ptrauth_blend GPR64:$Rd, imm64_0_65535:$imm),
-          (MOVKXi GPR64:$Rd, (trunc_imm imm64_0_65535:$imm), 48)>;
+          (PAUTH_BLEND GPR64:$Rd, (trunc_imm imm64_0_65535:$imm))>;
 def : Pat<(int_ptrauth_blend GPR64:$Rd, GPR64:$Rn),
           (BFMXri GPR64:$Rd, GPR64:$Rn, 16, 15)>;
-}
 
 //-----------------------------------------------------------------------------
 
diff --git a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
index 7576d2a899d1a..ec5db287571ee 100644
--- a/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
+++ b/llvm/lib/Target/AArch64/AArch64PointerAuth.cpp
@@ -46,6 +46,13 @@ class AArch64PointerAuth : public MachineFunctionPass {
   void authenticateLR(MachineFunction &MF,
                       MachineBasicBlock::iterator MBBI) const;
 
+  /// Stores blend(AddrDisc, IntDisc) to the Result register.
+  void emitBlend(MachineBasicBlock::iterator MBBI, Register Result,
+                 Register AddrDisc, unsigned IntDisc) const;
+
+  /// Expands PAUTH_BLEND pseudo instruction.
+  void expandPAuthBlend(MachineBasicBlock::iterator MBBI) const;
+
   bool checkAuthenticatedLR(MachineBasicBlock::iterator TI) const;
 };
 
@@ -295,6 +302,32 @@ bool AArch64PointerAuth::checkAuthenticatedLR(
   return true;
 }
 
+void AArch64PointerAuth::emitBlend(MachineBasicBlock::iterator MBBI,
+                                   Register Result, Register AddrDisc,
+                                   unsigned IntDisc) const {
+  MachineBasicBlock &MBB = *MBBI->getParent();
+  DebugLoc DL = MBBI->getDebugLoc();
+
+  if (Result != AddrDisc)
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), Result)
+        .addReg(AArch64::XZR)
+        .addReg(AddrDisc)
+        .addImm(0);
+
+  BuildMI(MBB, MBBI, DL, TII->get(AArch64::MOVKXi), Result)
+      .addReg(Result)
+      .addImm(IntDisc)
+      .addImm(48);
+}
+
+void AArch64PointerAuth::expandPAuthBlend(
+    MachineBasicBlock::iterator MBBI) const {
+  Register ResultReg = MBBI->getOperand(0).getReg();
+  Register AddrDisc = MBBI->getOperand(1).getReg();
+  unsigned IntDisc = MBBI->getOperand(2).getImm();
+  emitBlend(MBBI, ResultReg, AddrDisc, IntDisc);
+}
+
 bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
   const auto *MFnI = MF.getInfo<AArch64FunctionInfo>();
 
@@ -326,6 +359,7 @@ bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
         break;
       case AArch64::PAUTH_PROLOGUE:
       case AArch64::PAUTH_EPILOGUE:
+      case AArch64::PAUTH_BLEND:
         assert(!MI.isBundled());
         PAuthPseudoInstrs.push_back(MI.getIterator());
         break;
@@ -342,6 +376,9 @@ bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) {
       authenticateLR(MF, It);
       HasAuthenticationInstrs = true;
       break;
+    case AArch64::PAUTH_BLEND:
+      expandPAuthBlend(It);
+      break;
     default:
       llvm_unreachable("Unhandled opcode");
     }

[atrosinenko]

This patch does not change GlobalISel-specific implementation in AArch64InstructionSelector::selectIntrinsic. Looks like custom code is only needed for Intrinsic::ptrauth_blend because of a bug in TableGen backend for GlobalISel. I expect #74492 to make that custom handling of Intrinsic::ptrauth_blend unneeded. (This patch itself is unaffected by the bug because PAUTH_BLEND instruction defines GPR64:$disc and the name $disc is not used by the pattern).

[atrosinenko]

Ping.

The custom GISel-specific selector for ptrauth_blend was removed via #74492 and #75328.

[atrosinenko]

Ping.

[atrosinenko]

Ping.

Notes:

• this patch is arranged so that [AArch64][PAC] Implement code generation for @llvm.ptrauth.auth #72502 can reuse emitBlend method while expandPAuthBlend merely translates PAUTH_BLEND pseudo instruction (compared to [AArch64][PAC] Implement code generation for @llvm.ptrauth.auth #72502, added a dedicated method for uniformity)
• on enabling PAUTH_BLEND on cores without FEAT_PAUTH: in the discussion of the original patch on Phabricator, there were no strict opinions for either enabling or disabling it. Wouldn't enabling blend operation (and, afterwards, sign and auth using I-keys) useful for something like runtime libraries?

[atrosinenko]

Ping.

[asl]

The change looks reasonable to me. @ahmedbougacha @kbeyls ?

[kbeyls]

This also seems to make sense to me.
I have a minor comment about maybe adding a regression test case, but I'm happy with either way you'd make a decision on that

[kbeyls]

Looking at the one obvious regression test for the blend intrinsic, in "ptrauth-intrinsic-blend.ll", I seems to me that the case Result != AddrDisc isn't covered there, so the then part here isn't covered by a test case.
If my impression is correct, maybe it would make sense to add a test to that file to cover this case?

[atrosinenko]

Added a MIR-based test file, thank you.