Skip to content

Commit 54b4005

Browse files
atrosinenkoatrosinenko
committed
[AArch64][PAC] Rework discriminator analysis for calls and tail calls
Make use of fixupBlendComponents for AUTH_TCRETURN[_BTI] and for BLRA[_RVMARKER] pseudos the same way it is done for AUT/PAC/AUTPAC. This patch unifies discriminator analysis for DAGISel and GlobalISel and improves cross-BB analysis in case of DAGISel.
1 parent 42a8fe9 commit 54b4005

File tree

5 files changed

+56
-78
lines changed

5 files changed

+56
-78
lines changed

‎llvm/lib/Target/AArch64/AArch64ISelLowering.cpp‎

Lines changed: 22 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -347,40 +347,6 @@ static bool isZeroingInactiveLanes(SDValue Op) {
347347
}
348348
}
349349

350-
static std::tuple<SDValue, SDValue>
351-
extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
352-
SDLoc DL(Disc);
353-
SDValue AddrDisc;
354-
SDValue ConstDisc;
355-
356-
// If this is a blend, remember the constant and address discriminators.
357-
// Otherwise, it's either a constant discriminator, or a non-blended
358-
// address discriminator.
359-
if (Disc->getOpcode() == ISD::INTRINSIC_WO_CHAIN &&
360-
Disc->getConstantOperandVal(0) == Intrinsic::ptrauth_blend) {
361-
AddrDisc = Disc->getOperand(1);
362-
ConstDisc = Disc->getOperand(2);
363-
} else {
364-
ConstDisc = Disc;
365-
}
366-
367-
// If the constant discriminator (either the blend RHS, or the entire
368-
// discriminator value) isn't a 16-bit constant, bail out, and let the
369-
// discriminator be computed separately.
370-
const auto *ConstDiscN = dyn_cast<ConstantSDNode>(ConstDisc);
371-
if (!ConstDiscN || !isUInt<16>(ConstDiscN->getZExtValue()))
372-
return std::make_tuple(DAG->getTargetConstant(0, DL, MVT::i64), Disc);
373-
374-
// If there's no address discriminator, use NoRegister, which we'll later
375-
// replace with XZR, or directly use a Z variant of the inst. when available.
376-
if (!AddrDisc)
377-
AddrDisc = DAG->getRegister(AArch64::NoRegister, MVT::i64);
378-
379-
return std::make_tuple(
380-
DAG->getTargetConstant(ConstDiscN->getZExtValue(), DL, MVT::i64),
381-
AddrDisc);
382-
}
383-
384350
AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
385351
const AArch64Subtarget &STI)
386352
: TargetLowering(TM), Subtarget(&STI) {
@@ -3130,6 +3096,10 @@ void AArch64TargetLowering::fixupBlendComponents(
31303096
isUInt<16>(MaybeBlend->getOperand(1).getImm())) {
31313097
AddrDisc = AArch64::NoRegister;
31323098
IntDisc = MaybeBlend->getOperand(1).getImm();
3099+
} else if (MaybeBlend->getOpcode() == AArch64::COPY &&
3100+
MaybeBlend->getOperand(1).getReg() == AArch64::XZR) {
3101+
AddrDisc = AArch64::NoRegister;
3102+
IntDisc = 0;
31333103
}
31343104
}
31353105

@@ -3266,6 +3236,22 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
32663236
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
32673237
&AArch64::GPR64noipRegClass);
32683238
return BB;
3239+
case AArch64::AUTH_TCRETURN:
3240+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3241+
&AArch64::tcGPR64RegClass);
3242+
return BB;
3243+
case AArch64::AUTH_TCRETURN_BTI:
3244+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3245+
&AArch64::tcGPRnotx16x17RegClass);
3246+
return BB;
3247+
case AArch64::BLRA:
3248+
fixupBlendComponents(MI, BB, MI.getOperand(2), MI.getOperand(3),
3249+
&AArch64::GPR64noipRegClass);
3250+
return BB;
3251+
case AArch64::BLRA_RVMARKER:
3252+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3253+
&AArch64::GPR64noipRegClass);
3254+
return BB;
32693255
}
32703256
}
32713257

@@ -9518,18 +9504,13 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
95189504
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
95199505
"Invalid auth call key");
95209506

