[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707
Conversation
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.
ahmedbougacha
force-pushed
the
users/ahmedbougacha/ptrauth-instcombine-intrin-callee-safe
branch
from
903f149 to
e623f48
Compare
|
@llvm/pr-subscribers-backend-aarch64 @llvm/pr-subscribers-llvm-transforms Author: Ahmed Bougacha (ahmedbougacha) ChangesTry to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle: 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 bce1917. 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:
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)
|
There was a problem hiding this comment.
Changes are mostly OK as for me, but I think that several test cases are missing and also the code can be re-organized a little so we fail earlier for the majority of cases when such combining is not possible.
I'm OK with both adding tests as a part of this PR and submitting a follow-up patch with missing tests if adding them now is too time-consuming or there are other issues preventing that.
| @@ -0,0 +1,76 @@ | |||
| ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py | |||
There was a problem hiding this comment.
It's probably worth adding test cases for mismatching ptrauth parameters (key, discriminator, address discrimination) to ensure that combining is not performed when the parameters do not match. I like the approach for tests which is used in your other PR #94706. I'm OK with adding missing tests in a separate patch after merging this one.
| II->getOperand(4) != PtrAuthBundleOrNone->Inputs[1]) | ||
| return nullptr; | ||
|
|
||
| // Don't change the key used in the call; we don't know what's valid. |
There was a problem hiding this comment.
I suppose there should be a test for mismatching keys in resign intrinsic. I'm OK with adding missing tests in a separate patch after merging this one.
| // assuming the call bundle and the sign operands match. | ||
| case Intrinsic::ptrauth_resign: { | ||
| if (!PtrAuthBundleOrNone || | ||
| II->getOperand(3) != PtrAuthBundleOrNone->Inputs[0] || |
There was a problem hiding this comment.
Feel free to ignore: it's probably worth adding comments which say that operand N stands for X (e.g. operand 3 of intrinsic is a key to be resigned with). You already split the if into multiple parts instead of just using a giant condition in one if, so it looks like that there is not problem to split it even further and use multiple if statements with simple conditions and corresponding comments.
| IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0)); | ||
| if (!II) | ||
| return nullptr; | ||
|
|
There was a problem hiding this comment.
I'd probably add smth like this here to return early:
Intrinsic::ID IIID = II->getIntrinsicID();
if (IIID != Intrinsic::ptrauth_resign && IIID != Intrinsic::ptrauth_sign)
return nullptr;
In switch below, you can use smth like llvm_unreachable for default label.
| else | ||
| NewBundles.emplace_back(Bundle); | ||
| } | ||
|
|
There was a problem hiding this comment.
I'd probably add smth like this here to return early:
if (!PtrAuthBundleOrNone)
return nullptr;
With this, you can delete corresponding checks below.
| @@ -3665,6 +3665,78 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) { | |||
| return nullptr; | |||
| } | |||
|
|
|||
| Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) { | |||
| Value *Callee = Call.getCalledOperand(); | |||
There was a problem hiding this comment.
nit
| Value *Callee = Call.getCalledOperand(); | |
| const Value *Callee = Call.getCalledOperand(); |
| @@ -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); | |||
There was a problem hiding this comment.
nit
| auto *IPC = dyn_cast<IntToPtrInst>(Callee); | |
| const auto *IPC = dyn_cast<IntToPtrInst>(Callee); |
| if (!IPC || !IPC->isNoopCast(DL)) | ||
| return nullptr; | ||
|
|
||
| IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0)); |
There was a problem hiding this comment.
nit
| IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0)); | |
| const IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0)); |
|
|
||
| // 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. |
There was a problem hiding this comment.
Yes, both of these are bad options - not sure though what is worse, but I'm happy with having unauthenticated indirect calls here if everyone else is happy with this as well.
It's actually interesting, can we have such an IR in "wild nature" when: 1) generating IR from C/C++ code and not manually writing it; 2) not using ptrauth intrinsics in code explicitly? If yes, could you please provide an example of such code?
Try to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle:
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 bce1917. I'm comfortable doing that in my toolchain but I imagine others may disagree here; open to people's thoughts.