Skip to content

Commit 7af97c8

Browse files
atrosinenkoatrosinenko
committed
Prevent immediate modifier substitution
1 parent 29c91c8 commit 7af97c8

File tree

3 files changed

+93
-65
lines changed

3 files changed

+93
-65
lines changed

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

Lines changed: 63 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -3360,62 +3360,80 @@ AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
33603360
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
33613361
const DebugLoc &DL = MI.getDebugLoc();
33623362

3363-
MachineInstr *AddrInstr = nullptr;
3364-
int64_t Offset = 0;
3365-
// Try to find the address-setting instruction, accumulating the offset
3366-
// along the way. If any unknown pattern is found, keep everything as-is.
3367-
MachineOperand *CurOp = &MI.getOperand(1);
3368-
while (CurOp) {
3369-
MachineOperand *Def = MRI.getOneDef(CurOp->getReg());
3370-
if (!Def)
3371-
return BB;
3372-
MachineInstr *DefMI = Def->getParent();
3373-
assert(DefMI != nullptr);
3374-
3375-
switch (DefMI->getOpcode()) {
3376-
case AArch64::COPY:
3377-
CurOp = &DefMI->getOperand(1);
3378-
break;
3379-
case AArch64::ADDXri:
3380-
if (DefMI->getOperand(3).getImm() != 0) // shifts are not handled
3381-
return BB;
3382-
CurOp = &DefMI->getOperand(1);
3383-
Offset += DefMI->getOperand(2).getImm();
3384-
break;
3385-
case AArch64::MOVaddr:
3386-
case AArch64::LOADgotAUTH:
3387-
AddrInstr = DefMI;
3388-
CurOp = nullptr;
3389-
break;
3390-
default:
3391-
return BB;
3363+
// Find the unique register definition, skipping copies.
3364+
auto GetUniqueDef = [&MRI](Register Reg) {
3365+
for (;;) {
3366+
MachineInstr *Def = MRI.getUniqueVRegDef(Reg);
3367+
if (!Def || Def->getOpcode() != AArch64::COPY)
3368+
return Def;
3369+
3370+
Reg = Def->getOperand(1).getReg();
33923371
}
3393-
}
3372+
};
3373+
// Find the unique register definition, skipping copies and increments.
3374+
auto GetUniqueDefPlusOffset =
3375+
[GetUniqueDef](Register Reg, int64_t &Offset) -> MachineInstr * {
3376+
for (;;) {
3377+
MachineInstr *Def = GetUniqueDef(Reg);
3378+
if (!Def || Def->getOpcode() != AArch64::ADDXri)
3379+
return Def;
3380+
3381+
if (Def->getOperand(3).getImm() != 0)
3382+
return nullptr; // shifted immediates are not handled
3383+
Reg = Def->getOperand(1).getReg();
3384+
Offset += Def->getOperand(2).getImm();
3385+
}
3386+
};
3387+
3388+
// Try to find a known address-setting instruction, accumulating the offset
3389+
// along the way. If no known pattern is found, keep everything as-is.
3390+
3391+
int64_t AddrOffset = 0;
3392+
MachineInstr *AddrDefInstr =
3393+
GetUniqueDefPlusOffset(MI.getOperand(1).getReg(), AddrOffset);
3394+
if (!AddrDefInstr)
3395+
return BB;
33943396

3395-
unsigned NewOpcode = AddrInstr->getOpcode() == AArch64::LOADgotAUTH
3396-
? AArch64::LOADgotPAC
3397-
: AArch64::MOVaddrPAC;
3398-
MachineOperand &AddrOp = AddrInstr->getOperand(1);
3397+
unsigned NewOpcode;
3398+
if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
3399+
NewOpcode = AArch64::LOADgotPAC;
3400+
else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
3401+
NewOpcode = AArch64::MOVaddrPAC;
3402+
else
3403+
return BB; // Unknown opcode.
3404+
3405+
MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
33993406
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
3400-
Offset += AddrOp.getOffset();
3407+
const GlobalValue *GV = AddrOp.getGlobal();
3408+
AddrOffset += AddrOp.getOffset();
3409+
3410+
// Detect discriminator blend computation, if any.
3411+
Register RegDisc = isPACWithZeroDisc(MI.getOpcode())
3412+
? AArch64::XZR
3413+
: MI.getOperand(2).getReg();
3414+
unsigned IntDisc = 0;
3415+
MachineInstr *MaybeBlendDef =
3416+
RegDisc == AArch64::XZR ? nullptr : GetUniqueDef(RegDisc);
3417+
if (MaybeBlendDef && MaybeBlendDef->getOpcode() == AArch64::MOVKXi &&
3418+
MaybeBlendDef->getOperand(3).getImm() == 48) {
3419+
RegDisc = MaybeBlendDef->getOperand(1).getReg();
3420+
IntDisc = MaybeBlendDef->getOperand(2).getImm();
3421+
}
34013422

34023423
// MOVaddrPAC and LOADgotPAC pseudos are expanded so that they use X16/X17
34033424
// internally, thus their restrictions on the register class of $AddrDisc
34043425
// operand are stricter than those of real PAC* instructions.
3405-
// If the original instruction accepts a discriminator operand, make sure
3406-
// it is moved out of X16/X17.
3407-
Register DiscReg = AArch64::XZR;
3408-
if (!isPACWithZeroDisc(MI.getOpcode())) {
3409-
DiscReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3410-
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), DiscReg)
3411-
.addReg(MI.getOperand(2).getReg());
3426+
if (RegDisc != AArch64::XZR) {
3427+
Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3428+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), TmpReg).addReg(RegDisc);
3429+
RegDisc = TmpReg;
34123430
}
34133431

