[AArch64] Lower ptrauth.sign of constant as PtrAuthGlobalAddress.

This lets us guarantee the emission of one of our hardened constant
signed pointer materialization sequences, rather than emitting the
sign intrinsic as a naked sign, and exposing the raw pointer.

Do that in the builder to help guarantee that through the backend.

At the IR level, there aren't great ways to enforce that the constant
operand to the intrinsic stays there, but we might be able to do that
by generalizing immarg to be applicable to callsites rather than just
Intrinsic definitions.  That's all theoretical either way, and there's
rarely a good reason to hoist a constant op from an intrinsic anyway.

This is pretty tolerant of weird discriminators, which should cover the
interesting cases.

Author: ahmedbougacha

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -8005,6 +8005,65 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     setValue(&I, DAG.getNode(ISD::AND, sdl, PtrVT, Ptr, Mask));
     return;
   }
+  case Intrinsic::ptrauth_sign: {
+    Value *Ptr = I.getOperand(0);
+    Value *Disc = I.getOperand(2);
+
+    if (auto *PtrToInt = dyn_cast<PtrToIntOperator>(Ptr))
+      Ptr = PtrToInt->getOperand(0);
+
+    APInt PtrOff(DAG.getDataLayout().getPointerSizeInBits(), 0);
+    Ptr = Ptr->stripAndAccumulateConstantOffsets(DAG.getDataLayout(), PtrOff,
+                                                 /*AllowNonInbounds=*/true);
+
+    // We found that the raw pointer is a pointer constant + constant offset:
+    // we can turn it into the (hardened) signed constant pointer
+    // materialization sequence, via PtrAuthGlobalAddress.
+    // Otherwise, let it be handled as a plain raw intrinsic.
+    if (!isa<GlobalValue>(Ptr)) {
+      visitTargetIntrinsic(I, Intrinsic);
+      return;
+    }
+
+    auto *DiscC = dyn_cast<ConstantInt>(Disc);
+    Value *AddrDisc = nullptr;
+
+    auto *BlendInt = dyn_cast<IntrinsicInst>(Disc);
+    if (BlendInt && BlendInt->getIntrinsicID() == Intrinsic::ptrauth_blend) {
+      DiscC = dyn_cast<ConstantInt>(BlendInt->getOperand(1));
+      AddrDisc = BlendInt->getOperand(0);
+    }
+    // If we can't isolate the constant discriminator, treat the whole thing
+    // as a variable address discriminator.
+    if (!DiscC)
+      AddrDisc = Disc;
+
+    // We're extracting the addr-disc from a blend intrinsic, which could be
+    // in another block: make sure it isn't, because it might not have gotten
+    // exported to this block via a vreg.
+    // The blends should be sunk into the same block by CGP already.
+    // Ideally, to guarantee this, we should make ptrauth.sign have both
+    // address-discriminator and integer discriminator, and do the blend itself.
+    if (auto *AddrDiscInst = dyn_cast_or_null<Instruction>(AddrDisc)) {
+      if (AddrDiscInst->getParent() != I.getParent()) {
+        visitTargetIntrinsic(I, Intrinsic);
+        return;
+      }
+    }
+
+    SDValue NullPtr = DAG.getIntPtrConstant(0, getCurSDLoc());
+    SDValue DiscVal = DiscC ? getValue(DiscC) : NullPtr;
+    SDValue AddrDiscVal = AddrDisc ? getValue(AddrDisc) : NullPtr;
+
+    setValue(&I, DAG.getNode(ISD::PtrAuthGlobalAddress, getCurSDLoc(),
+                             NullPtr.getValueType(),
+                             DAG.getGlobalAddress(
+                                 cast<GlobalValue>(Ptr), getCurSDLoc(),
+                                 NullPtr.getValueType(), PtrOff.getSExtValue()),
+                             getValue(I.getOperand(1)), AddrDiscVal, DiscVal));
+    return;
+  }
+
   case Intrinsic::threadlocal_address: {
     setValue(&I, getValue(I.getOperand(0)));
     return;

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-sign-constant.ll

@@ -0,0 +1,241 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple arm64e-apple-darwin | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+define i8* @test_intrinsic() {
+; CHECK-LABEL: test_intrinsic:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 56)
+  %tmp2 = inttoptr i64 %tmp1 to i8*
+  ret i8* %tmp2
+}
+
+define i8* @test_intrinsic_weak() {
+; CHECK-LABEL: test_intrinsic_weak:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x0, l_g_weak$auth_ptr$da$56@PAGE
+; CHECK-NEXT:    ldr x0, [x0, l_g_weak$auth_ptr$da$56@PAGEOFF]
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g_weak to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 56)
+  %tmp2 = inttoptr i64 %tmp1 to i8*
+  ret i8* %tmp2
+}
+
+; Non-external symbols don't need to be accessed through the GOT: always prefer
+; the dynamic materialization sequence, with adrp+add rather than a GOT load.
+
+define i8* @test_intrinsic_strong_def() {
+; CHECK-LABEL: test_intrinsic_strong_def:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g_strong_def@PAGE
+; CHECK-NEXT:    add x16, x16, _g_strong_def@PAGEOFF
+; CHECK-NEXT:    pacdza x16
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g_strong_def to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 0)
+  %tmp2 = inttoptr i64 %tmp1 to i8*
+  ret i8* %tmp2
+}
+
+define i8* @test_intrinsic_bkey() {
+; CHECK-LABEL: test_intrinsic_bkey:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    mov x17, #78 ; =0x4e
+; CHECK-NEXT:    pacib x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 1, i64 78)
+  %tmp2 = inttoptr i64 %tmp1 to i8*
+  ret i8* %tmp2
+}
+
+define i8* @test_intrinsic_constantexpr() {
+; CHECK-LABEL: test_intrinsic_constantexpr:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 ptrtoint (i32* @g to i64), i32 2, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+define i8* @test_intrinsic_constantexpr_offset() {
+; CHECK-LABEL: test_intrinsic_constantexpr_offset:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    add x16, x16, #16
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 ptrtoint (i8* getelementptr (i8, i8* bitcast (i32* @g to i8*), i64 16) to i64), i32 2, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+define i8* @test_intrinsic_constantexpr_offset_neg() {
+; CHECK-LABEL: test_intrinsic_constantexpr_offset_neg:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    mov x17, #1 ; =0x1
+; CHECK-NEXT:    movk x17, #32769, lsl #16
+; CHECK-NEXT:    add x16, x16, x17
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 ptrtoint (i8* getelementptr (i8, i8* bitcast (i32* @g to i8*), i64 add (i64 2147483648, i64 65537)) to i64), i32 2, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+define i8* @test_intrinsic_non_constant(i8* %arg0) {
+; CHECK-LABEL: test_intrinsic_non_constant:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov w8, #56 ; =0x38
+; CHECK-NEXT:    pacda x0, x8
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i8* %arg0 to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 56)
+  %tmp2 = inttoptr i64 %tmp1 to i8*
+  ret i8* %tmp2
+}
+
+define i8* @test_intrinsic_blend_addr_disc(i8* %arg0) {
+; CHECK-LABEL: test_intrinsic_blend_addr_disc:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    mov x17, x0
+; CHECK-NEXT:    movk x17, #23, lsl #48
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g to i64
+  %tmp1 = ptrtoint i8* %arg0 to i64
+  %tmp2 = call i64 @llvm.ptrauth.blend(i64 %tmp1, i64 23)
+  %tmp3 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 %tmp2)
+  %tmp4 = inttoptr i64 %tmp3 to i8*
+  ret i8* %tmp4
+}
+
+define i8* @test_intrinsic_addr_disc(i8* %arg0) {
+; CHECK-LABEL: test_intrinsic_addr_disc:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g@GOTPAGE
+; CHECK-NEXT:    ldr x16, [x16, _g@GOTPAGEOFF]
+; CHECK-NEXT:    pacda x16, x0
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = ptrtoint i32* @g to i64
+  %tmp1 = ptrtoint i8* %arg0 to i64
+  %tmp2 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 %tmp1)
+  %tmp3 = inttoptr i64 %tmp2 to i8*
+  ret i8* %tmp3
+}
+
+define i8* @test_intrinsic_blend_addr_disc_cross_bb(i8* %arg0, i8* %arg1, i1 %arg2) {
+; CHECK-LABEL: test_intrinsic_blend_addr_disc_cross_bb:
+; CHECK:       ; %bb.0: ; %common.ret
+; CHECK-NEXT:    mov x8, x0
+; CHECK-NEXT:    movk x8, #23, lsl #48
+; CHECK-NEXT:  Lloh0:
+; CHECK-NEXT:    adrp x9, _g2@GOTPAGE
+; CHECK-NEXT:  Lloh1:
+; CHECK-NEXT:    ldr x9, [x9, _g2@GOTPAGEOFF]
+; CHECK-NEXT:  Lloh2:
+; CHECK-NEXT:    adrp x10, _g@GOTPAGE
+; CHECK-NEXT:  Lloh3:
+; CHECK-NEXT:    ldr x10, [x10, _g@GOTPAGEOFF]
+; CHECK-NEXT:    tst w2, #0x1
+; CHECK-NEXT:    csel x0, x10, x9, ne
+; CHECK-NEXT:    pacda x0, x8
+; CHECK-NEXT:    ret
+; CHECK-NEXT:    .loh AdrpLdrGot Lloh2, Lloh3
+; CHECK-NEXT:    .loh AdrpLdrGot Lloh0, Lloh1
+  %tmp0 = ptrtoint i8* %arg0 to i64
+  %tmp1 = call i64 @llvm.ptrauth.blend(i64 %tmp0, i64 23)
+  br i1 %arg2, label %bb1, label %bb2
+
+bb1:
+  %tmp2 = ptrtoint i32* @g to i64
+  %tmp3 = call i64 @llvm.ptrauth.sign(i64 %tmp2, i32 2, i64 %tmp1)
+  %tmp4 = inttoptr i64 %tmp3 to i8*
+  ret i8* %tmp4
+
+bb2:
+  %tmp5 = ptrtoint i32* @g2 to i64
+  %tmp6 = call i64 @llvm.ptrauth.sign(i64 %tmp5, i32 2, i64 %tmp1)
+  %tmp7 = inttoptr i64 %tmp6 to i8*
+  ret i8* %tmp7
+}
+
+define i8* @test_intrinsic_function() {
+; CHECK-LABEL: test_intrinsic_function:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _test_intrinsic_function@PAGE
+; CHECK-NEXT:    add x16, x16, _test_intrinsic_function@PAGEOFF
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacib x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 ptrtoint (i8* ()* @test_intrinsic_function to i64), i32 1, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+define i8* @test_intrinsic_constant_int() {
+; CHECK-LABEL: test_intrinsic_constant_int:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov x0, #0 ; =0x0
+; CHECK-NEXT:    mov w8, #56 ; =0x38
+; CHECK-NEXT:    pacda x0, x8
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 0, i32 2, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+define i8* @test_intrinsic_constantexpr_offset_strong_def() {
+; CHECK-LABEL: test_intrinsic_constantexpr_offset_strong_def:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    adrp x16, _g_strong_def@PAGE
+; CHECK-NEXT:    add x16, x16, _g_strong_def@PAGEOFF
+; CHECK-NEXT:    add x16, x16, #2
+; CHECK-NEXT:    mov x17, #56 ; =0x38
+; CHECK-NEXT:    pacda x16, x17
+; CHECK-NEXT:    mov x0, x16
+; CHECK-NEXT:    ret
+  %tmp0 = call i64 @llvm.ptrauth.sign(i64 ptrtoint (i8* getelementptr (i8, i8* bitcast (i32* @g_strong_def to i8*), i64 2) to i64), i32 2, i64 56)
+  %tmp1 = inttoptr i64 %tmp0 to i8*
+  ret i8* %tmp1
+}
+
+declare i64 @llvm.ptrauth.sign(i64, i32, i64)
+declare i64 @llvm.ptrauth.blend(i64, i64)
+
+@g = external global i32
+@g2 = external global i32
+
+@g_weak = extern_weak global i32
+
+@g_strong_def = constant i32 42