Skip to content

Commit be6d2af

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 63a240f commit be6d2af

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) {
@@ -3124,6 +3090,10 @@ void AArch64TargetLowering::fixupBlendComponents(
31243090
isUInt<16>(MaybeBlend->getOperand(1).getImm())) {
31253091
AddrDisc = AArch64::NoRegister;
31263092
IntDisc = MaybeBlend->getOperand(1).getImm();
3093+
} else if (MaybeBlend->getOpcode() == AArch64::COPY &&
3094+
MaybeBlend->getOperand(1).getReg() == AArch64::XZR) {
3095+
AddrDisc = AArch64::NoRegister;
3096+
IntDisc = 0;
31273097
}
31283098
}
31293099

@@ -3256,6 +3226,22 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
32563226
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
32573227
&AArch64::GPR64noipRegClass);
32583228
return BB;
3229+
case AArch64::AUTH_TCRETURN:
3230+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3231+
&AArch64::tcGPR64RegClass);
3232+
return BB;
3233+
case AArch64::AUTH_TCRETURN_BTI:
3234+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3235+
&AArch64::tcGPRnotx16x17RegClass);
3236+
return BB;
3237+
case AArch64::BLRA:
3238+
fixupBlendComponents(MI, BB, MI.getOperand(2), MI.getOperand(3),
3239+
&AArch64::GPR64noipRegClass);
3240+
return BB;
3241+
case AArch64::BLRA_RVMARKER:
3242+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3243+
&AArch64::GPR64noipRegClass);
3244+
return BB;
32593245
}
32603246
}
32613247

@@ -9508,18 +9494,13 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
95089494
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
95099495
"Invalid auth call key");
95109496

9511-
// Split the discriminator into address/integer components.
9512-
SDValue AddrDisc, IntDisc;
9513-
std::tie(IntDisc, AddrDisc) =
9514-
extractPtrauthBlendDiscriminators(CLI.PAI->Discriminator, &DAG);
9515-
95169497
if (Opc == AArch64ISD::CALL_RVMARKER)
95179498
Opc = AArch64ISD::AUTH_CALL_RVMARKER;
95189499
else
95199500
Opc = IsTailCall ? AArch64ISD::AUTH_TC_RETURN : AArch64ISD::AUTH_CALL;
95209501
Ops.push_back(DAG.getTargetConstant(Key, DL, MVT::i32));
9521-
Ops.push_back(IntDisc);
9522-
Ops.push_back(AddrDisc);
9502+
Ops.push_back(DAG.getTargetConstant(/*IntDisc=*/0, DL, MVT::i64));
9503+
Ops.push_back(CLI.PAI->Discriminator);
95239504
}
95249505

95259506
// 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
@@ -2080,6 +2080,7 @@ let Predicates = [HasPAuth] in {
20802080
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
20812081
let Defs = [X16,X17,LR];
20822082
let Uses = [SP];
2083+
let usesCustomInserter = 1;
20832084
}
20842085

20852086
def BLRA_RVMARKER : Pseudo<
@@ -2092,6 +2093,7 @@ let Predicates = [HasPAuth] in {
20922093
let isCall = 1;
20932094
let Defs = [X16,X17,LR];
20942095
let Uses = [SP];
2096+
let usesCustomInserter = 1;
20952097
}
20962098

