Skip to content

[MachineCP] Correctly handle register masks and sub-registers #122734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jan 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 91 additions & 55 deletions llvm/lib/CodeGen/MachineCopyPropagation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,32 @@ class CopyTracker {

DenseMap<MCRegUnit, CopyInfo> Copies;

// Memoised sets of register units which are preserved by each register mask,
// needed to efficiently remove copies which are invalidated by call
// instructions.
DenseMap<const uint32_t *, BitVector> RegMaskToPreservedRegUnits;

public:
/// Get the set of register units which are preserved by RegMaskOp.
BitVector &getPreservedRegUnits(const MachineOperand &RegMaskOp,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use LiveRegUnits instead? LRU should have featured similar functionalities.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking into the implementation of LiveRegUnits, it works in a similar way to this, but to re-use it we'd need to create a dummy LRU object for each call, which would probably have worse performance than this. The alternative would be to switch this whole pass over to LiveRegUnits, but that would be a lot of work, and I don't think it would be enough anyway because this pass tracks more than just liveness.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The alternative would be to switch this whole pass over to LiveRegUnits

Yeah, that's the direction we wanna take, so that Matt's concern in #122472 (comment) can be addressed too.

but that would be a lot of work, and I don't think it would be enough anyway because this pass tracks more than just liveness.

That's ok for now.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every pass tracking physical register liveness should be using LiveRegUnits. Way too many passes are doing their own hand rolled liveness tracking, and buggy

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the general principle, but I'm not sure that would help here, because LiveRegUnits only tracks liveness, but this pass also needs to track other information about the COPY instruction which defined that unit (if any). If that is possible, I don't think I'm qualified to do that re-write, since this is the first time I've looked at this pass. Are you OK with this fix as-is, or do you think that re-write needs to happen first?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm fine making progress here, but I'm stating as a general principle we should move every pass over to using LiveRegUnits

const TargetRegisterInfo &TRI) {
const uint32_t *RegMask = RegMaskOp.getRegMask();
auto Existing = RegMaskToPreservedRegUnits.find(RegMask);
if (Existing != RegMaskToPreservedRegUnits.end()) {
return Existing->second;
} else {
BitVector &PreservedRegUnits = RegMaskToPreservedRegUnits[RegMask];

PreservedRegUnits.resize(TRI.getNumRegUnits());
for (unsigned SafeReg = 0, E = TRI.getNumRegs(); SafeReg < E; ++SafeReg)
if (!RegMaskOp.clobbersPhysReg(SafeReg))
for (auto SafeUnit : TRI.regunits(SafeReg))
PreservedRegUnits.set(SafeUnit);

return PreservedRegUnits;
}
}

