[InstCombine] Combine ptrauth intrinsics into call operand bundles.

Since the "ptrauth" call operand bundles really describe a folded "llvm.ptrauth.auth" operation, there are optimizations we can make to elide redundant sign/auth pairs. Note that this needs to be careful about preserving the integrity of the intermediate unauthenticated values, as that's the sole purpose of having operand bundles in the first place.
swiftlang · ahmedbougacha · Oct 29, 2019 · Jun 26, 2019 · Jun 26, 2019 · Oct 21, 2019
commit a90507b2d353b243e3ee755f981ff4544b3908be
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -4214,6 +4214,81 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
   return nullptr;
 }
 
+Instruction *InstCombiner::tryCombinePtrAuthCall(CallBase &Call) {
+  Value *Callee = Call.getCalledValue();
+  auto *IPC = dyn_cast<IntToPtrInst>(Callee);
+  if (!IPC || !IPC->isNoopCast(DL))
+    return nullptr;
+
+  IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0));
+  if (!II)
+    return nullptr;
+
+  auto PtrAuthBundleOrNone = Call.getOperandBundle(LLVMContext::OB_ptrauth);
+  assert(Call.getNumOperandBundles() <= 1 &&
+         "unimplemented support for ptrauth and other bundle");
+
+  Value *NewCallee = nullptr;
+  SmallVector<OperandBundleDef, 1> NewBundles;
+  switch (II->getIntrinsicID()) {
+  default:
+    return nullptr;
+
+  // call(ptrauth_resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
+  // assuming the call bundle and the sign operands match.
+  case Intrinsic::ptrauth_resign: {
+    if (!PtrAuthBundleOrNone)
+      return nullptr;
+    auto PtrAuthBundle = *PtrAuthBundleOrNone;
+    if (II->getOperand(3) != PtrAuthBundle.Inputs[0] ||
+        II->getOperand(4) != PtrAuthBundle.Inputs[1])
+      return nullptr;
+
+    Value *NewBundleOps[] = {II->getOperand(1), II->getOperand(2)};
+    NewBundles.emplace_back("ptrauth", NewBundleOps);
+    NewCallee = II->getOperand(0);
+    break;
+  }
+
+  // call(ptrauth_sign(p)), ["ptrauth"()] ->  call p
+  // assuming the call bundle and the sign operands match.
+  case Intrinsic::ptrauth_sign: {
+    if (!PtrAuthBundleOrNone)
+      return nullptr;
+    auto PtrAuthBundle = *PtrAuthBundleOrNone;
+    if (II->getOperand(1) != PtrAuthBundle.Inputs[0] ||
+        II->getOperand(2) != PtrAuthBundle.Inputs[1])
+      return nullptr;
+    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)
+    return nullptr;
+
+  NewCallee = Builder.CreateBitOrPointerCast(NewCallee, Callee->getType());
+  CallBase *NewCall = nullptr;
+  if (auto *CI = dyn_cast<CallInst>(&Call)) {
+    NewCall = CallInst::Create(CI, NewBundles);
+  } else {
+    auto *IKI = cast<InvokeInst>(&Call);
+    NewCall = InvokeInst::Create(IKI, NewBundles);
+  }
+  NewCall->setCalledOperand(NewCallee);
+  return NewCall;
+}
+
 static void annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI) {
   unsigned NumArgs = Call.getNumArgOperands();
   ConstantInt *Op0C = dyn_cast<ConstantInt>(Call.getOperand(0));
@@ -4356,6 +4431,11 @@ Instruction *InstCombiner::visitCallBase(CallBase &Call) {
   if (IntrinsicInst *II = findInitTrampoline(Callee))
     return transformCallThroughTrampoline(Call, *II);
 
+  // Combine calls involving pointer authentication
+  if (Instruction *NewCall = tryCombinePtrAuthCall(Call))
+    return NewCall;
+
+
   PointerType *PTy = cast<PointerType>(Callee->getType());
   FunctionType *FTy = cast<FunctionType>(PTy->getElementType());
   if (FTy->isVarArg()) {

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -505,6 +505,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner
   bool transformConstExprCastCall(CallBase &Call);
   Instruction *transformCallThroughTrampoline(CallBase &Call,
                                               IntrinsicInst &Tramp);
+  Instruction *tryCombinePtrAuthCall(CallBase &Call);
 
   Value *simplifyMaskedLoad(IntrinsicInst &II);
   Instruction *simplifyMaskedStore(IntrinsicInst &II);

diff --git a/llvm/test/Transforms/InstCombine/ptrauth-call.ll b/llvm/test/Transforms/InstCombine/ptrauth-call.ll
@@ -0,0 +1,82 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i32 @test_ptrauth_call_resign(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_call_resign(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[P:%.*]] to i32 ()*
+; CHECK-NEXT:    [[TMP3:%.*]] = call i32 [[TMP1]]() [ "ptrauth"(i32 0, i64 1234) ]
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %tmp1 = call i64 @llvm.ptrauth.resign.i64(i64 %tmp0, i32 0, i64 1234, i32 2, i64 5678)
+  %tmp2 = inttoptr i64 %tmp1 to i32()*
+  %tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
+  ret i32 %tmp3
+}
+
+define i32 @test_ptrauth_call_resign_blend(i8** %pp) {
+; CHECK-LABEL: @test_ptrauth_call_resign_blend(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8** [[PP:%.*]] to i32 ()**
+; CHECK-NEXT:    [[TMP012:%.*]] = load i32 ()*, i32 ()** [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 [[TMP012]]() [ "ptrauth"(i32 0, i64 1234) ]
+; CHECK-NEXT:    ret i32 [[TMP6]]
+;
+  %tmp0 = load i8*, i8** %pp, align 8
+  %tmp1 = ptrtoint i8** %pp to i64
+  %tmp2 = ptrtoint i8* %tmp0 to i64
+  %tmp3 = call i64 @llvm.ptrauth.blend.i64(i64 %tmp1, i64 5678)
+  %tmp4 = call i64 @llvm.ptrauth.resign.i64(i64 %tmp2, i32 0, i64 1234, i32 1, i64 %tmp3)
+  %tmp5 = inttoptr i64 %tmp4 to i32()*
+  %tmp6 = call i32 %tmp5() [ "ptrauth"(i32 1, i64 %tmp3) ]
+  ret i32 %tmp6
+}
+
+define i32 @test_ptrauth_call_resign_blend_2(i8** %pp) {
+; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8** [[PP:%.*]] to i32 ()**
+; CHECK-NEXT:    [[TMP012:%.*]] = load i32 ()*, i32 ()** [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i8** [[PP]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = call i64 @llvm.ptrauth.blend.i64(i64 [[TMP1]], i64 5678)
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 [[TMP012]]() [ "ptrauth"(i32 1, i64 [[TMP3]]) ]
+; CHECK-NEXT:    ret i32 [[TMP6]]
+;
+  %tmp0 = load i8*, i8** %pp, align 8
+  %tmp1 = ptrtoint i8** %pp to i64
+  %tmp2 = ptrtoint i8* %tmp0 to i64
+  %tmp3 = call i64 @llvm.ptrauth.blend.i64(i64 %tmp1, i64 5678)
+  %tmp4 = call i64 @llvm.ptrauth.resign.i64(i64 %tmp2, i32 1, i64 %tmp3, i32 0, i64 1234)
+  %tmp5 = inttoptr i64 %tmp4 to i32()*
+  %tmp6 = call i32 %tmp5() [ "ptrauth"(i32 0, i64 1234) ]
+  ret i32 %tmp6
+}
+
+define i32 @test_ptrauth_call_auth(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_call_auth(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[P:%.*]] to i32 ()*
+; CHECK-NEXT:    [[TMP3:%.*]] = call i32 [[TMP1]]() [ "ptrauth"(i32 2, i64 5678) ]
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %tmp1 = call i64 @llvm.ptrauth.auth.i64(i64 %tmp0, i32 2, i64 5678)
+  %tmp2 = inttoptr i64 %tmp1 to i32()*
+  %tmp3 = call i32 %tmp2()
+  ret i32 %tmp3
+}
+
+define i32 @test_ptrauth_call_sign(i8* %p) {
+; CHECK-LABEL: @test_ptrauth_call_sign(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[P:%.*]] to i32 ()*
+; CHECK-NEXT:    [[TMP3:%.*]] = call i32 [[TMP1]]()
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+  %tmp0 = ptrtoint i8* %p to i64
+  %tmp1 = call i64 @llvm.ptrauth.sign.i64(i64 %tmp0, i32 2, i64 5678)
+  %tmp2 = inttoptr i64 %tmp1 to i32()*
+  %tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
+  ret i32 %tmp3
+}
+
+declare i64 @llvm.ptrauth.auth.i64(i64, i32, i64)
+declare i64 @llvm.ptrauth.sign.i64(i64, i32, i64)
+declare i64 @llvm.ptrauth.resign.i64(i64, i32, i64, i32, i64)
+declare i64 @llvm.ptrauth.blend.i64(i64, i64)