Skip to content

Commit 9224771

Browse files
ahmedbougachaahmedbougacha
committed
[InstCombine] Combine ptrauth constant callee into bundle. (llvm#94706)
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
1 parent ee9bfc7 commit 9224771

File tree

4 files changed

+192
-58
lines changed

4 files changed

+192
-58
lines changed

‎llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4125,6 +4125,34 @@ Instruction *InstCombinerImpl::tryCombinePtrAuthCall(CallBase &Call) {
41254125
return NewCall;
41264126
}
41274127

4128+
Instruction *InstCombinerImpl::foldPtrAuthConstantCallee(CallBase &Call) {
4129+
auto *CPA = dyn_cast<ConstantPtrAuth>(Call.getCalledOperand());
4130+
if (!CPA)
4131+
return nullptr;
4132+
4133+
auto *CalleeF = dyn_cast<Function>(CPA->getPointer());
4134+
// If the ptrauth constant isn't based on a function pointer, bail out.
4135+
if (!CalleeF)
4136+
return nullptr;
4137+
4138+
// Inspect the call ptrauth bundle to check it matches the ptrauth constant.
4139+
auto PAB = Call.getOperandBundle(LLVMContext::OB_ptrauth);
4140+
if (!PAB)
4141+
return nullptr;
4142+
4143+
auto *Key = cast<ConstantInt>(PAB->Inputs[0]);
4144+
Value *Discriminator = PAB->Inputs[1];
4145+
4146+
// If the bundle doesn't match, this is probably going to fail to auth.
4147+
if (!CPA->isKnownCompatibleWith(Key, Discriminator, DL))
4148+
return nullptr;
4149+
4150+
// If the bundle matches the constant, proceed in making this a direct call.
4151+
auto *NewCall = CallBase::removeOperandBundle(&Call, LLVMContext::OB_ptrauth);
4152+
NewCall->setCalledOperand(CalleeF);
4153+
return NewCall;
4154+
}
4155+
41284156
bool InstCombinerImpl::annotateAnyAllocSite(CallBase &Call,
41294157
const TargetLibraryInfo *TLI) {
41304158
// Note: We only handle cases which can't be driven from generic attributes
@@ -4289,6 +4317,10 @@ Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
42894317
if (Instruction *NewCall = tryCombinePtrAuthCall(Call))
42904318
return NewCall;
42914319

4320+
// Combine calls to ptrauth constants.
4321+
if (Instruction *NewCall = foldPtrAuthConstantCallee(Call))
4322+
return NewCall;
4323+
42924324
if (isa<InlineAsm>(Callee) && !Call.doesNotThrow()) {
42934325
InlineAsm *IA = cast<InlineAsm>(Callee);
42944326
if (!IA->canThrow()) {

‎llvm/lib/Transforms/InstCombine/InstCombineInternal.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,11 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
284284
IntrinsicInst &Tramp);
285285
Instruction *tryCombinePtrAuthCall(CallBase &Call);
286286

287+
/// Try to optimize a call to a ptrauth constant, into its ptrauth bundle:
288+
/// call(ptrauth(f)), ["ptrauth"()] -> call f
289+
/// as long as the key/discriminator are the same in constant and bundle.
290+
Instruction *foldPtrAuthConstantCallee(CallBase &Call);
291+
287292
// Return (a, b) if (LHS, RHS) is known to be (a, b) or (b, a).
288293
// Otherwise, return std::nullopt
289294
// Currently it matches:

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

Lines changed: 78 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,77 +1,97 @@
11
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
22
; RUN: opt < %s -passes=instcombine -S | FileCheck %s
33

4-
define i32 @test_ptrauth_call_resign(ptr %p) {
5-
; CHECK-LABEL: @test_ptrauth_call_resign(
6-
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 0, i64 1234) ]
7-
; CHECK-NEXT: ret i32 [[TMP3]]
4+
declare i64 @f(i32)
5+
declare ptr @f2(i32)
6+
7+
define i32 @test_ptrauth_call(i32 %a0) {
8+
; CHECK-LABEL: @test_ptrauth_call(
9+
; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
10+
; CHECK-NEXT: ret i32 [[V0]]
11+
;
12+
%v0 = call i32 ptrauth(ptr @f, i32 0)(i32 %a0) [ "ptrauth"(i32 0, i64 0) ]
13+
ret i32 %v0
14+
}
15+
16+
define i32 @test_ptrauth_call_disc(i32 %a0) {
17+
; CHECK-LABEL: @test_ptrauth_call_disc(
18+
; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
19+
; CHECK-NEXT: ret i32 [[V0]]
20+
;
21+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 1, i64 5678) ]
22+
ret i32 %v0
23+
}
24+
25+
@f_addr_disc.ref = constant ptr ptrauth(ptr @f, i32 1, i64 0, ptr @f_addr_disc.ref)
26+
27+
define i32 @test_ptrauth_call_addr_disc(i32 %a0) {
28+
; CHECK-LABEL: @test_ptrauth_call_addr_disc(
29+
; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
30+
; CHECK-NEXT: ret i32 [[V0]]
31+
;
32+
%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)) ]
33+
ret i32 %v0
34+
}
35+
36+
@f_both_disc.ref = constant ptr ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)
37+
38+
define i32 @test_ptrauth_call_blend(i32 %a0) {
39+
; CHECK-LABEL: @test_ptrauth_call_blend(
40+
; CHECK-NEXT: [[V0:%.*]] = call i32 @f(i32 [[A0:%.*]])
41+
; CHECK-NEXT: ret i32 [[V0]]
42+
;
43+
%v = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 1234)
44+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 %v) ]
45+
ret i32 %v0
46+
}
47+
48+
define i64 @test_ptrauth_call_cast(i32 %a0) {
49+
; CHECK-LABEL: @test_ptrauth_call_cast(
50+
; CHECK-NEXT: [[V0:%.*]] = call i64 @f2(i32 [[A0:%.*]])
51+
; CHECK-NEXT: ret i64 [[V0]]
852
;
9-
%tmp0 = ptrtoint ptr %p to i64
10-
%tmp1 = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 0, i64 1234, i32 2, i64 5678)
11-
%tmp2 = inttoptr i64 %tmp1 to ptr
12-
%tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
13-
ret i32 %tmp3
53+
%v0 = call i64 ptrauth(ptr @f2, i32 0)(i32 %a0) [ "ptrauth"(i32 0, i64 0) ]
54+
ret i64 %v0
1455
}
1556

