[InstCombine] Combine ptrauth constant callee into bundle. #94706
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Try to optimize a call to a ptrauth constant, into its ptrauth bundle: call(ptrauth(f)), ["ptrauth"()] -> call f as long as the key/discriminator are the same in constant and bundle.
ahmedbougacha
force-pushed
the
users/ahmedbougacha/ptrauth-instcombine-constant-callee
branch
from
2b18790 to
de79c48
Compare
Member
|
@llvm/pr-subscribers-backend-aarch64 @llvm/pr-subscribers-llvm-transforms Author: Ahmed Bougacha (ahmedbougacha) ChangesTry to optimize a call to a ptrauth constant, into its ptrauth bundle: as long as the key/discriminator are the same in constant and bundle. Full diff: https://github.com/llvm/llvm-project/pull/94706.diff 3 Files Affected:
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 436cdbff75669..64f3038d94f94 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3665,6 +3665,34 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
return nullptr;
}
+Instruction *InstCombinerImpl::foldPtrAuthConstantCallee(CallBase &Call) {
+ auto *CPA = dyn_cast<ConstantPtrAuth>(Call.getCalledOperand());
+ if (!CPA)
+ return nullptr;
+
+ auto *CalleeF = dyn_cast<Function>(CPA->getPointer()->stripPointerCasts());
+ // If the ptrauth constant isn't based on a function pointer, bail out.
+ if (!CalleeF)
+ return nullptr;
+
+ // Inspect the call ptrauth bundle to check it matches the ptrauth constant.
+ auto PAB = Call.getOperandBundle(LLVMContext::OB_ptrauth);
+ if (!PAB)
+ return nullptr;
+
+ auto *Key = cast<ConstantInt>(PAB->Inputs[0]);
+ Value *Discriminator = PAB->Inputs[1];
+
+ // If the bundle doesn't match, this is probably going to fail to auth.
+ if (!CPA->isKnownCompatibleWith(Key, Discriminator, DL))
+ return nullptr;
+
+ // If the bundle matches the constant, proceed in making this a direct call.
+ auto *NewCall = CallBase::removeOperandBundle(&Call, LLVMContext::OB_ptrauth);
+ NewCall->setCalledOperand(CalleeF);
+ 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 +3840,10 @@ Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
if (IntrinsicInst *II = findInitTrampoline(Callee))
return transformCallThroughTrampoline(Call, *II);
+ // Combine calls to ptrauth constants.
+ if (Instruction *NewCall = foldPtrAuthConstantCallee(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..9268cbe594d90 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -282,6 +282,11 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
Instruction *transformCallThroughTrampoline(CallBase &Call,
IntrinsicInst &Tramp);
+ /// Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
+ /// call(ptrauth(f)), ["ptrauth"()] -> call f
+ /// as long as the key/discriminator are the same in constant and bundle.
+ Instruction *foldPtrAuthConstantCallee(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-call.ll b/llvm/test/Transforms/InstCombine/ptrauth-call.ll
new file mode 100644
index 0000000000000..b4363b528d4e2
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/ptrauth-call.ll
@@ -0,0 +1,89 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+
+declare i64 @f(i32)
+declare ptr @f2(i32)
+
+define i32 @test_ptrauth_call(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call(
+; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v0 = call i32 ptrauth(ptr @f, i32 0)(i32 %a0) [ "ptrauth"(i32 0, i64 0) ]
+ ret i32 %v0
+}
+
+define i32 @test_ptrauth_call_disc(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_disc(
+; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 1, i64 5678) ]
+ ret i32 %v0
+}
+
+@f_addr_disc.ref = constant ptr ptrauth(ptr @f, i32 1, i64 0, ptr @f_addr_disc.ref)
+
+define i32 @test_ptrauth_call_addr_disc(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_addr_disc(
+; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 0, ptr @f_addr_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 ptrtoint (ptr @f_addr_disc.ref to i64)) ]
+ ret i32 %v0
+}
+
+@f_both_disc.ref = constant ptr ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)
+
+define i32 @test_ptrauth_call_blend(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_blend(
+; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 1234)
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 %v) ]
+ ret i32 %v0
+}
+
+define i64 @test_ptrauth_call_cast(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_cast(
+; CHECK-NEXT: [[V0:%.*]] = call ptr @f2(i32 [[A0:%.*]])
+; CHECK-NEXT: [[TMP1:%.*]] = ptrtoint ptr [[V0]] to i64
+; CHECK-NEXT: ret i64 [[TMP1]]
+;
+ %v0 = call i64 ptrauth(ptr @f2, i32 0)(i32 %a0) [ "ptrauth"(i32 0, i64 0) ]
+ ret i64 %v0
+}
+
+define i32 @test_ptrauth_call_mismatch_key(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_mismatch_key(
+; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 5678)(i32 [[A0:%.*]]) [ "ptrauth"(i32 0, i64 5678) ]
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 0, i64 5678) ]
+ ret i32 %v0
+}
+
+define i32 @test_ptrauth_call_mismatch_disc(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_mismatch_disc(
+; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 5678)(i32 [[A0:%.*]]) [ "ptrauth"(i32 1, i64 0) ]
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 1, i64 0) ]
+ ret i32 %v0
+}
+
+define i32 @test_ptrauth_call_mismatch_blend(i32 %a0) {
+; CHECK-LABEL: @test_ptrauth_call_mismatch_blend(
+; CHECK-NEXT: [[V:%.*]] = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 0)
+; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 [[A0:%.*]]) [ "ptrauth"(i32 1, i64 [[V]]) ]
+; CHECK-NEXT: ret i32 [[V0]]
+;
+ %v = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 0)
+ %v0 = call i32 ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 %v) ]
+ ret i32 %v0
+}
+
+declare i64 @llvm.ptrauth.blend(i64, i64)
|
nikic
reviewed
Jun 12, 2024
kovdan01
approved these changes
Jun 17, 2024
llvmbot
added
the
llvm:instcombine
github-project-automation
bot
moved this from In Progress
to Done
in Pointer Authentication Tasks
ahmedbougacha
deleted the
users/ahmedbougacha/ptrauth-instcombine-constant-callee
branch
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/162/builds/26905 Here is the relevant piece of the build log for the reference |
ahmedbougacha
added a commit
to swiftlang/llvm-project
that referenced
this pull request
Jul 17, 2025
Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
call(ptrauth(f)), ["ptrauth"()] -> call f
as long as the key/discriminator are the same in constant and bundle.
Reapplies 42d2ae1 after it was dropped wholesale in d85524f.
Also renames the old (unrelated) test to match the upstream rename.
rdar://156047963
ahmedbougacha
added a commit
to swiftlang/llvm-project
that referenced
this pull request
Jul 17, 2025
Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
call(ptrauth(f)), ["ptrauth"()] -> call f
as long as the key/discriminator are the same in constant and bundle.
Reapplies 42d2ae1 after it was dropped wholesale in d85524f.
Also renames the old (unrelated) test to match the upstream rename.
rdar://156047963
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
as long as the key/discriminator are the same in constant and bundle.