Skip to content

Commit 5e3a6e4

Browse files
atrosinenkoatrosinenko
committed
Prevent immediate modifier substitution
1 parent 936ded1 commit 5e3a6e4

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
@@ -3069,62 +3069,80 @@ AArch64TargetLowering::tryRewritingPAC(MachineInstr &MI,
30693069
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();
30703070
const DebugLoc &DL = MI.getDebugLoc();
30713071

3072-
MachineInstr *AddrInstr = nullptr;
3073-
int64_t Offset = 0;
3074-
// Try to find the address-setting instruction, accumulating the offset
3075-
// along the way. If any unknown pattern is found, keep everything as-is.
3076-
MachineOperand *CurOp = &MI.getOperand(1);
3077-
while (CurOp) {
3078-
MachineOperand *Def = MRI.getOneDef(CurOp->getReg());
3079-
if (!Def)
3080-
return BB;
3081-
MachineInstr *DefMI = Def->getParent();
3082-
assert(DefMI != nullptr);
3083-
3084-
switch (DefMI->getOpcode()) {
3085-
case AArch64::COPY:
3086-
CurOp = &DefMI->getOperand(1);
3087-
break;
3088-
case AArch64::ADDXri:
3089-
if (DefMI->getOperand(3).getImm() != 0) // shifts are not handled
3090-
return BB;
3091-
CurOp = &DefMI->getOperand(1);
3092-
Offset += DefMI->getOperand(2).getImm();
3093-
break;
3094-
case AArch64::MOVaddr:
3095-
case AArch64::LOADgotAUTH:
3096-
AddrInstr = DefMI;
3097-
CurOp = nullptr;
3098-
break;
3099-
default:
3100-
return BB;
3072+
// Find the unique register definition, skipping copies.
3073+
auto GetUniqueDef = [&MRI](Register Reg) {
3074+
for (;;) {
3075+
MachineInstr *Def = MRI.getUniqueVRegDef(Reg);
3076+
if (!Def || Def->getOpcode() != AArch64::COPY)
3077+
return Def;
3078+
3079+
Reg = Def->getOperand(1).getReg();
31013080
}
3102-
}
3081+
};
3082+
// Find the unique register definition, skipping copies and increments.
3083+
auto GetUniqueDefPlusOffset =
3084+
[GetUniqueDef](Register Reg, int64_t &Offset) -> MachineInstr * {
3085+
for (;;) {
3086+
MachineInstr *Def = GetUniqueDef(Reg);
3087+
if (!Def || Def->getOpcode() != AArch64::ADDXri)
3088+
return Def;
3089+
3090+
if (Def->getOperand(3).getImm() != 0)
3091+
return nullptr; // shifted immediates are not handled
3092+
Reg = Def->getOperand(1).getReg();
3093+
Offset += Def->getOperand(2).getImm();
3094+
}
3095+
};
3096+
3097+
// Try to find a known address-setting instruction, accumulating the offset
3098+
// along the way. If no known pattern is found, keep everything as-is.
3099+
3100+
int64_t AddrOffset = 0;
3101+
MachineInstr *AddrDefInstr =
3102+
GetUniqueDefPlusOffset(MI.getOperand(1).getReg(), AddrOffset);
3103+
if (!AddrDefInstr)
3104+
return BB;
31033105

3104-
unsigned NewOpcode = AddrInstr->getOpcode() == AArch64::LOADgotAUTH
3105-
? AArch64::LOADgotPAC
3106-
: AArch64::MOVaddrPAC;
3107-
MachineOperand &AddrOp = AddrInstr->getOperand(1);
3106+
unsigned NewOpcode;
3107+
if (AddrDefInstr->getOpcode() == AArch64::LOADgotAUTH)
3108+
NewOpcode = AArch64::LOADgotPAC;
3109+
else if (AddrDefInstr->getOpcode() == AArch64::MOVaddr)
3110+
NewOpcode = AArch64::MOVaddrPAC;
3111+
else
3112+
return BB; // Unknown opcode.
3113+
3114+
MachineOperand &AddrOp = AddrDefInstr->getOperand(1);
31083115
unsigned TargetFlags = AddrOp.getTargetFlags() & ~AArch64II::MO_PAGE;
3109-
Offset += AddrOp.getOffset();
3116+
const GlobalValue *GV = AddrOp.getGlobal();
3117+
AddrOffset += AddrOp.getOffset();
3118+
3119+
// Detect discriminator blend computation, if any.
3120+
Register RegDisc = isPACWithZeroDisc(MI.getOpcode())
3121+
? AArch64::XZR
3122+
: MI.getOperand(2).getReg();
3123+
unsigned IntDisc = 0;
3124+
MachineInstr *MaybeBlendDef =
3125+
RegDisc == AArch64::XZR ? nullptr : GetUniqueDef(RegDisc);
3126+
if (MaybeBlendDef && MaybeBlendDef->getOpcode() == AArch64::MOVKXi &&
3127+
MaybeBlendDef->getOperand(3).getImm() == 48) {
3128+
RegDisc = MaybeBlendDef->getOperand(1).getReg();
3129+
IntDisc = MaybeBlendDef->getOperand(2).getImm();
3130+
}
31103131

31113132
// MOVaddrPAC and LOADgotPAC pseudos are expanded so that they use X16/X17
31123133
// internally, thus their restrictions on the register class of $AddrDisc
31133134
// operand are stricter than those of real PAC* instructions.
3114-
// If the original instruction accepts a discriminator operand, make sure
3115-
// it is moved out of X16/X17.
3116-
Register DiscReg = AArch64::XZR;
3117-
if (!isPACWithZeroDisc(MI.getOpcode())) {
3118-
DiscReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3119-
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), DiscReg)
3120-
.addReg(MI.getOperand(2).getReg());
3135+
if (RegDisc != AArch64::XZR) {
3136+
Register TmpReg = MRI.createVirtualRegister(&AArch64::GPR64noipRegClass);
3137+
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), TmpReg).addReg(RegDisc);
3138+
RegDisc = TmpReg;
31213139
}
31223140

31233141
BuildMI(*BB, MI, DL, TII->get(NewOpcode))
3124-
.addGlobalAddress(AddrOp.getGlobal(), Offset, TargetFlags)
3142+
.addGlobalAddress(GV, AddrOffset, TargetFlags)
31253143
.addImm(getKeyForPACOpcode(MI.getOpcode()))
3126-
.addReg(DiscReg)
3127-
.addImm(0);
3144+
.addReg(RegDisc)
3145+
.addImm(IntDisc);
31283146

31293147
BuildMI(*BB, MI, DL, TII->get(AArch64::COPY), MI.getOperand(0).getReg())
31303148
.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)