diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td index 6beacea31cb00..6018d662fe514 100644 --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td @@ -1581,6 +1581,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 { @@ -9317,12 +9320,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 7509afaeb5fef..d8a2853f68e02 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; }; @@ -359,6 +366,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(); @@ -390,6 +423,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; @@ -406,6 +440,9 @@ bool AArch64PointerAuth::runOnMachineFunction(MachineFunction &MF) { authenticateLR(MF, It); HasAuthenticationInstrs = true; break; + case AArch64::PAUTH_BLEND: + expandPAuthBlend(It); + break; default: llvm_unreachable("Unhandled opcode"); } diff --git a/llvm/test/CodeGen/AArch64/ptrauth-pseudo-instructions.mir b/llvm/test/CodeGen/AArch64/ptrauth-pseudo-instructions.mir new file mode 100644 index 0000000000000..d7fe1953deb47 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/ptrauth-pseudo-instructions.mir @@ -0,0 +1,27 @@ +# RUN: llc -mtriple=aarch64--- -run-pass=aarch64-ptrauth -verify-machineinstrs %s -o - | FileCheck %s + +# Test the corner cases that cannot be reliably tested using LLVM IR as input. + +--- | + define i64 @blend_untied(i64 %unused, i64 %ptr_arg) { + ret i64 0 + } +... +--- +# Check that the input register is copied to the output one, if not tied. + +name: blend_untied +tracksRegLiveness: true +body: | + bb.0: + liveins: $lr, $x0, $x1 + $x0 = PAUTH_BLEND $x1, 42 + RET undef $lr + +# CHECK: liveins: $lr, $x0, $x1 +# CHECK-NEXT: {{^ +$}} +# CHECK-NEXT: $x0 = ORRXrs $xzr, $x1, 0 +# CHECK-NEXT: $x0 = MOVKXi $x0, 42, 48 +# CHECK-NEXT: RET undef $lr + +...