Skip to content

[AIEX] Extend Staged 2D/3D regalloc to avoid spills #685

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.

Sign up for GitHub

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

Open
wants to merge 9 commits into
base: aie-public
Choose a base branch
from
+1,534 −253
Open

[AIEX] Extend Staged 2D/3D regalloc to avoid spills #685

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
afdc992
[AIE2P] Extend allocation filter to include only Target MIs
andcarminati Oct 9, 2025
12da76d
[AIEX] Fix LIs in SuperRegRewriter that are related to unallocated re…
andcarminati Oct 10, 2025
23c18b4
[AIEX][NFC] Refator some AIESuperReg utilities
andcarminati Oct 9, 2025
264a491
[AIE2P] Add a base test to detect missing undef flags in copy expansion.
andcarminati Oct 30, 2025
9975813
[AIEX] Add missing undef flags in copy expansion
andcarminati Oct 31, 2025
1968bf5
[AIE2P] Add a base lit test with unallocated 2D/3D regs before main G…
andcarminati Oct 21, 2025
3527e24
[AIEX] Add a Pass to expand unallocated 2D/3D into individual ones
andcarminati Oct 10, 2025
fbaf155
[AIE2P] Add base tests exposing cycles in copy bundles
andcarminati Oct 14, 2025
925f9d9
[AIEX] Expand copy bundles for unallocated dest VRegs
andcarminati Oct 13, 2025
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
32 changes: 29 additions & 3 deletions llvm/lib/Target/AIE/AIESuperRegRewriter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
//
//===----------------------------------------------------------------------===//

