[llvm] [InstCombine] Combine ptrauth intrin. callee into same-key bundle. (PR #94707)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Jun 11 17:05:22 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Ahmed Bougacha (ahmedbougacha)
<details>
<summary>Changes</summary>
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
```
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.)
Generating a plain call to a raw unauthenticated pointer is generally undesirable, but if we ended up seeing a naked ptrauth.sign in the first place, we already have suspicious code. Unauthenticated calls are also easier to spot than naked signs, so let the indirect call shine.
Note that there is an arguably unsafe extension to this, where we don't bother checking that the key in bundle and intrinsic are the same (and also allow folding away an auth into a bundle.)
This can end up generating calls with a bundle that has an invalid key (which an informed frontend wouldn't have otherwise done), which can be problematic. The C that generates that is straightforward but arguably unreasonable. That wouldn't be an issue if we were to bite the bullet and make these fully AArch64-specific, allowing key knowledge to be embedded here.
I have a branch that does the "unsafe" combines in bce1917e28e47df58c5eb5e78437d24e2726e90b. I'm comfortable doing that in my toolchain but I imagine others may disagree here; open to people's thoughts.
---
Full diff: https://github.com/llvm/llvm-project/pull/94707.diff
3 Files Affected:
- (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+76)
- (modified) llvm/lib/Transforms/InstCombine/InstCombineInternal.h (+8)
- (added) llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll (+76)
``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 436cdbff75669..8c32f411320ab 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3665,6 +3665,78 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
return nullptr;
}
+Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) {
+ Value *Callee = Call.getCalledOperand();
+ 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;
+
+ // Isolate the ptrauth bundle from the others.
+ std::optional<OperandBundleUse> PtrAuthBundleOrNone;
+ SmallVector<OperandBundleDef, 2> NewBundles;
+ for (unsigned BI = 0, BE = Call.getNumOperandBundles(); BI != BE; ++BI) {
+ OperandBundleUse Bundle = Call.getOperandBundleAt(BI);
+ if (Bundle.getTagID() == LLVMContext::OB_ptrauth)
+ PtrAuthBundleOrNone = Bundle;
+ else
+ NewBundles.emplace_back(Bundle);
+ }
+
+ Value *NewCallee = nullptr;
+ 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 ||
+ II->getOperand(3) != PtrAuthBundleOrNone->Inputs[0] ||
+ 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);
+ break;
+ }
+
+ // call(ptrauth.sign(p)), ["ptrauth"()] -> call p
+ // assuming the call bundle and the sign operands match.
+ // Non-ptrauth indirect calls are undesirable, but so is ptrauth.sign.
+ case Intrinsic::ptrauth_sign: {
+ if (!PtrAuthBundleOrNone ||
+ II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0] ||
+ II->getOperand(2) != PtrAuthBundleOrNone->Inputs[1])
+ return nullptr;
+ 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;
+}
+
bool InstCombinerImpl::annotateAnyAllocSite(CallBase &Call,
const TargetLibraryInfo *TLI) {
// Note: We only handle cases which can't be driven from generic attributes
@@ -3812,6 +3884,10 @@ Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
if (IntrinsicInst *II = findInitTrampoline(Callee))
return transformCallThroughTrampoline(Call, *II);
+ // Combine calls involving pointer authentication intrinsics.
+ if (Instruction *NewCall = foldPtrAuthIntrinsicCallee(Call))
+ return NewCall;
+
if (isa<InlineAsm>(Callee) && !Call.doesNotThrow()) {
InlineAsm *IA = cast<InlineAsm>(Callee);
if (!IA->canThrow()) {
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 984f02bcccad7..f290f446b424e 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -282,6 +282,14 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
Instruction *transformCallThroughTrampoline(CallBase &Call,
IntrinsicInst &Tramp);
+ /// 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
+ /// 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.)
+ Instruction *foldPtrAuthIntrinsicCallee(CallBase &Call);
+
// Return (a, b) if (LHS, RHS) is known to be (a, b) or (b, a).
// Otherwise, return std::nullopt
// Currently it matches:
diff --git a/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll b/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll
new file mode 100644
index 0000000000000..900fd43f01602
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll
@@ -0,0 +1,76 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+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: 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)
+ %v2 = inttoptr i64 %v1 to i32()*
+ %v3 = call i32 %v2() [ "ptrauth"(i32 1, 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: 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)
+ %v5 = inttoptr i64 %v4 to i32()*
+ %v6 = call i32 %v5() [ "ptrauth"(i32 1, i64 %v3) ]
+ ret i32 %v6
+}
+
+define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
+; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
+; 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: 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)
+ %v5 = inttoptr i64 %v4 to i32()*
+ %v6 = call i32 %v5() [ "ptrauth"(i32 0, i64 1234) ]
+ ret i32 %v6
+}
+
+define i32 @test_ptrauth_call_sign(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_sign(
+; CHECK-NEXT: [[V3:%.*]] = call i32 [[P:%.*]]()
+; CHECK-NEXT: ret i32 [[V3]]
+;
+ %v0 = ptrtoint ptr %p to i64
+ %v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
+ %v2 = inttoptr i64 %v1 to i32()*
+ %v3 = call i32 %v2() [ "ptrauth"(i32 2, i64 5678) ]
+ ret i32 %v3
+}
+
+define i32 @test_ptrauth_call_sign_otherbundle(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_sign_otherbundle(
+; CHECK-NEXT: [[V3:%.*]] = call i32 [[P:%.*]]() [ "somebundle"(ptr null), "otherbundle"(i64 0) ]
+; CHECK-NEXT: ret i32 [[V3]]
+;
+ %v0 = ptrtoint ptr %p to i64
+ %v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
+ %v2 = inttoptr i64 %v1 to i32()*
+ %v3 = call i32 %v2() [ "somebundle"(ptr null), "ptrauth"(i32 2, i64 5678), "otherbundle"(i64 0) ]
+ ret i32 %v3
+}
+
+declare i64 @llvm.ptrauth.sign(i64, i32, i64)
+declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
+declare i64 @llvm.ptrauth.blend(i64, i64)
``````````
</details>
https://github.com/llvm/llvm-project/pull/94707
More information about the llvm-commits
mailing list