Skip to content

Commit 7bf2578

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 612623c commit 7bf2578

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
@@ -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) {
@@ -3156,6 +3122,12 @@ void AArch64TargetLowering::fixupPtrauthDiscriminator(
31563122
IntDisc = DiscMI->getOperand(1).getImm();
31573123
}
31583124
break;
3125+
case AArch64::COPY:
3126+
if (DiscMI->getOperand(1).getReg() == AArch64::XZR) {
3127+
AddrDisc = AArch64::NoRegister;
3128+
IntDisc = 0;
3129+
}
3130+
break;
31593131
}
31603132
}
31613133

@@ -3292,6 +3264,22 @@ MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
32923264
fixupPtrauthDiscriminator(MI, BB, MI.getOperand(4), MI.getOperand(5),
32933265
&AArch64::GPR64noipRegClass);
32943266
return BB;
3267+
case AArch64::AUTH_TCRETURN:
3268+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3269+
&AArch64::tcGPR64RegClass);
3270+
return BB;
3271+
case AArch64::AUTH_TCRETURN_BTI:
3272+
fixupBlendComponents(MI, BB, MI.getOperand(3), MI.getOperand(4),
3273+
&AArch64::tcGPRnotx16x17RegClass);
3274+
return BB;
3275+
case AArch64::BLRA:
3276+
fixupBlendComponents(MI, BB, MI.getOperand(2), MI.getOperand(3),
3277+
&AArch64::GPR64noipRegClass);
3278+
return BB;
3279+
case AArch64::BLRA_RVMARKER:
3280+
fixupBlendComponents(MI, BB, MI.getOperand(4), MI.getOperand(5),
3281+
&AArch64::GPR64noipRegClass);
3282+
return BB;
32953283
}
32963284
}
32973285

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

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

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

20872088
def BLRA_RVMARKER : Pseudo<
@@ -2094,6 +2095,7 @@ let Predicates = [HasPAuth] in {
20942095
let isCall = 1;
20952096
let Defs = [X16,X17,LR];
20962097
let Uses = [SP];
2098+
let usesCustomInserter = 1;
20972099
}
20982100

20992101
// BRA pseudo, generalized version of BRAA/BRAB/Z.
@@ -2239,7 +2241,7 @@ let Predicates = [HasPAuth] in {
22392241
// make sure at least one register is usable as a scratch one - for that
22402242
// purpose, use tcGPRnotx16x17 register class for one of the operands.
22412243
let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Size = 16,
2242-
Defs = [X16,X17], Uses = [SP] in {
2244+
usesCustomInserter = 1, Defs = [X16,X17], Uses = [SP] in {
22432245
def AUTH_TCRETURN
22442246
: Pseudo<(outs), (ins tcGPRnotx16x17:$dst, i32imm:$FPDiff, i32imm:$Key,
22452247
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)