Skip to content

Commit 6e3bbe5

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 d29ded0 commit 6e3bbe5

File tree

5 files changed

+58
-78
lines changed

5 files changed

+58
-78
lines changed

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

Lines changed: 24 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -350,40 +350,6 @@ static bool isZeroingInactiveLanes(SDValue Op) {
350350
}
351351
}
352352

353-
static std::tuple<SDValue, SDValue>
354-
extractPtrauthBlendDiscriminators(SDValue Disc, SelectionDAG *DAG) {
355-
SDLoc DL(Disc);
356-
SDValue AddrDisc;
357-
SDValue ConstDisc;
358-
359-
// If this is a blend, remember the constant and address discriminators.
360-
// Otherwise, it's either a constant discriminator, or a non-blended
361-
// address discriminator.
362-
if (Disc->getOpcode() == ISD::INTRINSIC_WO_CHAIN &&
363-
Disc->getConstantOperandVal(0) == Intrinsic::ptrauth_blend) {
364-
AddrDisc = Disc->getOperand(1);
365-
ConstDisc = Disc->getOperand(2);
366-
} else {
367-
ConstDisc = Disc;
368-
}
369-
370-
// If the constant discriminator (either the blend RHS, or the entire
371-
// discriminator value) isn't a 16-bit constant, bail out, and let the
372-
// discriminator be computed separately.
373-
const auto *ConstDiscN = dyn_cast<ConstantSDNode>(ConstDisc);
374-
if (!ConstDiscN || !isUInt<16>(ConstDiscN->getZExtValue()))
375-
return std::make_tuple(DAG->getTargetConstant(0, DL, MVT::i64), Disc);
376-
377-
// If there's no address discriminator, use NoRegister, which we'll later
378-
// replace with XZR, or directly use a Z variant of the inst. when available.
379-
if (!AddrDisc)
380-
AddrDisc = DAG->getRegister(AArch64::NoRegister, MVT::i64);
381-
382-
return std::make_tuple(
383-
DAG->getTargetConstant(ConstDiscN->getZExtValue(), DL, MVT::i64),
384-
AddrDisc);
385-
}
386-
387353
AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
388354
const AArch64Subtarget &STI)
389355
: TargetLowering(TM), Subtarget(&STI) {
@@ -3159,6 +3125,12 @@ void AArch64TargetLowering::fixupPtrauthDiscriminator(
31593125
IntDisc = DiscMI->getOperand(1).getImm();
31603126
}
31613127
break;
3128+
case AArch64::COPY:
3129+
if (DiscMI->getOperand(1).getReg() == AArch64::XZR) {
3130+
AddrDisc = AArch64::NoRegister;
3131+
IntDisc = 0;
3132+
}
3133+
break;
31623134
}
31633135
}
31643136

@@ -3295,6 +3267,22 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
32953267
fixupPtrauthDiscriminator(MI, BB, MI.getOperand(4), MI.getOperand(5),
32963268
&AArch64::GPR64noipRegClass);
32973269
return BB;
3270+
case AArch64::AUTH_TCRETURN:
3271+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3272+
&AArch64::tcGPR64RegClass);
3273+
return BB;
3274+
case AArch64::AUTH_TCRETURN_BTI:
3275+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3276+
&AArch64::tcGPRnotx16x17RegClass);
3277+
return BB;
3278+
case AArch64::BLRA:
3279+
fixupBlendComponents(MI, BB, MI.getOperand(2), MI.getOperand(3),
3280+
&AArch64::GPR64noipRegClass);
3281+
return BB;
3282+
case AArch64::BLRA_RVMARKER:
3283+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3284+
&AArch64::GPR64noipRegClass);
3285+
return BB;
32983286
}
32993287
}
33003288

@@ -9548,18 +9536,13 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
95489536
assert((Key == AArch64PACKey::IA || Key == AArch64PACKey::IB) &&
95499537
"Invalid auth call key");
95509538