16-
define i32 @test_ptrauth_call_resign_blend(ptr %pp) {
17-
; CHECK-LABEL: @test_ptrauth_call_resign_blend(
18-
; CHECK-NEXT: [[TMP01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
19-
; CHECK-NEXT: [[TMP6:%.*]] = call i32 [[TMP01]]() [ "ptrauth"(i32 0, i64 1234) ]
20-
; CHECK-NEXT: ret i32 [[TMP6]]
57+
define i32 @test_ptrauth_call_mismatch_key(i32 %a0) {
58+
; CHECK-LABEL: @test_ptrauth_call_mismatch_key(
59+
; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 5678)(i32 [[A0:%.*]]) [ "ptrauth"(i32 0, i64 5678) ]
60+
; CHECK-NEXT: ret i32 [[V0]]
2161
;
22-
%tmp0 = load ptr, ptr %pp, align 8
23-
%tmp1 = ptrtoint ptr %pp to i64
24-
%tmp2 = ptrtoint ptr %tmp0 to i64
25-
%tmp3 = call i64 @llvm.ptrauth.blend(i64 %tmp1, i64 5678)
26-
%tmp4 = call i64 @llvm.ptrauth.resign(i64 %tmp2, i32 0, i64 1234, i32 1, i64 %tmp3)
27-
%tmp5 = inttoptr i64 %tmp4 to ptr
28-
%tmp6 = call i32 %tmp5() [ "ptrauth"(i32 1, i64 %tmp3) ]
29-
ret i32 %tmp6
62+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 0, i64 5678) ]
63+
ret i32 %v0
3064
}
3165

32-
define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
33-
; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
34-
; CHECK-NEXT: [[TMP01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
35-
; CHECK-NEXT: [[TMP1:%.*]] = ptrtoint ptr [[PP]] to i64
36-
; CHECK-NEXT: [[TMP3:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[TMP1]], i64 5678)
37-
; CHECK-NEXT: [[TMP6:%.*]] = call i32 [[TMP01]]() [ "ptrauth"(i32 1, i64 [[TMP3]]) ]
38-
; CHECK-NEXT: ret i32 [[TMP6]]
66+
define i32 @test_ptrauth_call_mismatch_disc(i32 %a0) {
67+
; CHECK-LABEL: @test_ptrauth_call_mismatch_disc(
68+
; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 5678)(i32 [[A0:%.*]]) [ "ptrauth"(i32 1, i64 0) ]
69+
; CHECK-NEXT: ret i32 [[V0]]
3970
;
40-
%tmp0 = load ptr, ptr %pp, align 8
41-
%tmp1 = ptrtoint ptr %pp to i64
42-
%tmp2 = ptrtoint ptr %tmp0 to i64
43-
%tmp3 = call i64 @llvm.ptrauth.blend(i64 %tmp1, i64 5678)
44-
%tmp4 = call i64 @llvm.ptrauth.resign(i64 %tmp2, i32 1, i64 %tmp3, i32 0, i64 1234)
45-
%tmp5 = inttoptr i64 %tmp4 to ptr
46-
%tmp6 = call i32 %tmp5() [ "ptrauth"(i32 0, i64 1234) ]
47-
ret i32 %tmp6
71+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 5678)(i32 %a0) [ "ptrauth"(i32 1, i64 0) ]
72+
ret i32 %v0
4873
}
4974

50-
define i32 @test_ptrauth_call_auth(ptr %p) {
51-
; CHECK-LABEL: @test_ptrauth_call_auth(
52-
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 2, i64 5678) ]
53-
; CHECK-NEXT: ret i32 [[TMP3]]
75+
define i32 @test_ptrauth_call_mismatch_blend(i32 %a0) {
76+
; CHECK-LABEL: @test_ptrauth_call_mismatch_blend(
77+
; CHECK-NEXT: [[V:%.*]] = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 0)
78+
; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 [[A0:%.*]]) [ "ptrauth"(i32 1, i64 [[V]]) ]
79+
; CHECK-NEXT: ret i32 [[V0]]
5480
;
55-
%tmp0 = ptrtoint ptr %p to i64
56-
%tmp1 = call i64 @llvm.ptrauth.auth(i64 %tmp0, i32 2, i64 5678)
57-
%tmp2 = inttoptr i64 %tmp1 to ptr
58-
%tmp3 = call i32 %tmp2()
59-
ret i32 %tmp3
81+
%v = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_both_disc.ref to i64), i64 0)
82+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 %v) ]
83+
ret i32 %v0
6084
}
6185

