[InstCombine] Combine ptrauth sign/resign + auth/resign intrinsics.

[ab]

Summary

Quoting the comment:

(sign|resign) + (auth|resign) can be folded by omitting the middle
sign+auth component if the key and discriminator match.

[apazos]

llvm/test/Transforms/InstCombine/ptrauth-intrinsics.ll

12

missing test cases for the combo resign - auth just like you have for the combo sign - auth

[ab]

• add test for resign+auth -> auth

llvm/test/Transforms/InstCombine/ptrauth-intrinsics.ll

12

Good catch, added!

[kristof.beyls]

LGTM, thanks!

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (match(II->getArgOperand(0), m_OneUse(m_FNeg(m_Value(X))))) {
       // sin(-x) --> -sin(x)
       Value *NewSin = Builder.CreateUnaryIntrinsic(Intrinsic::sin, X, II);
       Instruction *FNeg = UnaryOperator::CreateFNeg(NewSin);
       FNeg->copyFastMathFlags(II);
       return FNeg;
     }
     break;
   }
+  case Intrinsic::ptrauth_auth:
+  case Intrinsic::ptrauth_resign: {
+    // (sign|resign) + (auth|resign) can be folded by omitting the middle
+    // sign+auth component if the key and discriminator match.
+    bool NeedSign = II->getIntrinsicID() == Intrinsic::ptrauth_resign;
+    Value *Key = II->getArgOperand(1);
+    Value *Disc = II->getArgOperand(2);
+
+    // AuthKey will be the key we need to end up authenticating against in
+    // whatever we replace this sequence with.
+    Value *AuthKey = nullptr, *AuthDisc = nullptr, *BasePtr;
+    if (auto CI = dyn_cast<CallBase>(II->getArgOperand(0))) {
+      BasePtr = CI->getArgOperand(0);
+      if (CI->getIntrinsicID() == Intrinsic::ptrauth_sign) {
+        if (CI->getArgOperand(1) != Key || CI->getArgOperand(2) != Disc)
+          break;
+      } else if (CI->getIntrinsicID() == Intrinsic::ptrauth_resign) {
+        if (CI->getArgOperand(3) != Key || CI->getArgOperand(4) != Disc)
+          break;
+        AuthKey = CI->getArgOperand(1);
+        AuthDisc = CI->getArgOperand(2);
+      } else
+        break;
+    } else
+      break;
 
+    unsigned NewIntrin;
+    if (AuthKey && NeedSign) {
+      // resign(0,1) + resign(1,2) = resign(0, 2)
+      NewIntrin = Intrinsic::ptrauth_resign;
+    } else if (AuthKey) {
+      // resign(0,1) + auth(1) = auth(0)
+      NewIntrin = Intrinsic::ptrauth_auth;
+    } else if (NeedSign) {
+      // sign(0) + resign(0, 1) = sign(1)
+      NewIntrin = Intrinsic::ptrauth_sign;
+    } else {
+      // sign(0) + auth(0) = nop
+      replaceInstUsesWith(*II, BasePtr);
+      eraseInstFromFunction(*II);
+      return nullptr;
+    }
+
+    SmallVector<Value *, 4> CallArgs;
+    CallArgs.push_back(BasePtr);
+    if (AuthKey) {
+      CallArgs.push_back(AuthKey);
+      CallArgs.push_back(AuthDisc);
+    }
+
+    if (NeedSign) {
+      CallArgs.push_back(II->getArgOperand(3));
+      CallArgs.push_back(II->getArgOperand(4));
+    }
+
+    Function *NewFn = Intrinsic::getDeclaration(II->getModule(), NewIntrin);
+    return CallInst::Create(NewFn, CallArgs);
+  }
   case Intrinsic::arm_neon_vtbl1:
   case Intrinsic::aarch64_neon_tbl1:
     if (Value *V = simplifyNeonTbl1(*II, Builder))
       return replaceInstUsesWith(*II, V);
     break;
 
   case Intrinsic::arm_neon_vmulls:
   case Intrinsic::arm_neon_vmullu:

llvm/test/Transforms/InstCombine/ptrauth-intrinsics.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i64 @test_ptrauth_nop(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_nop(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    ret i64 [[TMP0]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 1, i64 1234)
+  %authed = call i64 @llvm.ptrauth.auth(i64 %signed, i32 1, i64 1234)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_nop_mismatch(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_nop_mismatch(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[SIGNED:%.*]] = call i64 @llvm.ptrauth.sign(i64 [[TMP0]], i32 1, i64 1234)
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.auth(i64 [[SIGNED]], i32 1, i64 10)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 1, i64 1234)
+  %authed = call i64 @llvm.ptrauth.auth(i64 %signed, i32 1, i64 10)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_nop_mismatch_keys(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_nop_mismatch_keys(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[SIGNED:%.*]] = call i64 @llvm.ptrauth.sign(i64 [[TMP0]], i32 0, i64 1234)
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.auth(i64 [[SIGNED]], i32 1, i64 1234)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 0, i64 1234)
+  %authed = call i64 @llvm.ptrauth.auth(i64 %signed, i32 1, i64 1234)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_sign_resign(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_sign_resign(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.sign(i64 [[TMP0]], i32 0, i64 42)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 1, i64 1234)
+  %authed = call i64 @llvm.ptrauth.resign(i64 %signed, i32 1, i64 1234, i32 0, i64 42)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_resign_resign(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_resign_resign(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.resign(i64 [[TMP0]], i32 1, i64 1234, i32 1, i64 3141)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 1, i64 1234, i32 0, i64 42)
+  %authed = call i64 @llvm.ptrauth.resign(i64 %signed, i32 0, i64 42, i32 1, i64 3141)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_resign_auth(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_resign_auth(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.auth(i64 [[TMP0]], i32 1, i64 1234)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 1, i64 1234, i32 0, i64 42)
+  %authed = call i64 @llvm.ptrauth.auth(i64 %signed, i32 0, i64 42)
+  ret i64 %authed
+}
+
+define i64 @test_ptrauth_resign_auth_mismatch(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_resign_auth_mismatch(
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint i8* [[P:%.*]] to i64
+; CHECK-NEXT:    [[SIGNED:%.*]] = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 1, i64 1234, i32 0, i64 10)
+; CHECK-NEXT:    [[AUTHED:%.*]] = call i64 @llvm.ptrauth.auth(i64 %signed, i32 0, i64 42)
+; CHECK-NEXT:    ret i64 [[AUTHED]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %signed = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 1, i64 1234, i32 0, i64 10)
+  %authed = call i64 @llvm.ptrauth.auth(i64 %signed, i32 0, i64 42)
+  ret i64 %authed
+}
+
+declare i64 @llvm.ptrauth.auth(i64, i32, i64)
+declare i64 @llvm.ptrauth.sign(i64, i32, i64)
+declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
 No newline at end of file