20972099
// BRA pseudo, generalized version of BRAA/BRAB/Z.
@@ -2221,7 +2223,7 @@ let Predicates = [HasPAuth] in {
22212223
// make sure at least one register is usable as a scratch one - for that
22222224
// purpose, use tcGPRnotx16x17 register class for one of the operands.
22232225
let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
2224-
Defs = [X16,X17], Uses = [SP] in {
2226+
usesCustomInserter = 1, Defs = [X16,X17], Uses = [SP] in {
22252227
def AUTH_TCRETURN
22262228
: Pseudo<(outs), (ins tcGPRnotx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
22272229
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
@@ -196,9 +196,9 @@ define i32 @test_call_omit_extra_moves(ptr %objptr) #0 {
196196
; CHECK-NEXT: mov x17, x0
197197
; CHECK-NEXT: movk x17, #6503, lsl #48
198198
; CHECK-NEXT: autda x16, x17
199-
; CHECK-NEXT: ldr x8, [x16]
199+
; CHECK-NEXT: ldr x9, [x16]
200200
; CHECK-NEXT: movk x16, #34646, lsl #48
201-
; CHECK-NEXT: blraa x8, x16
201+
; CHECK-NEXT: blraa x9, x16
202202
; CHECK-NEXT: mov w0, #42
203203
; DARWIN-NEXT: ldp x29, x30, [sp], #16
204204
; ELF-NEXT: ldr x30, [sp], #16

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

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -477,7 +477,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
477477
; DAGISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY3]]
478478
; DAGISEL-NEXT: {{ $}}
479479
; DAGISEL-NEXT: bb.2.exit:
480-
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
480+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
481+
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 42, [[COPY4]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
481482
;
482483
; GISEL-LABEL: name: auth_tcreturn_blend_components
483484
; GISEL: bb.1.entry:
@@ -487,8 +488,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
487488
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprnotx16x17 = COPY $x20
488489
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
489490
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
490-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
491-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
491+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
492+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgpr64 = MOVKXi [[LDRXui]], 42, 48
492493
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
493494
; GISEL-NEXT: B %bb.2
494495
; GISEL-NEXT: {{ $}}
@@ -498,7 +499,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
498499
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
499500
; GISEL-NEXT: {{ $}}
500501
; GISEL-NEXT: bb.3.exit:
501-
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
502+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
503+
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
502504
entry:
503505
%addrdisc = load i64, ptr @discvar
504506
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -536,7 +538,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
536538
; DAGISEL-NEXT: {{ $}}
537539
; DAGISEL-NEXT: bb.2.exit:
538540
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgprx16x17 = COPY [[COPY1]]
539-
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
541+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
542+
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 42, [[COPY5]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
540543
;
541544
; GISEL-LABEL: name: auth_tcreturn_bti_blend_components
542545
; GISEL: bb.1.entry:
@@ -546,8 +549,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
546549
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprx16x17 = COPY $x20
547550
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
548551
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
549-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgprnotx16x17 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
550-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
552+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
553+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgprnotx16x17 = MOVKXi [[LDRXui]], 42, 48
551554
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
552555
; GISEL-NEXT: B %bb.2
553556
; GISEL-NEXT: {{ $}}
@@ -557,7 +560,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
557560
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
558561
; GISEL-NEXT: {{ $}}
559562
; GISEL-NEXT: bb.3.exit:
560-
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
563+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
564+
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
561565
entry:
562566
%addrdisc = load i64, ptr @discvar
563567
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -595,10 +599,11 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
595599
; DAGISEL-NEXT: {{ $}}
596600
; DAGISEL-NEXT: bb.2.exit:
597601
; DAGISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
598-
; 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
602+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
603+
; 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
599604
; DAGISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
600-
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
601-
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
605+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:gpr64sp = COPY $x0
606+
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY5]], 123, 0
602607
; DAGISEL-NEXT: $x0 = COPY [[ADDXri]]
603608
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
604609
;
@@ -611,21 +616,23 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
611616
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
612617
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
613618
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
614-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
619+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64 = MOVKXi [[LDRXui]], 42, 48
615620
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
616621
; GISEL-NEXT: B %bb.2
617622
; GISEL-NEXT: {{ $}}
618623
; GISEL-NEXT: bb.2.next:
619624
; GISEL-NEXT: successors: %bb.3(0x80000000)
620625
; GISEL-NEXT: {{ $}}
621-
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
626+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64common = COPY [[MOVKXi]]
627+
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY2]]
622628
; GISEL-NEXT: {{ $}}
623629
; GISEL-NEXT: bb.3.exit:
624630
; GISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
625-
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
631+
; GISEL-NEXT: [[COPY3:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
632+
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[COPY3]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
626633
; GISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
627-
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64sp = COPY $x0
628-
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY2]], 123, 0
634+
; GISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
635+
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
629636
; GISEL-NEXT: $x0 = COPY [[ADDXri]]
630637
; GISEL-NEXT: RET_ReallyLR implicit $x0
631638
entry:

0 commit comments

Comments
 (0)