Expand Down Expand Up @@ -153,6 +153,7 @@ bool AIESuperRegRewriter::runOnMachineFunction(MachineFunction &MF) {

// Collect already-assigned VRegs that can be split into smaller ones.
LLVM_DEBUG(VRM.dump());
LLVM_DEBUG(LIS.dump());
for (unsigned VRegIdx = 0, End = MRI.getNumVirtRegs(); VRegIdx != End;
++VRegIdx) {
Register Reg = Register::index2VirtReg(VRegIdx);
Expand Down Expand Up @@ -208,7 +209,8 @@ static LaneBitmask getLiveLanesAt(SlotIndex Index, Register Reg,
/// Rewrite a full copy into multiple copies using the subregs in \p CopySubRegs
static void rewriteFullCopy(MachineInstr &MI, const std::set<int> &CopySubRegs,
LiveIntervals &LIS, const TargetInstrInfo &TII,
const TargetRegisterInfo &TRI) {
const TargetRegisterInfo &TRI, VirtRegMap &VRM,
LiveRegMatrix &LRM) {
assert(MI.isFullCopy());
SlotIndex CopyIndex = LIS.getInstructionIndex(MI);
LLVM_DEBUG(dbgs() << " Changing full copy at " << CopyIndex << ": " << MI);
Expand All @@ -217,6 +219,8 @@ static void rewriteFullCopy(MachineInstr &MI, const std::set<int> &CopySubRegs,
LaneBitmask LiveSrcLanes = getLiveLanesAt(CopyIndex, SrcReg, LIS);

LIS.removeVRegDefAt(LIS.getInterval(DstReg), CopyIndex.getRegSlot());

SmallSet<Register, 8> RegistersToRepair;
for (int SubRegIdx : CopySubRegs) {
if ((LiveSrcLanes & TRI.getSubRegIndexLaneMask(SubRegIdx)).none()) {
LLVM_DEBUG(dbgs() << " Skip undef subreg "
Expand All @@ -232,10 +236,32 @@ static void rewriteFullCopy(MachineInstr &MI, const std::set<int> &CopySubRegs,
LLVM_DEBUG(dbgs() << " to " << *PartCopy);
LIS.InsertMachineInstrInMaps(*PartCopy);
LIS.getInterval(PartCopy->getOperand(0).getReg());
RegistersToRepair.insert(PartCopy->getOperand(1).getReg());
Copy link
Collaborator Author

@andcarminati andcarminati Nov 17, 2025

Choose a reason for hiding this comment

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

Comment about repair (only r)

}

LIS.RemoveMachineInstrFromMaps(MI);
MI.eraseFromParent();
// As we don't handle all registers now (selective LI filter),
// We should make sure that all LiveIntervals are correct.
// If we dont't repair, MI will compose the LIs of some registers,
Copy link
Collaborator

@F-Stuckmann F-Stuckmann Nov 3, 2025

Choose a reason for hiding this comment

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

nit: don't

Copy link
Collaborator

@F-Stuckmann F-Stuckmann Nov 3, 2025

Choose a reason for hiding this comment

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

What do you mean with compose? the dead MI will block some LI ?

Copy link
Collaborator Author

@andcarminati andcarminati Nov 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

A dead MI will have an invalid SlotIndex that will appear in some instruction`s LI, what is wrong. The original copy should not appear in any LI.

// what is not correct because MI was deleted.
for (Register R : RegistersToRepair) {

if (!LIS.hasInterval(R))
continue;

if (VRM.hasPhys(R)) {
const MCRegister PhysReg = VRM.getPhys(R);
const LiveInterval &OldLI = LIS.getInterval(R);
LRM.unassign(OldLI);
LIS.removeInterval(R);
const LiveInterval &LI = LIS.createAndComputeVirtRegInterval(R);
LRM.assign(LI, PhysReg);
} else {
LIS.removeInterval(R);
LIS.createAndComputeVirtRegInterval(R);
}
}
}

void AIESuperRegRewriter::rewriteSuperReg(
Expand All @@ -260,7 +286,7 @@ void AIESuperRegRewriter::rewriteSuperReg(
for (MachineInstr &MI : make_early_inc_range(MRI.reg_instructions(Reg))) {
if (MI.isFullCopy())
rewriteFullCopy(MI, TRI.getSubRegSplit(MRI.getRegClass(Reg)->getID()),
LIS, *TII, TRI);
LIS, *TII, TRI, VRM, LRM);
}

LLVM_DEBUG(dbgs() << " Splitting range " << LIS.getInterval(Reg) << "\n");
Expand Down
16 changes: 8 additions & 8 deletions llvm/test/CodeGen/AIE/aie2p/ra/staged-ra-spill.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
# (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates

# RUN: llc -O2 -mtriple=aie2p -verify-machineinstrs --aie-staged-ra -start-before=greedy -stop-after=virtregrewriter %s -o - \
# RUN: | FileCheck %s --check-prefix=RA
Expand All @@ -24,23 +24,23 @@ body: |
; RA-NEXT: renamable $dn0 = LDA_dms_lda_idx_imm renamable $p1, 0
; RA-NEXT: renamable $m0 = LDA_dms_lda_idx_imm renamable $p1, 4
; RA-NEXT: renamable $dj0 = LDA_dms_lda_idx_imm renamable $p1, 8
; RA-NEXT: ST_D_SPILL renamable $d0, %stack.0, implicit $sp :: (store (s128) into %stack.0, align 4)
; RA-NEXT: ST_D_SPILL renamable $d0, %stack.1, implicit $sp :: (store (s128) into %stack.1, align 4)
; RA-NEXT: renamable $dj0 = MOV_PD_imm11_pseudo 12
; RA-NEXT: renamable $r0 = LDA_dms_lda_idx renamable $p1, killed renamable $dj0
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.0, implicit $sp :: (load (s128) from %stack.0, align 4)
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.1, implicit $sp :: (load (s128) from %stack.1, align 4)
; RA-NEXT: renamable $dc0 = COPY killed renamable $r0
; RA-NEXT: ST_D_SPILL killed renamable $d0, %stack.0, implicit $sp :: (store (s128) into %stack.0, align 4)
; RA-NEXT: ST_D_SPILL killed renamable $d0, %stack.1, implicit $sp :: (store (s128) into %stack.1, align 4)
; RA-NEXT: renamable $dn0 = LDA_dms_lda_idx_imm renamable $p1, 16
; RA-NEXT: renamable $m0 = LDA_dms_lda_idx_imm renamable $p1, 20
; RA-NEXT: renamable $dj0 = LDA_dms_lda_idx_imm renamable $p1, 24
; RA-NEXT: renamable $dc0 = LDA_dms_lda_idx_imm killed renamable $p1, 28
; RA-NEXT: ST_D_SPILL killed renamable $d0, %stack.1, implicit $sp :: (store (s128) into %stack.1, align 4)
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.0, implicit $sp :: (load (s128) from %stack.0, align 4)
; RA-NEXT: $p0, $dc0 = PADDA_2D_split killed $p0, $m0, $dn0, $dj0, $dc0
; RA-NEXT: ST_D_SPILL killed renamable $d0, %stack.0, implicit $sp :: (store (s128) into %stack.0, align 4)
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.1, implicit $sp :: (load (s128) from %stack.1, align 4)
; RA-NEXT: $p0, dead $dc0 = PADDA_2D_split killed $p0, $m0, $dn0, $dj0, $dc0
; RA-NEXT: $p0, $dc0 = PADDA_2D_split killed $p0, killed $m0, killed $dn0, killed $dj0, killed $dc0
; RA-NEXT: ST_D_SPILL renamable $d0, %stack.1, implicit $sp :: (store (s128) into %stack.1, align 4)
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.0, implicit $sp :: (load (s128) from %stack.0, align 4)
; RA-NEXT: $p0, dead $dc0 = PADDA_2D_split killed $p0, killed $m0, killed $dn0, killed $dj0, killed $dc0
; RA-NEXT: renamable $d0 = LDA_D_SPILL %stack.1, implicit $sp :: (load (s128) from %stack.1, align 4)
; RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc0, implicit $d1, implicit $d2, implicit $d3, implicit $d4, implicit $d5, implicit $d6, implicit $d7
%20:ep = COPY $p0
%21:ep = COPY $p1
Expand Down
27 changes: 16 additions & 11 deletions llvm/test/CodeGen/AIE/staged-ra-rewrite.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
# (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates

# RUN: llc -O2 -mtriple=aie2 -verify-machineinstrs -run-pass=greedy,aie-superreg-rewrite %s -o - | FileCheck %s --check-prefix=AIE2-VREGS
# RUN: llc -O2 -mtriple=aie2 -verify-machineinstrs --aie-staged-ra -start-before=greedy -stop-after=virtregrewriter %s -o - \
Expand Down Expand Up @@ -291,15 +291,16 @@ body: |
; AIE2P-RA-NEXT: renamable $r0 = LDA_dms_lda_idx_imm renamable $p1, 4
; AIE2P-RA-NEXT: renamable $r1 = LDA_dms_lda_idx_imm renamable $p1, 8
; AIE2P-RA-NEXT: renamable $r2 = LDA_dms_lda_idx_imm killed renamable $p1, 12
; AIE2P-RA-NEXT: renamable $dn1 = COPY killed renamable $r0
; AIE2P-RA-NEXT: renamable $dj1 = COPY killed renamable $r1
; AIE2P-RA-NEXT: renamable $dc1 = COPY killed renamable $r2
; AIE2P-RA-NEXT: renamable $dn0 = COPY killed renamable $r0
; AIE2P-RA-NEXT: renamable $m0 = COPY killed renamable $m1
; AIE2P-RA-NEXT: renamable $dj0 = COPY killed renamable $r1
; AIE2P-RA-NEXT: renamable $dc0 = COPY killed renamable $r2
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: bb.1:
; AIE2P-RA-NEXT: liveins: $dc1, $dj1, $dn1, $m1, $p0
; AIE2P-RA-NEXT: liveins: $dc0, $dj0, $dn0, $m0, $p0
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: $p0, $dc1 = PADDA_2D_split killed $p0, killed $m1, killed $dn1, killed $dj1, killed $dc1
; AIE2P-RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc1
; AIE2P-RA-NEXT: $p0, $dc0 = PADDA_2D_split killed $p0, killed $m0, killed $dn0, killed $dj0, killed $dc0
; AIE2P-RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc0
bb.1.entry:
liveins: $p0, $p1, $d1
%20:ep = COPY $p0
Expand Down Expand Up @@ -444,10 +445,14 @@ body: |
; AIE2P-RA-NEXT: renamable $dc1 = MOV_PD_imm11_pseudo 0
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: bb.1:
; AIE2P-RA-NEXT: liveins: $dc1, $dj1, $dn1, $m1, $p0
; AIE2P-RA-NEXT: liveins: $d1:0x0000000000200E00, $p0
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: $p0, $dc1 = PADDA_2D_split killed $p0, killed $m1, killed $dn1, killed $dj1, killed $dc1
; AIE2P-RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc1
; AIE2P-RA-NEXT: renamable $dc0 = COPY renamable $dc1
; AIE2P-RA-NEXT: renamable $dn0 = COPY renamable $dn1
; AIE2P-RA-NEXT: renamable $dj0 = COPY renamable $dj1
; AIE2P-RA-NEXT: renamable $m0 = COPY killed renamable $m1
; AIE2P-RA-NEXT: $p0, $dc0 = PADDA_2D_split killed $p0, killed $m0, killed $dn0, killed $dj0, killed $dc0
; AIE2P-RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc0
bb.1.entry:
liveins: $p0, $m1, $dn1, $dj1
%20:ep = COPY $p0
Expand Down Expand Up @@ -623,7 +628,7 @@ body: |
; AIE2P-RA-NEXT: renamable $dc0 = LDA_dms_lda_idx_imm killed renamable $p1, 12
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: bb.1:
; AIE2P-RA-NEXT: liveins: $dc0, $dj0, $dn0, $p0
; AIE2P-RA-NEXT: liveins: $d0:0x0000000000000E00, $p0
; AIE2P-RA-NEXT: {{ $}}
; AIE2P-RA-NEXT: $p0, $dc0 = PADDA_2D_split killed $p0, undef $m0, killed $dn0, killed $dj0, killed $dc0
; AIE2P-RA-NEXT: PseudoRET implicit $lr, implicit killed renamable $p0, implicit killed renamable $dc0
Expand Down