9551-
// Split the discriminator into address/integer components.
9552-
SDValue AddrDisc, IntDisc;
9553-
std::tie(IntDisc, AddrDisc) =
9554-
extractPtrauthBlendDiscriminators(CLI.PAI->Discriminator, &DAG);
9555-
95569539
if (Opc == AArch64ISD::CALL_RVMARKER)
95579540
Opc = AArch64ISD::AUTH_CALL_RVMARKER;
95589541
else
95599542
Opc = IsTailCall ? AArch64ISD::AUTH_TC_RETURN : AArch64ISD::AUTH_CALL;
95609543
Ops.push_back(DAG.getTargetConstant(Key, DL, MVT::i32));
9561-
Ops.push_back(IntDisc);
9562-
Ops.push_back(AddrDisc);
9544+
Ops.push_back(DAG.getTargetConstant(/*IntDisc=*/0, DL, MVT::i64));
9545+
Ops.push_back(CLI.PAI->Discriminator);
95639546
}
95649547

95659548
// 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
@@ -2083,6 +2083,7 @@ let Predicates = [HasPAuth] in {
20832083
let Size = 12; // 4 fixed + 8 variable, to compute discriminator.
20842084
let Defs = [X16,X17,LR];
20852085
let Uses = [SP];
2086+
let usesCustomInserter = 1;
20862087
}
20872088

20882089
def BLRA_RVMARKER : Pseudo<
@@ -2095,6 +2096,7 @@ let Predicates = [HasPAuth] in {
20952096
let isCall = 1;
20962097
let Defs = [X16,X17,LR];
20972098
let Uses = [SP];
2099+
let usesCustomInserter = 1;
20982100
}
20992101