/// Mark all of the given registers and their subregisters as unavailable for
/// copying.
void markRegsUnavailable(ArrayRef<MCRegister> Regs,
Expand Down Expand Up @@ -164,64 +189,70 @@ class CopyTracker {
Copies.erase(Unit);
}

/// Clobber a single register, removing it from the tracker's copy maps.
void clobberRegister(MCRegister Reg, const TargetRegisterInfo &TRI,
const TargetInstrInfo &TII, bool UseCopyInstr) {
for (MCRegUnit Unit : TRI.regunits(Reg)) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about we keep this method and adding a filter BitVector to indicate which regunit should be really clobbered?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this way is clearer than adding an optional parameter which changes the behaviour of the existing function.

auto I = Copies.find(Unit);
if (I != Copies.end()) {
// When we clobber the source of a copy, we need to clobber everything
// it defined.
markRegsUnavailable(I->second.DefRegs, TRI);
// When we clobber the destination of a copy, we need to clobber the
// whole register it defined.
if (MachineInstr *MI = I->second.MI) {
std::optional<DestSourcePair> CopyOperands =
isCopyInstr(*MI, TII, UseCopyInstr);

MCRegister Def = CopyOperands->Destination->getReg().asMCReg();
MCRegister Src = CopyOperands->Source->getReg().asMCReg();

markRegsUnavailable(Def, TRI);

// Since we clobber the destination of a copy, the semantic of Src's
// "DefRegs" to contain Def is no longer effectual. We will also need
// to remove the record from the copy maps that indicates Src defined
// Def. Failing to do so might cause the target to miss some
// opportunities to further eliminate redundant copy instructions.
// Consider the following sequence during the
// ForwardCopyPropagateBlock procedure:
// L1: r0 = COPY r9 <- TrackMI
// L2: r0 = COPY r8 <- TrackMI (Remove r9 defined r0 from tracker)
// L3: use r0 <- Remove L2 from MaybeDeadCopies
// L4: early-clobber r9 <- Clobber r9 (L2 is still valid in tracker)
// L5: r0 = COPY r8 <- Remove NopCopy
for (MCRegUnit SrcUnit : TRI.regunits(Src)) {
auto SrcCopy = Copies.find(SrcUnit);
if (SrcCopy != Copies.end() && SrcCopy->second.LastSeenUseInCopy) {
// If SrcCopy defines multiple values, we only need
// to erase the record for Def in DefRegs.
for (auto itr = SrcCopy->second.DefRegs.begin();
itr != SrcCopy->second.DefRegs.end(); itr++) {
if (*itr == Def) {
SrcCopy->second.DefRegs.erase(itr);
// If DefReg becomes empty after removal, we can remove the
// SrcCopy from the tracker's copy maps. We only remove those
// entries solely record the Def is defined by Src. If an
// entry also contains the definition record of other Def'
// registers, it cannot be cleared.
if (SrcCopy->second.DefRegs.empty() && !SrcCopy->second.MI) {
Copies.erase(SrcCopy);
}
break;
/// Clobber a single register unit, removing it from the tracker's copy maps.
void clobberRegUnit(MCRegUnit Unit, const TargetRegisterInfo &TRI,
const TargetInstrInfo &TII, bool UseCopyInstr) {
auto I = Copies.find(Unit);
if (I != Copies.end()) {
// When we clobber the source of a copy, we need to clobber everything
// it defined.
markRegsUnavailable(I->second.DefRegs, TRI);
// When we clobber the destination of a copy, we need to clobber the
// whole register it defined.
if (MachineInstr *MI = I->second.MI) {
std::optional<DestSourcePair> CopyOperands =
isCopyInstr(*MI, TII, UseCopyInstr);

MCRegister Def = CopyOperands->Destination->getReg().asMCReg();
MCRegister Src = CopyOperands->Source->getReg().asMCReg();

markRegsUnavailable(Def, TRI);

// Since we clobber the destination of a copy, the semantic of Src's
// "DefRegs" to contain Def is no longer effectual. We will also need
// to remove the record from the copy maps that indicates Src defined
// Def. Failing to do so might cause the target to miss some
// opportunities to further eliminate redundant copy instructions.
// Consider the following sequence during the
// ForwardCopyPropagateBlock procedure:
// L1: r0 = COPY r9 <- TrackMI
// L2: r0 = COPY r8 <- TrackMI (Remove r9 defined r0 from tracker)
// L3: use r0 <- Remove L2 from MaybeDeadCopies
// L4: early-clobber r9 <- Clobber r9 (L2 is still valid in tracker)
// L5: r0 = COPY r8 <- Remove NopCopy
for (MCRegUnit SrcUnit : TRI.regunits(Src)) {
auto SrcCopy = Copies.find(SrcUnit);
if (SrcCopy != Copies.end() && SrcCopy->second.LastSeenUseInCopy) {
// If SrcCopy defines multiple values, we only need
// to erase the record for Def in DefRegs.
for (auto itr = SrcCopy->second.DefRegs.begin();
itr != SrcCopy->second.DefRegs.end(); itr++) {
if (*itr == Def) {
SrcCopy->second.DefRegs.erase(itr);
// If DefReg becomes empty after removal, we can remove the
// SrcCopy from the tracker's copy maps. We only remove those
// entries solely record the Def is defined by Src. If an
// entry also contains the definition record of other Def'
// registers, it cannot be cleared.
if (SrcCopy->second.DefRegs.empty() && !SrcCopy->second.MI) {
Copies.erase(SrcCopy);
}
break;
}
}
}
}
// Now we can erase the copy.
Copies.erase(I);
}
// Now we can erase the copy.
Copies.erase(I);
}
}

/// Clobber a single register, removing it from the tracker's copy maps.
void clobberRegister(MCRegister Reg, const TargetRegisterInfo &TRI,
const TargetInstrInfo &TII, bool UseCopyInstr) {
for (MCRegUnit Unit : TRI.regunits(Reg)) {
clobberRegUnit(Unit, TRI, TII, UseCopyInstr);
}
}

Expand Down Expand Up @@ -960,6 +991,9 @@ void MachineCopyPropagation::ForwardCopyPropagateBlock(MachineBasicBlock &MBB) {
// a large set of registers. Treat clobbered registers the same way as
// defined registers.
if (RegMask) {
BitVector &PreservedRegUnits =
Tracker.getPreservedRegUnits(*RegMask, *TRI);

// Erase any MaybeDeadCopies whose destination register is clobbered.
for (SmallSetVector<MachineInstr *, 8>::iterator DI =
MaybeDeadCopies.begin();
Expand All @@ -978,9 +1012,11 @@ void MachineCopyPropagation::ForwardCopyPropagateBlock(MachineBasicBlock &MBB) {
LLVM_DEBUG(dbgs() << "MCP: Removing copy due to regmask clobbering: ";
MaybeDead->dump());

// Make sure we invalidate any entries in the copy maps before erasing
// the instruction.
Tracker.clobberRegister(Reg, *TRI, *TII, UseCopyInstr);
// Invalidate all entries in the copy map which are not preserved by
// this register mask.
for (unsigned RegUnit : TRI->regunits(Reg))
if (!PreservedRegUnits.test(RegUnit))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this might cause an invalid pointer problem. Some of the RegUnit in CopyInfo is still referencing the MaybeDead using pointer. Then after we call MaybeDead->eraseFromParent(); in following 2 lines, the MI pointer might become invalidated , we will then run into problem in next iteration

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll take a look at this, do you have a test case which triggers it?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We do, but unfortunately the MIR contains some register class that are downstream only, so we can't provide it now.
The following diff should fix it; however, I don't have reduced MIR available now.

diff --git a/llvm/lib/CodeGen/MachineCopyPropagation.cpp b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
index d44b064dcb4b..70dd7d04a153 100644
--- a/llvm/lib/CodeGen/MachineCopyPropagation.cpp
+++ b/llvm/lib/CodeGen/MachineCopyPropagation.cpp
@@ -1009,18 +1009,30 @@ void MachineCopyPropagation::ForwardCopyPropagateBlock(MachineBasicBlock &MBB) {
           continue;
         }
 
-        LLVM_DEBUG(dbgs() << "MCP: Removing copy due to regmask clobbering: ";
-                   MaybeDead->dump());
-
         // Invalidate all entries in the copy map which are not preserved by
         // this register mask.
-        for (unsigned RegUnit : TRI->regunits(Reg))
+        bool MIRefedinCopyInfo = false;
+        for (unsigned RegUnit : TRI->regunits(Reg)) {
           if (!PreservedRegUnits.test(RegUnit))
             Tracker.clobberRegUnit(RegUnit, *TRI, *TII, UseCopyInstr);
+          else {
+            if (MaybeDead == Tracker.findCopyForUnit(RegUnit, *TRI)) {
+              MIRefedinCopyInfo = true;
+            }
+          }
+        }
 
         // erase() will return the next valid iterator pointing to the next
         // element after the erased one.
         DI = MaybeDeadCopies.erase(DI);
+
+        // Preserved by RegMask, DO NOT remove copy
+        if (MIRefedinCopyInfo)
+          continue;
+
+        LLVM_DEBUG(dbgs() << "MCP: Removing copy due to regmask clobbering: ";
+                   MaybeDead->dump());
+
         MaybeDead->eraseFromParent();
         Changed = true;
         ++NumDeletes;

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great, could you send that as a PR? I think it should be possible to hit this with the ARM floating-point register classes, but I've not managed to work out the details yet.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tracker.clobberRegUnit(RegUnit, *TRI, *TII, UseCopyInstr);

// erase() will return the next valid iterator pointing to the next
// element after the erased one.
Expand Down
30 changes: 30 additions & 0 deletions llvm/test/CodeGen/AArch64/machine-cp-sub-reg.mir
Original file line number Diff line number Diff line change
@@ -1,6 +1,17 @@
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
# RUN: llc -o - %s --run-pass=machine-cp -mcp-use-is-copy-instr -mtriple=arm64-apple-macos --verify-machineinstrs | FileCheck %s

--- |
declare void @foo()

define void @test() {
unreachable
}
define void @test2() {
unreachable
}
...

---
name: test
tracksRegLiveness: true
Expand Down Expand Up @@ -30,3 +41,22 @@ body: |

RET undef $lr, implicit $x0
...
---
name: test2
tracksRegLiveness: true
body: |
bb.0:
liveins: $q14, $d29, $x0, $x1
; CHECK-LABEL: name: test2
; CHECK: liveins: $q14, $d29, $x0, $x1
; CHECK-NEXT: {{ $}}
; CHECK-NEXT: renamable $d8 = COPY killed renamable $d29
; CHECK-NEXT: BL @foo, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
; CHECK-NEXT: renamable $b0 = SMAXVv8i8v killed renamable $d8, implicit-def $q0
; CHECK-NEXT: RET_ReallyLR implicit $b0
renamable $q8 = COPY renamable $q14
renamable $d8 = COPY killed renamable $d29
BL @foo, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
renamable $b0 = SMAXVv8i8v killed renamable $d8, implicit-def $q0
RET_ReallyLR implicit $b0
...
Loading