62-
define i32 @test_ptrauth_call_sign(ptr %p) {
63-
; CHECK-LABEL: @test_ptrauth_call_sign(
64-
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]()
65-
; CHECK-NEXT: ret i32 [[TMP3]]
86+
define i32 @test_ptrauth_call_mismatch_blend_addr(i32 %a0) {
87+
; CHECK-LABEL: @test_ptrauth_call_mismatch_blend_addr(
88+
; CHECK-NEXT: [[V:%.*]] = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_addr_disc.ref to i64), i64 1234)
89+
; CHECK-NEXT: [[V0:%.*]] = call i32 ptrauth (ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 [[A0:%.*]]) [ "ptrauth"(i32 1, i64 [[V]]) ]
90+
; CHECK-NEXT: ret i32 [[V0]]
6691
;
67-
%tmp0 = ptrtoint ptr %p to i64
68-
%tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 5678)
69-
%tmp2 = inttoptr i64 %tmp1 to ptr
70-
%tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
71-
ret i32 %tmp3
92+
%v = call i64 @llvm.ptrauth.blend(i64 ptrtoint (ptr @f_addr_disc.ref to i64), i64 1234)
93+
%v0 = call i32 ptrauth(ptr @f, i32 1, i64 1234, ptr @f_both_disc.ref)(i32 %a0) [ "ptrauth"(i32 1, i64 %v) ]
94+
ret i32 %v0
7295
}
7396

