Skip to content

Commit e2a071e

Browse files
ostannardostannard
authored
[MachineCP] Correctly handle register masks and sub-registers (#122472)
When passing an instruction with a register mask, the machine copy propagation pass was dropping the information about some copy instructions which define a register which is preserved by the mask, because that register overlaps a register which is partially clobbered by it. This resulted in a miscompilation for AArch64, because this caused a live copy to be considered dead. The fix is to clobber register masks by finding the set of reg units which is preserved by the mask, and clobbering all units not in that set.
1 parent d03f35f commit e2a071e

File tree

2 files changed

+111
-57
lines changed

2 files changed

+111
-57
lines changed

‎llvm/lib/CodeGen/MachineCopyPropagation.cpp

Lines changed: 80 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -164,67 +164,91 @@ class CopyTracker {
164164
Copies.erase(Unit);
165165
}
166166

167-
/// Clobber a single register, removing it from the tracker's copy maps.
168-
void clobberRegister(MCRegister Reg, const TargetRegisterInfo &TRI,
169-
const TargetInstrInfo &TII, bool UseCopyInstr) {
170-
for (MCRegUnit Unit : TRI.regunits(Reg)) {
171-
auto I = Copies.find(Unit);
172-
if (I != Copies.end()) {
173-
// When we clobber the source of a copy, we need to clobber everything
174-
// it defined.
175-
markRegsUnavailable(I->second.DefRegs, TRI);
176-
// When we clobber the destination of a copy, we need to clobber the
177-
// whole register it defined.
178-
if (MachineInstr *MI = I->second.MI) {
179-
std::optional<DestSourcePair> CopyOperands =
180-
isCopyInstr(*MI, TII, UseCopyInstr);
181-
182-
MCRegister Def = CopyOperands->Destination->getReg().asMCReg();
183-
MCRegister Src = CopyOperands->Source->getReg().asMCReg();
184-
185-
markRegsUnavailable(Def, TRI);
186-
187-
// Since we clobber the destination of a copy, the semantic of Src's
188-
// "DefRegs" to contain Def is no longer effectual. We will also need
189-
// to remove the record from the copy maps that indicates Src defined
190-
// Def. Failing to do so might cause the target to miss some
191-
// opportunities to further eliminate redundant copy instructions.
192-
// Consider the following sequence during the
193-
// ForwardCopyPropagateBlock procedure:
194-
// L1: r0 = COPY r9 <- TrackMI
195-
// L2: r0 = COPY r8 <- TrackMI (Remove r9 defined r0 from tracker)
196-
// L3: use r0 <- Remove L2 from MaybeDeadCopies
197-
// L4: early-clobber r9 <- Clobber r9 (L2 is still valid in tracker)
198-
// L5: r0 = COPY r8 <- Remove NopCopy
199-
for (MCRegUnit SrcUnit : TRI.regunits(Src)) {
200-
auto SrcCopy = Copies.find(SrcUnit);
201-
if (SrcCopy != Copies.end() && SrcCopy->second.LastSeenUseInCopy) {
202-
// If SrcCopy defines multiple values, we only need
203-
// to erase the record for Def in DefRegs.
204-
for (auto itr = SrcCopy->second.DefRegs.begin();
205-
itr != SrcCopy->second.DefRegs.end(); itr++) {
206-
if (*itr == Def) {
207-
SrcCopy->second.DefRegs.erase(itr);
208-
// If DefReg becomes empty after removal, we can remove the
209-
// SrcCopy from the tracker's copy maps. We only remove those
210-
// entries solely record the Def is defined by Src. If an
211-
// entry also contains the definition record of other Def'
212-
// registers, it cannot be cleared.
213-
if (SrcCopy->second.DefRegs.empty() && !SrcCopy->second.MI) {
214-
Copies.erase(SrcCopy);
215-
}
216-
break;
167+
/// Clobber a single register unit, removing it from the tracker's copy maps.
168+
void clobberRegUnit(MCRegUnit Unit, const TargetRegisterInfo &TRI,
169+
const TargetInstrInfo &TII, bool UseCopyInstr) {
170+
auto I = Copies.find(Unit);
171+
if (I != Copies.end()) {
172+
// When we clobber the source of a copy, we need to clobber everything
173+
// it defined.
174+
markRegsUnavailable(I->second.DefRegs, TRI);
175+
// When we clobber the destination of a copy, we need to clobber the
176+
// whole register it defined.
177+
if (MachineInstr *MI = I->second.MI) {
178+
std::optional<DestSourcePair> CopyOperands =
179+
isCopyInstr(*MI, TII, UseCopyInstr);
180+
181+
MCRegister Def = CopyOperands->Destination->getReg().asMCReg();
182+
MCRegister Src = CopyOperands->Source->getReg().asMCReg();
183+
184+
markRegsUnavailable(Def, TRI);
185+
186+
// Since we clobber the destination of a copy, the semantic of Src's
187+
// "DefRegs" to contain Def is no longer effectual. We will also need
188+
// to remove the record from the copy maps that indicates Src defined
189+
// Def. Failing to do so might cause the target to miss some
190+
// opportunities to further eliminate redundant copy instructions.
191+
// Consider the following sequence during the
192+
// ForwardCopyPropagateBlock procedure:
193+
// L1: r0 = COPY r9 <- TrackMI
194+
// L2: r0 = COPY r8 <- TrackMI (Remove r9 defined r0 from tracker)
195+
// L3: use r0 <- Remove L2 from MaybeDeadCopies
196+
// L4: early-clobber r9 <- Clobber r9 (L2 is still valid in tracker)
197+
// L5: r0 = COPY r8 <- Remove NopCopy
198+
for (MCRegUnit SrcUnit : TRI.regunits(Src)) {
199+
auto SrcCopy = Copies.find(SrcUnit);
200+
if (SrcCopy != Copies.end() && SrcCopy->second.LastSeenUseInCopy) {
201+
// If SrcCopy defines multiple values, we only need
202+
// to erase the record for Def in DefRegs.
203+
for (auto itr = SrcCopy->second.DefRegs.begin();
204+
itr != SrcCopy->second.DefRegs.end(); itr++) {
205+
if (*itr == Def) {
206+
SrcCopy->second.DefRegs.erase(itr);
207+
// If DefReg becomes empty after removal, we can remove the
208+
// SrcCopy from the tracker's copy maps. We only remove those
209+
// entries solely record the Def is defined by Src. If an
210+
// entry also contains the definition record of other Def'
211+
// registers, it cannot be cleared.
212+
if (SrcCopy->second.DefRegs.empty() && !SrcCopy->second.MI) {
213+
Copies.erase(SrcCopy);
217214
}
215+
break;
218216
}
219217
}
220218
}
221219
}
222-
// Now we can erase the copy.
223-
Copies.erase(I);
224220
}
221+
// Now we can erase the copy.
222+
Copies.erase(I);
225223
}
226224
}
227225

226+
/// Clobber a single register, removing it from the tracker's copy maps.
227+
void clobberRegister(MCRegister Reg, const TargetRegisterInfo &TRI,
228+
const TargetInstrInfo &TII, bool UseCopyInstr) {
229+
for (MCRegUnit Unit : TRI.regunits(Reg)) {
230+
clobberRegUnit(Unit, TRI, TII, UseCopyInstr);
231+
}
232+
}
233+
234+
/// Clobber all registers which are not preserved by RegMask, removing them
235+
/// from the tracker's copy maps.
236+
void clobberRegistersExceptMask(const MachineOperand *RegMask,
237+
const TargetRegisterInfo &TRI,
238+
const TargetInstrInfo &TII,
239+
bool UseCopyInstr) {
240+
BitVector SafeRegUnits(TRI.getNumRegUnits());
241+
242+
for (unsigned SafeReg = 0, E = TRI.getNumRegs(); SafeReg < E; ++SafeReg)
243+
if (!RegMask->clobbersPhysReg(SafeReg))
244+
for (auto SafeUnit : TRI.regunits(SafeReg))
245+
SafeRegUnits.set(SafeUnit);
246+
247+
for (unsigned Unit = 0, E = TRI.getNumRegUnits(); Unit < E; ++Unit)
248+
if (!SafeRegUnits.test(Unit))
249+
clobberRegUnit(Unit, TRI, TII, UseCopyInstr);
250+
}
251+
228252
/// Track copy's src users, and return false if that can't be done.
229253
/// We can only track if we have a COPY instruction which source is
230254
/// the same as the Reg.
@@ -960,6 +984,10 @@ void MachineCopyPropagation::ForwardCopyPropagateBlock(MachineBasicBlock &MBB) {
960984
// a large set of registers. Treat clobbered registers the same way as
961985
// defined registers.
962986
if (RegMask) {
987+
// Invalidate all entries in the copy map which are not preserved by this
988+
// register mask.
989+
Tracker.clobberRegistersExceptMask(RegMask, *TRI, *TII, UseCopyInstr);
990+
963991
// Erase any MaybeDeadCopies whose destination register is clobbered.
964992
for (SmallSetVector<MachineInstr *, 8>::iterator DI =
965993
MaybeDeadCopies.begin();
@@ -978,10 +1006,6 @@ void MachineCopyPropagation::ForwardCopyPropagateBlock(MachineBasicBlock &MBB) {
9781006
LLVM_DEBUG(dbgs() << "MCP: Removing copy due to regmask clobbering: ";
9791007
MaybeDead->dump());
9801008

981-
// Make sure we invalidate any entries in the copy maps before erasing
982-
// the instruction.
983-
Tracker.clobberRegister(Reg, *TRI, *TII, UseCopyInstr);
984-
9851009
// erase() will return the next valid iterator pointing to the next
9861010
// element after the erased one.
9871011
DI = MaybeDeadCopies.erase(DI);

‎llvm/test/CodeGen/AArch64/machine-cp-sub-reg.mir

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,16 @@
11
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
2-
# RUN: llc -o - %s --run-pass=machine-cp -mcp-use-is-copy-instr -mtriple=arm64-apple-macos --verify-machineinstrs | FileCheck %s
2+
# RUN: llc -o - %s --run-pass=machine-cp -mcp-use-is-copy-instr -mtriple=arm64-apple-macos | FileCheck %s
3+
4+
--- |
5+
declare void @foo()
6+
7+
define void @test() {
8+
unreachable
9+
}
10+
define void @test2() {
11+
unreachable
12+
}
13+
...
314

415
---
516
name: test
@@ -30,3 +41,22 @@ body: |
3041
3142
RET undef $lr, implicit $x0
3243
...
44+
---
45+
name: test2
46+
tracksRegLiveness: true
47+
body: |
48+
bb.0:
49+
liveins: $q14, $d29, $x0, $x1
50+
; CHECK-LABEL: name: test2
51+
; CHECK: liveins: $q14, $d29, $x0, $x1
52+
; CHECK-NEXT: {{ $}}
53+
; CHECK-NEXT: renamable $d8 = COPY killed renamable $d29
54+
; CHECK-NEXT: BL @foo, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
55+
; CHECK-NEXT: renamable $b0 = SMAXVv8i8v killed renamable $d8, implicit-def $q0
56+
; CHECK-NEXT: RET_ReallyLR implicit $b0
57+
renamable $q8 = COPY renamable $q14
58+
renamable $d8 = COPY killed renamable $d29
59+
BL @foo, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
60+
renamable $b0 = SMAXVv8i8v killed renamable $d8, implicit-def $q0
61+
RET_ReallyLR implicit $b0
62+
...

0 commit comments

Comments
 (0)