21002102
// BRA pseudo, generalized version of BRAA/BRAB/Z.
@@ -2240,7 +2242,7 @@ let Predicates = [HasPAuth] in {
22402242
// make sure at least one register is usable as a scratch one - for that
22412243
// purpose, use tcGPRnotx16x17 register class for one of the operands.
22422244
let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
2243-
Defs = [X16,X17], Uses = [SP] in {
2245+
usesCustomInserter = 1, Defs = [X16,X17], Uses = [SP] in {
22442246
def AUTH_TCRETURN
22452247
: Pseudo<(outs), (ins tcGPRnotx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
22462248
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
@@ -618,7 +618,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
618618
; DAGISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY3]]
619619
; DAGISEL-NEXT: {{ $}}
620620
; DAGISEL-NEXT: bb.2.exit:
621-
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
621+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
622+
; DAGISEL-NEXT: AUTH_TCRETURN [[COPY1]], 0, 1, 42, [[COPY4]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
622623
;
623624
; GISEL-LABEL: name: auth_tcreturn_blend_components
624625
; GISEL: bb.1.entry:
@@ -628,8 +629,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
628629
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprnotx16x17 = COPY $x20
629630
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
630631
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
631-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
632-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
632+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
633+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgpr64 = MOVKXi [[LDRXui]], 42, 48
633634
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
634635
; GISEL-NEXT: B %bb.2
635636
; GISEL-NEXT: {{ $}}
@@ -639,7 +640,8 @@ define preserve_nonecc i64 @auth_tcreturn_blend_components(ptr %callee, i1 %cond
639640
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
640641
; GISEL-NEXT: {{ $}}
641642
; GISEL-NEXT: bb.3.exit:
642-
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
643+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgpr64 = COPY [[LDRXui]]
644+
; GISEL-NEXT: AUTH_TCRETURN [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
643645
entry:
644646
%addrdisc = load i64, ptr @discvar
645647
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -677,7 +679,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
677679
; DAGISEL-NEXT: {{ $}}
678680
; DAGISEL-NEXT: bb.2.exit:
679681
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:tcgprx16x17 = COPY [[COPY1]]
680-
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 0, [[COPY2]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
682+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
683+
; DAGISEL-NEXT: AUTH_TCRETURN_BTI [[COPY4]], 0, 1, 42, [[COPY5]], csr_aarch64_noneregs, implicit-def dead $x16, implicit-def dead $x17, implicit $sp
681684
;
682685
; GISEL-LABEL: name: auth_tcreturn_bti_blend_components
683686
; GISEL: bb.1.entry:
@@ -687,8 +690,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
687690
; GISEL-NEXT: [[COPY:%[0-9]+]]:tcgprx16x17 = COPY $x20
688691
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
689692
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
690-
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:tcgprnotx16x17 = 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
693+
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
694+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:tcgprnotx16x17 = MOVKXi [[LDRXui]], 42, 48
692695
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
693696
; GISEL-NEXT: B %bb.2
694697
; GISEL-NEXT: {{ $}}
@@ -698,7 +701,8 @@ define preserve_nonecc i64 @auth_tcreturn_bti_blend_components(ptr %callee, i1 %
698701
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
699702
; GISEL-NEXT: {{ $}}
700703
; GISEL-NEXT: bb.3.exit:
701-
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
704+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:tcgprnotx16x17 = COPY [[LDRXui]]
705+
; GISEL-NEXT: AUTH_TCRETURN_BTI [[COPY]], 0, 1, 42, [[COPY2]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit $sp
702706
entry:
703707
%addrdisc = load i64, ptr @discvar
704708
%disc = call i64 @llvm.ptrauth.blend(i64 %addrdisc, i64 42)
@@ -736,10 +740,11 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
736740
; DAGISEL-NEXT: {{ $}}
737741
; DAGISEL-NEXT: bb.2.exit:
738742
; DAGISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
739-
; 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
743+
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
744+
; 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
740745
; DAGISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
741-
; DAGISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
742-
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
746+
; DAGISEL-NEXT: [[COPY5:%[0-9]+]]:gpr64sp = COPY $x0
747+
; DAGISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY5]], 123, 0
743748
; DAGISEL-NEXT: $x0 = COPY [[ADDXri]]
744749
; DAGISEL-NEXT: RET_ReallyLR implicit $x0
745750
;
@@ -752,21 +757,23 @@ define preserve_nonecc i64 @blra_blend_components(ptr %callee, i1 %cond.b) {
752757
; GISEL-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $w0
753758
; GISEL-NEXT: [[ADRP:%[0-9]+]]:gpr64common = ADRP target-flags(aarch64-page) @discvar
754759
; GISEL-NEXT: [[LDRXui:%[0-9]+]]:gpr64 = LDRXui [[ADRP]], target-flags(aarch64-pageoff, aarch64-nc) @discvar :: (dereferenceable load (s64) from @discvar)
755-
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64common = MOVKXi [[LDRXui]], 42, 48
760+
; GISEL-NEXT: [[MOVKXi:%[0-9]+]]:gpr64 = MOVKXi [[LDRXui]], 42, 48
756761
; GISEL-NEXT: TBZW [[COPY1]], 0, %bb.3
757762
; GISEL-NEXT: B %bb.2
758763
; GISEL-NEXT: {{ $}}
759764
; GISEL-NEXT: bb.2.next:
760765
; GISEL-NEXT: successors: %bb.3(0x80000000)
761766
; GISEL-NEXT: {{ $}}
762-
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[MOVKXi]]
767+
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64common = COPY [[MOVKXi]]
768+
; GISEL-NEXT: INLINEASM &nop, 1 /* sideeffect attdialect */, 3866633 /* reguse:GPR64common */, [[COPY2]]
763769
; GISEL-NEXT: {{ $}}
764770
; GISEL-NEXT: bb.3.exit:
765771
; GISEL-NEXT: ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
766-
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[LDRXui]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
772+
; GISEL-NEXT: [[COPY3:%[0-9]+]]:gpr64noip = COPY [[LDRXui]]
773+
; GISEL-NEXT: BLRA [[COPY]], 1, 42, [[COPY3]], csr_aarch64_noneregs, implicit-def $x16, implicit-def $x17, implicit-def $lr, implicit $sp, implicit-def $x0
767774
; GISEL-NEXT: ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
768-
; GISEL-NEXT: [[COPY2:%[0-9]+]]:gpr64sp = COPY $x0
769-
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY2]], 123, 0
775+
; GISEL-NEXT: [[COPY4:%[0-9]+]]:gpr64sp = COPY $x0
776+
; GISEL-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY4]], 123, 0
770777
; GISEL-NEXT: $x0 = COPY [[ADDXri]]
771778
; GISEL-NEXT: RET_ReallyLR implicit $x0
772779
entry:

0 commit comments

Comments
 (0)