9521-
// Split the discriminator into address/integer components.
9522-
SDValue AddrDisc, IntDisc;
9523-
std::tie(IntDisc, AddrDisc) =
9524-
extractPtrauthBlendDiscriminators(CLI.PAI->Discriminator, &DAG);
9525-
95269507
if (Opc == AArch64ISD::CALL_RVMARKER)
95279508
Opc = AArch64ISD::AUTH_CALL_RVMARKER;
95289509
else
95299510
Opc = IsTailCall ? AArch64ISD::AUTH_TC_RETURN : AArch64ISD::AUTH_CALL;
95309511
Ops.push_back(DAG.getTargetConstant(Key, DL, MVT::i32));
9531-
Ops.push_back(IntDisc);
9532-
Ops.push_back(AddrDisc);
9512+
Ops.push_back(DAG.getTargetConstant(/*IntDisc=*/0, DL, MVT::i64));
9513+
Ops.push_back(CLI.PAI->Discriminator);
95339514
}
95349515

95359516
// Add argument registers to the end of the list so that they are known live

‎llvm/lib/Target/AArch64/AArch64InstrInfo.td‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2086,6 +2086,7 @@ let Predicates = [HasPAuth] in {
20862086
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
20872087
let Defs = [X16,X17,LR];
20882088
let Uses = [SP];
2089+
let usesCustomInserter = 1;
20892090
}
20902091

20912092
def BLRA_RVMARKER : Pseudo<
@@ -2098,6 +2099,7 @@ let Predicates = [HasPAuth] in {
20982099
let isCall = 1;
20992100
let Defs = [X16,X17,LR];
21002101
let Uses = [SP];
2102+
let usesCustomInserter = 1;
21012103
}
21022104

