[llvm] b3272f5 - [AArch64][PAC] Select MOVK for ptrauth.blend intrinsic.

[Ahmed Bougacha via llvm-commits]

Author: Ahmed Bougacha
Date: 2023-06-26T09:43:37-07:00
New Revision: b3272f5ddbfd2b4a3219df223392af7a09571cc1

URL: https://github.com/llvm/llvm-project/commit/b3272f5ddbfd2b4a3219df223392af7a09571cc1
DIFF: https://github.com/llvm/llvm-project/commit/b3272f5ddbfd2b4a3219df223392af7a09571cc1.diff

LOG: [AArch64][PAC] Select MOVK for ptrauth.blend intrinsic.

Blend combines two discriminator values used by other ptrauth ops.
On AArch64 here, it does that by replacing the high 16 bits of the
LHS with the low 16 bits of the RHS.

Usually the RHS is a constant, which lets us do this efficiently in
a single MOVK.  When the RHS isn't constant, we can do a BFI.

In a sense, this is implementing an ABI decision (how to lower the
software construct of "blend"), but if there are interesting variants to
consider, this could be made object-file-format-specific in some way.

Differential Revision: https://reviews.llvm.org/D132384

Added: 
    llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll

Modified: 
    llvm/lib/Target/AArch64/AArch64InstrFormats.td
    llvm/lib/Target/AArch64/AArch64InstrInfo.td
    llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64InstrFormats.td b/llvm/lib/Target/AArch64/AArch64InstrFormats.td
index 5d4baf67fa1a8..362bade748828 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrFormats.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrFormats.td
@@ -889,6 +889,10 @@ def timm32_0_65535 : Operand<i32>, TImmLeaf<i32, [{
 def timm64_0_65535 : Operand<i64>, TImmLeaf<i64, [{
   return ((uint64_t)Imm) < 65536;
 }]>;
+
+def imm64_0_65535 : Operand<i64>, ImmLeaf<i64, [{
+  return ((uint64_t)Imm) < 65536;
+}]>;
 }
 
 def imm0_255 : Operand<i32>, ImmLeaf<i32, [{

diff  --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 5da5006c28c9e..973b90a5ded87 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -8849,6 +8849,18 @@ let Predicates = [HasMOPS, HasMTE], Defs = [NZCV], Size = 12, mayLoad = 0, maySt
                                           [], "$Rd = $Rd_wb,$Rn = $Rn_wb">, Sched<[]>;
 }
 
+//-----------------------------------------------------------------------------
+// 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)>;
+def : Pat<(int_ptrauth_blend GPR64:$Rd, GPR64:$Rn),
+          (BFMXri GPR64:$Rd, GPR64:$Rn, 16, 15)>;
+}
+
+//-----------------------------------------------------------------------------
+
 // This gets lowered into an instruction sequence of 20 bytes
 let Defs = [X16, X17], mayStore = 1, isCodeGenOnly = 1, Size = 20 in
 def StoreSwiftAsyncContext

diff  --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index d7d2085289c6d..657053602d89f 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -5976,6 +5976,25 @@ bool AArch64InstructionSelector::selectIntrinsic(MachineInstr &I,
     I.eraseFromParent();
     return true;
   }
+  case Intrinsic::ptrauth_blend: {
+    MachineFunction &MF = *I.getParent()->getParent();
+    auto RHS = getIConstantVRegVal(I.getOperand(3).getReg(), MRI);
+    if (RHS && (RHS->getZExtValue() <= 0xffff)) {
+      I.setDesc(TII.get(AArch64::MOVKXi));
+      I.removeOperand(3);
+      I.removeOperand(1);
+      MachineInstrBuilder(MF, I)
+          .addImm(RHS->getZExtValue() & 0xffff)
+          .addImm(48)
+          .constrainAllUses(TII, TRI, RBI);
+    } else {
+      I.setDesc(TII.get(AArch64::BFMXri));
+      I.removeOperand(1);
+      MachineInstrBuilder(MF, I).addImm(16).addImm(15).constrainAllUses(
+          TII, TRI, RBI);
+    }
+    return true;
+  }
   case Intrinsic::frameaddress:
   case Intrinsic::returnaddress: {
     MachineFunction &MF = *I.getParent()->getParent();

diff  --git a/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll b/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll
new file mode 100644
index 0000000000000..e60b3dfeaf225
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=0 | FileCheck %s
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=1 -global-isel-abort=1 | FileCheck %s
+
+define i64 @test_blend(i64 %arg, i64 %arg1) {
+; CHECK-LABEL: test_blend:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    bfi x0, x1, #48, #16
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 %arg, i64 %arg1)
+  ret i64 %tmp
+}
+
+define i64 @test_blend_constant(i64 %arg) {
+; CHECK-LABEL: test_blend_constant:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    movk x0, #12345, lsl #48
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 %arg, i64 12345)
+  ret i64 %tmp
+}
+
+; Blend isn't commutative.
+define i64 @test_blend_constant_swapped(i64 %arg) {
+; CHECK-LABEL: test_blend_constant_swapped:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov w8, #12345
+; CHECK-NEXT:    bfi x8, x0, #48, #16
+; CHECK-NEXT:    mov x0, x8
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 12345, i64 %arg)
+  ret i64 %tmp
+}
+
+; Blends of constants wider than 16 bits truncate the constant.
+define i64 @test_blend_constant_wide(i64 %arg) {
+; CHECK-LABEL: test_blend_constant_wide:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov w8, #65536
+; CHECK-NEXT:    bfi x0, x8, #48, #16
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 %arg, i64 65536)
+  ret i64 %tmp
+}
+
+declare i64 @llvm.ptrauth.blend(i64, i64)