34143432
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3415-
.addGlobalAddress(AddrOp.getGlobal(), Offset, TargetFlags)
3433+
.addGlobalAddress(GV, AddrOffset, TargetFlags)
34163434
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3417-
.addReg(DiscReg)
3418-
.addImm(0);
3435+
.addReg(RegDisc)
3436+
.addImm(IntDisc);
34193437

34203438
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
34213439
.addReg(AArch64::X16);

‎llvm/test/CodeGen/AArch64/GlobalISel/ptrauth-constant-in-code.ll‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -96,29 +96,34 @@ define ptr @foo() {
9696
@const_table_local = dso_local constant [3 x ptr] [ptr null, ptr null, ptr null]
9797
@const_table_got = constant [3 x ptr] [ptr null, ptr null, ptr null]
9898

99+
; Test that after post-processing in finalize-isel, MOVaddrPAC (or LOADgotPAC,
100+
; respectively) has both $AddrDisc and $Disc operands set. MOVaddr (or LOADgotAUTH,
101+
; respectively) and MOVKXi are not used anymore and are dead-code-eliminated
102+
; by the later passes.
103+
99104
define void @store_signed_const_local(ptr %dest) {
100105
; ISEL-MIR-LABEL: name: store_signed_const_local
101106
; ISEL-MIR: body:
102107
; ISEL-MIR: %0:gpr64common = COPY $x0
103108
; ISEL-MIR-NEXT: %10:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
104109
; ISEL-MIR-NEXT: %2:gpr64common = MOVKXi %0, 1234, 48
105-
; ISEL-MIR-NEXT: %15:gpr64noip = COPY %2
106-
; ISEL-MIR-NEXT: MOVaddrPAC @const_table_local + 8, 2, %15, 0, implicit-def $x16, implicit-def $x17
110+
; ISEL-MIR-NEXT: %15:gpr64noip = COPY %0
111+
; ISEL-MIR-NEXT: MOVaddrPAC @const_table_local + 8, 2, %15, 1234, implicit-def $x16, implicit-def $x17
107112
; ISEL-MIR-NEXT: %4:gpr64 = COPY $x16
108113
; ISEL-MIR-NEXT: %14:gpr64 = COPY %4
109114
; ISEL-MIR-NEXT: STRXui %14, %0, 0 :: (store (p0) into %ir.dest)
110115
; ISEL-MIR-NEXT: RET_ReallyLR
111116
;
112117
; ISEL-ASM-LABEL: store_signed_const_local:
113118
; ISEL-ASM-NEXT: .cfi_startproc
114-
; ISEL-ASM-NEXT: mov x8, x0
115-
; ISEL-ASM-NEXT: movk x8, #1234, lsl #48
116119
; ISEL-ASM-ELF-NEXT: adrp x16, const_table_local
117120
; ISEL-ASM-ELF-NEXT: add x16, x16, :lo12:const_table_local
118121
; ISEL-ASM-MACHO-NEXT: adrp x16, _const_table_local@PAGE
119122
; ISEL-ASM-MACHO-NEXT: add x16, x16, _const_table_local@PAGEOFF
120123
; ISEL-ASM-NEXT: add x16, x16, #8
121-
; ISEL-ASM-NEXT: pacda x16, x8
124+
; ISEL-ASM-NEXT: mov x17, x0
125+
; ISEL-ASM-NEXT: movk x17, #1234, lsl #48
126+
; ISEL-ASM-NEXT: pacda x16, x17
122127
; ISEL-ASM-NEXT: str x16, [x0]
123128
; ISEL-ASM-NEXT: ret
124129
%dest.i = ptrtoint ptr %dest to i64
@@ -136,17 +141,15 @@ define void @store_signed_const_got(ptr %dest) {
136141
; ISEL-MIR-ELF-NEXT: %7:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got
137142
; ISEL-MIR-ELF-NEXT: %6:gpr64common = ADDXri %7, 8, 0
138143
; ISEL-MIR-ELF-NEXT: %2:gpr64common = MOVKXi %0, 1234, 48
139-
; ISEL-MIR-ELF-NEXT: %12:gpr64noip = COPY %2
140-
; ISEL-MIR-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
144+
; ISEL-MIR-ELF-NEXT: %12:gpr64noip = COPY %0
145+
; ISEL-MIR-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %12, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
141146
; ISEL-MIR-ELF-NEXT: %4:gpr64 = COPY $x16
142147
; ISEL-MIR-ELF-NEXT: %10:gpr64 = COPY %4
143148
; ISEL-MIR-ELF-NEXT: STRXui %10, %0, 0 :: (store (p0) into %ir.dest)
144149
; ISEL-MIR-ELF-NEXT: RET_ReallyLR
145150
;
146151
; ISEL-ASM-ELF-LABEL: store_signed_const_got:
147152
; ISEL-ASM-ELF-NEXT: .cfi_startproc
148-
; ISEL-ASM-ELF-NEXT: mov x8, x0
149-
; ISEL-ASM-ELF-NEXT: movk x8, #1234, lsl #48
150153
; ISEL-ASM-ELF-NEXT: adrp x17, :got_auth:const_table_got
151154
; ISEL-ASM-ELF-NEXT: add x17, x17, :got_auth_lo12:const_table_got
152155
; ISEL-ASM-ELF-NEXT: ldr x16, [x17]
@@ -158,7 +161,9 @@ define void @store_signed_const_got(ptr %dest) {
158161
; ISEL-ASM-ELF-NEXT: brk #0xc472
159162
; ISEL-ASM-ELF-NEXT: .Lauth_success_0:
160163
; ISEL-ASM-ELF-NEXT: add x16, x16, #8
161-
; ISEL-ASM-ELF-NEXT: pacda x16, x8
164+
; ISEL-ASM-ELF-NEXT: mov x17, x0
165+
; ISEL-ASM-ELF-NEXT: movk x17, #1234, lsl #48
166+
; ISEL-ASM-ELF-NEXT: pacda x16, x17
162167
; ISEL-ASM-ELF-NEXT: str x16, [x0]
163168
; ISEL-ASM-ELF-NEXT: ret
164169
%dest.i = ptrtoint ptr %dest to i64

‎llvm/test/CodeGen/AArch64/ptrauth-constant-in-code.ll‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -83,28 +83,33 @@ define ptr @foo() {
8383
@const_table_local = dso_local constant [3 x ptr] [ptr null, ptr null, ptr null]
8484
@const_table_got = constant [3 x ptr] [ptr null, ptr null, ptr null]
8585

86+
; Test that after post-processing in finalize-isel, MOVaddrPAC (or LOADgotPAC,
87+
; respectively) has both $AddrDisc and $Disc operands set. MOVaddr (or LOADgotAUTH,
88+
; respectively) and MOVKXi are not used anymore and are dead-code-eliminated
89+
; by the later passes.
90+
8691
define void @store_signed_const_local(ptr %dest) {
8792
; ISEL-MIR-LABEL: name: store_signed_const_local
8893
; ISEL-MIR: body:
8994
; ISEL-MIR: %0:gpr64common = COPY $x0
9095
; ISEL-MIR-NEXT: %1:gpr64common = MOVKXi %0, 1234, 48
9196
; ISEL-MIR-NEXT: %2:gpr64common = MOVaddr target-flags(aarch64-page) @const_table_local + 8, target-flags(aarch64-pageoff, aarch64-nc) @const_table_local + 8
92-
; ISEL-MIR-NEXT: %4:gpr64noip = COPY %1
93-
; ISEL-MIR-NEXT: MOVaddrPAC @const_table_local + 8, 2, %4, 0, implicit-def $x16, implicit-def $x17
97+
; ISEL-MIR-NEXT: %4:gpr64noip = COPY %0
98+
; ISEL-MIR-NEXT: MOVaddrPAC @const_table_local + 8, 2, %4, 1234, implicit-def $x16, implicit-def $x17
9499
; ISEL-MIR-NEXT: %3:gpr64 = COPY $x16
95100
; ISEL-MIR-NEXT: STRXui killed %3, %0, 0 :: (store (s64) into %ir.dest)
96101
; ISEL-MIR-NEXT: RET_ReallyLR
97102
;
98103
; ISEL-ASM-LABEL: store_signed_const_local:
99104
; ISEL-ASM-NEXT: .cfi_startproc
100-
; ISEL-ASM-NEXT: mov x8, x0
101-
; ISEL-ASM-NEXT: movk x8, #1234, lsl #48
102105
; ISEL-ASM-ELF-NEXT: adrp x16, const_table_local
103106
; ISEL-ASM-ELF-NEXT: add x16, x16, :lo12:const_table_local
104107
; ISEL-ASM-MACHO-NEXT: adrp x16, _const_table_local@PAGE
105108
; ISEL-ASM-MACHO-NEXT: add x16, x16, _const_table_local@PAGEOFF
106109
; ISEL-ASM-NEXT: add x16, x16, #8
107-
; ISEL-ASM-NEXT: pacda x16, x8
110+
; ISEL-ASM-NEXT: mov x17, x0
111+
; ISEL-ASM-NEXT: movk x17, #1234, lsl #48
112+
; ISEL-ASM-NEXT: pacda x16, x17
108113
; ISEL-ASM-NEXT: str x16, [x0]
109114
; ISEL-ASM-NEXT: ret
110115
%dest.i = ptrtoint ptr %dest to i64
@@ -122,16 +127,14 @@ define void @store_signed_const_got(ptr %dest) {
122127
; ISEL-MIR-ELF-NEXT: %1:gpr64common = MOVKXi %0, 1234, 48
123128
; ISEL-MIR-ELF-NEXT: %2:gpr64common = LOADgotAUTH target-flags(aarch64-got) @const_table_got, implicit-def dead $x16, implicit-def dead $x17, implicit-def dead $nzcv
124129
; ISEL-MIR-ELF-NEXT: %3:gpr64common = ADDXri killed %2, 8, 0
125-
; ISEL-MIR-ELF-NEXT: %5:gpr64noip = COPY %1
126-
; ISEL-MIR-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %5, 0, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
130+
; ISEL-MIR-ELF-NEXT: %5:gpr64noip = COPY %0
131+
; ISEL-MIR-ELF-NEXT: LOADgotPAC target-flags(aarch64-got) @const_table_got + 8, 2, %5, 1234, implicit-def $x16, implicit-def $x17, implicit-def $nzcv
127132
; ISEL-MIR-ELF-NEXT: %4:gpr64 = COPY $x16
128133
; ISEL-MIR-ELF-NEXT: STRXui killed %4, %0, 0 :: (store (s64) into %ir.dest)
129134
; ISEL-MIR-ELF-NEXT: RET_ReallyLR
130135
;
131136
; ISEL-ASM-ELF-LABEL: store_signed_const_got:
132137
; ISEL-ASM-ELF-NEXT: .cfi_startproc
133-
; ISEL-ASM-ELF-NEXT: mov x8, x0
134-
; ISEL-ASM-ELF-NEXT: movk x8, #1234, lsl #48
135138
; ISEL-ASM-ELF-NEXT: adrp x17, :got_auth:const_table_got
136139
; ISEL-ASM-ELF-NEXT: add x17, x17, :got_auth_lo12:const_table_got
137140
; ISEL-ASM-ELF-NEXT: ldr x16, [x17]
@@ -143,7 +146,9 @@ define void @store_signed_const_got(ptr %dest) {
143146
; ISEL-ASM-ELF-NEXT: brk #0xc472
144147
; ISEL-ASM-ELF-NEXT: .Lauth_success_0:
145148
; ISEL-ASM-ELF-NEXT: add x16, x16, #8
146-
; ISEL-ASM-ELF-NEXT: pacda x16, x8
149+
; ISEL-ASM-ELF-NEXT: mov x17, x0
150+
; ISEL-ASM-ELF-NEXT: movk x17, #1234, lsl #48
151+
; ISEL-ASM-ELF-NEXT: pacda x16, x17
147152
; ISEL-ASM-ELF-NEXT: str x16, [x0]
148153
; ISEL-ASM-ELF-NEXT: ret
149154
%dest.i = ptrtoint ptr %dest to i64

0 commit comments

Comments
 (0)