74-
declare i64 @llvm.ptrauth.auth(i64, i32, i64)
75-
declare i64 @llvm.ptrauth.sign(i64, i32, i64)
76-
declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
7797
declare i64 @llvm.ptrauth.blend(i64, i64)

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
2+
; RUN: opt < %s -passes=instcombine -S | FileCheck %s
3+
4+
define i32 @test_ptrauth_call_resign(ptr %p) {
5+
; CHECK-LABEL: @test_ptrauth_call_resign(
6+
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 0, i64 1234) ]
7+
; CHECK-NEXT: ret i32 [[TMP3]]
8+
;
9+
%tmp0 = ptrtoint ptr %p to i64
10+
%tmp1 = call i64 @llvm.ptrauth.resign(i64 %tmp0, i32 0, i64 1234, i32 2, i64 5678)
11+
%tmp2 = inttoptr i64 %tmp1 to ptr
12+
%tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
13+
ret i32 %tmp3
14+
}
15+
16+
define i32 @test_ptrauth_call_resign_blend(ptr %pp) {
17+
; CHECK-LABEL: @test_ptrauth_call_resign_blend(
18+
; CHECK-NEXT: [[TMP01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
19+
; CHECK-NEXT: [[TMP6:%.*]] = call i32 [[TMP01]]() [ "ptrauth"(i32 0, i64 1234) ]
20+
; CHECK-NEXT: ret i32 [[TMP6]]
21+
;
22+
%tmp0 = load ptr, ptr %pp, align 8
23+
%tmp1 = ptrtoint ptr %pp to i64
24+
%tmp2 = ptrtoint ptr %tmp0 to i64
25+
%tmp3 = call i64 @llvm.ptrauth.blend(i64 %tmp1, i64 5678)
26+
%tmp4 = call i64 @llvm.ptrauth.resign(i64 %tmp2, i32 0, i64 1234, i32 1, i64 %tmp3)
27+
%tmp5 = inttoptr i64 %tmp4 to ptr
28+
%tmp6 = call i32 %tmp5() [ "ptrauth"(i32 1, i64 %tmp3) ]
29+
ret i32 %tmp6
30+
}
31+
32+
define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
33+
; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
34+
; CHECK-NEXT: [[TMP01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
35+
; CHECK-NEXT: [[TMP1:%.*]] = ptrtoint ptr [[PP]] to i64
36+
; CHECK-NEXT: [[TMP3:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[TMP1]], i64 5678)
37+
; CHECK-NEXT: [[TMP6:%.*]] = call i32 [[TMP01]]() [ "ptrauth"(i32 1, i64 [[TMP3]]) ]
38+
; CHECK-NEXT: ret i32 [[TMP6]]
39+
;
40+
%tmp0 = load ptr, ptr %pp, align 8
41+
%tmp1 = ptrtoint ptr %pp to i64
42+
%tmp2 = ptrtoint ptr %tmp0 to i64
43+
%tmp3 = call i64 @llvm.ptrauth.blend(i64 %tmp1, i64 5678)
44+
%tmp4 = call i64 @llvm.ptrauth.resign(i64 %tmp2, i32 1, i64 %tmp3, i32 0, i64 1234)
45+
%tmp5 = inttoptr i64 %tmp4 to ptr
46+
%tmp6 = call i32 %tmp5() [ "ptrauth"(i32 0, i64 1234) ]
47+
ret i32 %tmp6
48+
}
49+
50+
define i32 @test_ptrauth_call_auth(ptr %p) {
51+
; CHECK-LABEL: @test_ptrauth_call_auth(
52+
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 2, i64 5678) ]
53+
; CHECK-NEXT: ret i32 [[TMP3]]
54+
;
55+
%tmp0 = ptrtoint ptr %p to i64
56+
%tmp1 = call i64 @llvm.ptrauth.auth(i64 %tmp0, i32 2, i64 5678)
57+
%tmp2 = inttoptr i64 %tmp1 to ptr
58+
%tmp3 = call i32 %tmp2()
59+
ret i32 %tmp3
60+
}
61+
62+
define i32 @test_ptrauth_call_sign(ptr %p) {
63+
; CHECK-LABEL: @test_ptrauth_call_sign(
64+
; CHECK-NEXT: [[TMP3:%.*]] = call i32 [[P:%.*]]()
65+
; CHECK-NEXT: ret i32 [[TMP3]]
66+
;
67+
%tmp0 = ptrtoint ptr %p to i64
68+
%tmp1 = call i64 @llvm.ptrauth.sign(i64 %tmp0, i32 2, i64 5678)
69+
%tmp2 = inttoptr i64 %tmp1 to ptr
70+
%tmp3 = call i32 %tmp2() [ "ptrauth"(i32 2, i64 5678) ]
71+
ret i32 %tmp3
72+
}
73+
74+
declare i64 @llvm.ptrauth.auth(i64, i32, i64)
75+
declare i64 @llvm.ptrauth.sign(i64, i32, i64)
76+
declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
77+
declare i64 @llvm.ptrauth.blend(i64, i64)

0 commit comments

Comments
 (0)