[InstCombine] Combine ptrauth intrin. callee into diff. key bundle.

Try to optimize a call to the result of a ptrauth intrinsic, potentially
into the ptrauth call bundle:
- call(ptrauth.resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
- call(ptrauth.sign(p)),   ["ptrauth"()] ->  call p
- call(ptrauth.auth(p))                  ->  call p, ["ptrauth"()]
as long as the key/discriminator are the same in sign and auth-bundle.

This relaxes the existing combine in one major way: it can generate calls
with ptrauth bundles with a different key from the existing call.
Without knowledge of the target keys, this may be unsound.

That also allows folding ptrauth.auth into the call bundle: we can't
generally do that if we don't know whether the auth key is valid for
calls.

Author: ahmedbougacha

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -3681,10 +3681,6 @@ Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) {
         II->getOperand(4) != PtrAuthBundleOrNone->Inputs[1])
       return nullptr;
 
-    // Don't change the key used in the call; we don't know what's valid.
-    if (II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0])
-      return nullptr;
-
     Value *NewBundleOps[] = {II->getOperand(1), II->getOperand(2)};
     NewBundles.emplace_back("ptrauth", NewBundleOps);
     NewCallee = II->getOperand(0);
@@ -3702,6 +3698,16 @@ Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) {
     NewCallee = II->getOperand(0);
     break;
   }
+
+  // call(ptrauth.auth(p)) ->  call p, ["ptrauth"()]
+  case Intrinsic::ptrauth_auth: {
+    if (PtrAuthBundleOrNone)
+      return nullptr;
+    Value *NewBundleOps[] = {II->getOperand(1), II->getOperand(2)};
+    NewBundles.emplace_back("ptrauth", NewBundleOps);
+    NewCallee = II->getOperand(0);
+    break;
+  }
   }
 
   if (!NewCallee)

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -286,8 +286,8 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   /// into the ptrauth call bundle:
   /// - call(ptrauth.resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
   /// - call(ptrauth.sign(p)),   ["ptrauth"()] ->  call p
-  /// as long as the key/discriminator are the same in sign and auth-bundle,
-  /// and we don't change the key in the bundle (to a potentially-invalid key.)
+  /// - call(ptrauth.auth(p))                  ->  call p, ["ptrauth"()]
+  /// as long as the key/discriminator are the same in sign and auth-bundle.
   Instruction *foldPtrAuthIntrinsicCallee(CallBase &Call);
 
   // Return (a, b) if (LHS, RHS) is known to be (a, b) or (b, a).

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

@@ -3,27 +3,27 @@
 
 define i32 @test_ptrauth_call_resign(ptr %p) {
 ; CHECK-LABEL: @test_ptrauth_call_resign(
-; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 1, i64 1234) ]
+; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 0, i64 1234) ]
 ; CHECK-NEXT:    ret i32 [[V3]]
 ;
   %v0 = ptrtoint ptr %p to i64
-  %v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 1, i64 1234, i32 1, i64 5678)
+  %v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 0, i64 1234, i32 2, i64 5678)
   %v2 = inttoptr i64 %v1 to i32()*
-  %v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
+  %v3 = call i32 %v2() [ "ptrauth"(i32 2, i64 5678) ]
   ret i32 %v3
 }
 
 define i32 @test_ptrauth_call_resign_blend(ptr %pp) {
 ; CHECK-LABEL: @test_ptrauth_call_resign_blend(
 ; CHECK-NEXT:    [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
-; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 1, i64 1234) ]
+; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 0, i64 1234) ]
 ; CHECK-NEXT:    ret i32 [[V6]]
 ;
   %v0 = load ptr, ptr %pp, align 8
   %v1 = ptrtoint ptr %pp to i64
   %v2 = ptrtoint ptr %v0 to i64
   %v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
-  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 1, i64 1234, i32 1, i64 %v3)
+  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 0, i64 1234, i32 1, i64 %v3)
   %v5 = inttoptr i64 %v4 to i32()*
   %v6 = call i32 %v5() [ "ptrauth"(i32 1, i64 %v3) ]
   ret i32 %v6
@@ -34,19 +34,31 @@ define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
 ; CHECK-NEXT:    [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
 ; CHECK-NEXT:    [[V1:%.*]] = ptrtoint ptr [[PP]] to i64
 ; CHECK-NEXT:    [[V3:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[V1]], i64 5678)
-; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 0, i64 [[V3]]) ]
+; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 1, i64 [[V3]]) ]
 ; CHECK-NEXT:    ret i32 [[V6]]
 ;
   %v0 = load ptr, ptr %pp, align 8
   %v1 = ptrtoint ptr %pp to i64
   %v2 = ptrtoint ptr %v0 to i64
   %v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
-  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 0, i64 %v3, i32 0, i64 1234)
+  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 1, i64 %v3, i32 0, i64 1234)
   %v5 = inttoptr i64 %v4 to i32()*
   %v6 = call i32 %v5() [ "ptrauth"(i32 0, i64 1234) ]
   ret i32 %v6
 }
 
+define i32 @test_ptrauth_call_auth(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_auth(
+; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 2, i64 5678) ]
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v0 = ptrtoint ptr %p to i64
+  %v1 = call i64 @llvm.ptrauth.auth(i64 %v0, i32 2, i64 5678)
+  %v2 = inttoptr i64 %v1 to i32()*
+  %v3 = call i32 %v2()
+  ret i32 %v3
+}
+
 define i32 @test_ptrauth_call_sign(ptr %p) {
 ; CHECK-LABEL: @test_ptrauth_call_sign(
 ; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]()