21032105
// BRA pseudo, generalized version of BRAA/BRAB/Z.
@@ -2242,7 +2244,7 @@ let Predicates = [HasPAuth] in {
22422244
// make sure at least one register is usable as a scratch one - for that
22432245
// purpose, use tcGPRnotx16x17 register class for one of the operands.
22442246
let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
2245-
Defs = [X16,X17], Uses = [SP] in {
2247+
usesCustomInserter = 1, Defs = [X16,X17], Uses = [SP] in {
22462248
def AUTH_TCRETURN
22472249
: Pseudo<(outs), (ins tcGPRnotx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
22482250
i64imm:$Disc, tcGPR64:$AddrDisc),

‎llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1169,15 +1169,9 @@ bool AArch64CallLowering::lowerTailCall(
11691169
Info.PAI->Key == AArch64PACKey::IB) &&
11701170
"Invalid auth call key");
11711171
MIB.addImm(Info.PAI->Key);
1172-
1173-
Register AddrDisc = 0;
1174-
uint16_t IntDisc = 0;
1175-
std::tie(IntDisc, AddrDisc) =
1176-
extractPtrauthBlendDiscriminators(Info.PAI->Discriminator, MRI);
1177-
1178-
MIB.addImm(IntDisc);
1179-
MIB.addUse(AddrDisc);
1180-
if (AddrDisc != AArch64::NoRegister) {
1172+
MIB.addImm(/*IntDisc=*/0);
1173+
MIB.addUse(Info.PAI->Discriminator);
1174+
if (Info.PAI->Discriminator != AArch64::NoRegister) {
11811175
MIB->getOperand(4).setReg(constrainOperandRegClass(
11821176
MF, *TRI, MRI, *MF.getSubtarget().getInstrInfo(),
11831177
*MF.getSubtarget().getRegBankInfo(), *MIB, MIB->getDesc(),
@@ -1443,15 +1437,9 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
14431437
Info.PAI->Key == AArch64PACKey::IB) &&
14441438
"Invalid auth call key");
14451439
MIB.addImm(Info.PAI->Key);
1446-
1447-
Register AddrDisc = 0;
1448-
uint16_t IntDisc = 0;
1449-
std::tie(IntDisc, AddrDisc) =
1450-
extractPtrauthBlendDiscriminators(Info.PAI->Discriminator, MRI);
1451-
1452-
MIB.addImm(IntDisc);
1453-
MIB.addUse(AddrDisc);
1454-
if (AddrDisc != AArch64::NoRegister) {
1440+
MIB.addImm(/*IntDisc=*/0);
1441+
MIB.addUse(Info.PAI->Discriminator);
1442+
if (Info.PAI->Discriminator != AArch64::NoRegister) {
14551443
constrainOperandRegClass(MF, *TRI, MRI, *MF.getSubtarget().getInstrInfo(),
14561444
*MF.getSubtarget().getRegBankInfo(), *MIB,
14571445
MIB->getDesc(), MIB->getOperand(CalleeOpNo + 3),

‎llvm/test/CodeGen/AArch64/ptrauth-call.ll‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -202,9 +202,9 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
202202
; DARWIN-NEXT: mov x17, x0
203203
; DARWIN-NEXT: movk x17, #6503, lsl #48
204204
; DARWIN-NEXT: autda x16, x17
205-
; DARWIN-NEXT: ldr x8, [x16]
205+
; DARWIN-NEXT: ldr x9, [x16]
206206
; DARWIN-NEXT: movk x16, #34646, lsl #48
207-
; DARWIN-NEXT: blraa x8, x16
207+
; DARWIN-NEXT: blraa x9, x16
208208
; DARWIN-NEXT: mov w0, #42
209209
; DARWIN-NEXT: ldp x29, x30, [sp], #16
210210
; ELF-NEXT: str x30, [sp, #-16]!

‎llvm/test/CodeGen/AArch64/ptrauth-isel.ll‎

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -554,7 +554,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
554554
; DAGISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY3]]
555555
; DAGISEL-NEXT: {{ $}}
556556
; DAGISEL-NEXT: bb.2.exit:
557-
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
557+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
558+
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 42, [[COPY4]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
558559
;
559560
; GISEL-LABEL: name: auth_tcreturn_blend_components
560561
; GISEL: bb.1.entry:
@@ -564,8 +565,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
564565
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprnotx16x17 = COPY $x20
565566
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
566567
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
567-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
568-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
568+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
569+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgpr64 = MOVKXi [[LDRXui]], 42, 48
569570
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
570571
; GISEL-NEXT: B %bb.2
571572
; GISEL-NEXT: {{ $}}
@@ -575,7 +576,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
575576
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
576577
; GISEL-NEXT: {{ $}}
577578
; GISEL-NEXT: bb.3.exit:
578-
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
579+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
580+
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
579581
entry:
580582
%addrdisc = load i64, ptr @discvar
581583
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -613,7 +615,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
613615
; DAGISEL-NEXT: {{ $}}
614616
; DAGISEL-NEXT: bb.2.exit:
615617
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgprx16x17 = COPY [[COPY1]]
616-
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
618+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
619+
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 42, [[COPY5]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
617620
;
618621
; GISEL-LABEL: name: auth_tcreturn_bti_blend_components
619622
; GISEL: bb.1.entry:
@@ -623,8 +626,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
623626
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprx16x17 = COPY $x20
624627
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
625628
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
626-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgprnotx16x17 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
627-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
629+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
630+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgprnotx16x17 = MOVKXi [[LDRXui]], 42, 48
628631
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
629632
; GISEL-NEXT: B %bb.2
630633
; GISEL-NEXT: {{ $}}
@@ -634,7 +637,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
634637
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
635638
; GISEL-NEXT: {{ $}}
636639
; GISEL-NEXT: bb.3.exit:
637-
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
640+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
641+
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
638642
entry:
639643
%addrdisc = load i64, ptr @discvar
640644
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -672,10 +676,11 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
672676
; DAGISEL-NEXT: {{ $}}
673677
; DAGISEL-NEXT: bb.2.exit:
674678
; DAGISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
675-
; DAGISEL-NEXT: BLRA [[COPY1]], 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
679+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
680+
; DAGISEL-NEXT: BLRA [[COPY1]], 1, 42, [[COPY4]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
676681
; DAGISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
677-
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
678-
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
682+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:gpr64sp = COPY $x0
683+
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY5]], 123, 0
679684
; DAGISEL-NEXT: $x0 = COPY [[ADDXri]]
680685
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
681686
;
@@ -688,21 +693,23 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
688693
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
689694
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
690695
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
691-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
696+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64 = MOVKXi [[LDRXui]], 42, 48
692697
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
693698
; GISEL-NEXT: B %bb.2
694699
; GISEL-NEXT: {{ $}}
695700
; GISEL-NEXT: bb.2.next:
696701
; GISEL-NEXT: successors: %bb.3(0x80000000)
697702
; GISEL-NEXT: {{ $}}
698-
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
703+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64common = COPY [[MOVKXi]]
704+
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY2]]
699705
; GISEL-NEXT: {{ $}}
700706
; GISEL-NEXT: bb.3.exit:
701707
; GISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
702-
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
708+
; GISEL-NEXT: [[COPY3:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
709+
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[COPY3]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
703710
; GISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
704-
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64sp = COPY $x0
705-
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY2]], 123, 0
711+
; GISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
712+
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
706713
; GISEL-NEXT: $x0 = COPY [[ADDXri]]
707714
; GISEL-NEXT: RET_ReallyLR implicit $x0
708715
entry:

0 commit comments

Comments
 (0)