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

andcarminati · 2025-10-21T15:14:42Z

This work is intended to avoid 2D/3D (when possible) register spills.

The idea and rationale behind this work is in a previous Draft PR: #442.

To review, I recommend to follow this PR commit by commit.

Credits also for the co-author @krishnamtibrewala.

andcarminati · 2025-10-21T15:15:04Z

QoR results:

krishnamtibrewala · 2025-10-21T15:39:04Z

Thanks you @andcarminati, very much !!

llvm/lib/Target/AIE/aie2p/AIE2PTargetMachine.cpp

krishnamtibrewala · 2025-10-21T16:02:34Z

Also do you think following commit will help ?
a681b6e

andcarminati · 2025-10-22T08:31:57Z

Also do you think following commit will help ? a681b6e

Maybe yes! As mentioned before, I prefer to keep just the minimal necessary changes. We can test after, on top of this PR.

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp

llvm/lib/Target/AIE/aie2p/AIE2PTargetMachine.cpp

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp

F-Stuckmann · 2025-11-03T14:28:02Z

llvm/lib/Target/AIE/AIESuperRegUtils.h

+                                      const AIEBaseRegisterInfo &TRI,
+                                      std::set<Register> &VisitedVRegs);
+
+SmallSet<int, 8> getRewritableSubRegs(Register Reg,


nit: can we have a comment what this function does?

Yes, it is a refactor but it is always good to document.

Could we swap the documentation here?
I think the high level function should get the documentation, the actual implementation function (the one above this) should get the slimmed documentation.

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

llvm/test/CodeGen/AIE/aie2p/ra/staged-ra-spill.mir

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

F-Stuckmann · 2025-11-21T07:20:22Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+}
+
+/// Rewrite a full copy into multiple copies using the subregs in \p CopySubRegs
+void rewriteFullCopy(MachineInstr &CopyMI, LiveIntervals &LIS,


note: in the current state we will never rewrite subreg copies, e.g. copies of d0 or d4 in the case of a 3d_0 register.
Do you know why they are not relevant for the Superregrewriter?

I guess this will never occur because it will be prevented by. getRewritableSubRegs.

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

F-Stuckmann · 2025-11-21T08:25:42Z

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

+       ++VRegIdx) {
+    const Register Reg = Register::index2VirtReg(VRegIdx);
+
+    // Ignore un-used od already allocated registers.


F-Stuckmann · 2025-11-21T08:37:37Z

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

+//
+// (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+//
+//===----------------------------------------------------------------------===//


What registers aren't assigned after the previous 2d/3d reg allocs?
Are we changing here the copies and moves not the padds?
Could you Comment why this pass is necessary and the previous passes do not pick this up?

I will complement this part.

F-Stuckmann · 2025-11-21T14:58:56Z

LGTM @martien-de-jong what do you say

F-Stuckmann · 2025-11-26T15:04:18Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+    int SubReg = RegOp.getSubReg();
+    assert(SubReg);
+    RegOp.setReg(SubRegToVReg[SubReg]);
+    RegOp.setSubReg(0);


Note: 3D and 2D reg alloc will have already allocated the the 3D and 2D users except for the copies.
We spill them in line 178, so that we only work on subregs.
The users we now encounter are the subreg copies from the rewriteFullCopy

Please note that rewriteFullCopy is to prevent failures during the rewrite. We can rewrite also operands in load and move (also immediate) instructions. See instructions covered by ItinRegClassPair.

Now we filter by register class and usage. Basically, we exclude here instructions like copies and non-2D/3D ones. Co-Authored-By: Krishnam Tibrewala <[email protected]>

…gisters Co-Authored-By: Krishnam Tibrewala <[email protected]>

Co-Authored-By: Krishnam Tibrewala <[email protected]>

The goal of this test is to check if we properly insert undef flag on the def side of a expanded full copy. On a sub-register def operand, it refers to the part of the register that isn't written. A sub-register def implicitly reads the other parts of the register being redefined unless the <undef> flag is set, and a missing flag can force the related register to be inserted in liveout set of the predecessors block, causing dominance problems. Co-Authored-By: Krishnam Tibrewala <[email protected]>

This will handle properly use of non-dominating definitions. We also change the handling of the destination registers in two parts: *Copy expansion: we replace the ogininal index by the index of the first lane copy to avoid the creation LRs with just one instruction, in this way we keep que LI correct. *Rewrite: reset dead flags if necessary. Co-Authored-By: Krishnam Tibrewala <[email protected]>

…reedy run Co-Authored-By: Krishnam Tibrewala <[email protected]>

If we don't need a full register, we can expand to individual lanes. Co-Authored-By: Krishnam Tibrewala <[email protected]>

Co-Authored-By: Krishnam Tibrewala <[email protected]>

This avoids cycles in bundles that appear in VirtRegRewriter. We also update LIs related to src and dst operands of those expanded copies. Co-Authored-By: Krishnam Tibrewala <[email protected]>

andcarminati requested review from F-Stuckmann, SagarMaheshwari99, abhinay-anubola, abnikant, katerynamuts, khallouh, konstantinschwarz, martien-de-jong, niwinanto and stephenneuendorffer as code owners October 21, 2025 15:14

krishnamtibrewala reviewed Oct 21, 2025

View reviewed changes

llvm/lib/Target/AIE/aie2p/AIE2PTargetMachine.cpp Show resolved Hide resolved

andcarminati force-pushed the andreu.extend.2d3d.allocation branch from 31b7e71 to a27561f Compare October 22, 2025 08:35

F-Stuckmann reviewed Oct 28, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp Outdated Show resolved Hide resolved

andcarminati force-pushed the andreu.extend.2d3d.allocation branch from a27561f to e124649 Compare October 31, 2025 14:07

andcarminati requested a review from mludevid as a code owner October 31, 2025 14:07

andcarminati force-pushed the andreu.extend.2d3d.allocation branch 3 times, most recently from 80e6f7c to 4c47705 Compare October 31, 2025 14:36

F-Stuckmann reviewed Nov 3, 2025

View reviewed changes

llvm/lib/Target/AIE/aie2p/AIE2PTargetMachine.cpp Outdated Show resolved Hide resolved

F-Stuckmann reviewed Nov 3, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp Outdated Show resolved Hide resolved

F-Stuckmann reviewed Nov 3, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Show resolved Hide resolved

F-Stuckmann reviewed Nov 3, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Outdated Show resolved Hide resolved

F-Stuckmann reviewed Nov 3, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Outdated Show resolved Hide resolved

andcarminati force-pushed the andreu.extend.2d3d.allocation branch 2 times, most recently from a88a541 to ab8c08d Compare November 7, 2025 13:02

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp Outdated Show resolved Hide resolved

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp Outdated Show resolved Hide resolved

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Show resolved Hide resolved

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/test/CodeGen/AIE/aie2p/ra/staged-ra-spill.mir Outdated Show resolved Hide resolved

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Outdated Show resolved Hide resolved

andcarminati force-pushed the andreu.extend.2d3d.allocation branch 3 times, most recently from 9e25710 to 8dba52e Compare November 19, 2025 10:11

F-Stuckmann reviewed Nov 21, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Show resolved Hide resolved

F-Stuckmann reviewed Nov 21, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegUtils.cpp Show resolved Hide resolved

F-Stuckmann reviewed Nov 21, 2025

View reviewed changes

andcarminati force-pushed the andreu.extend.2d3d.allocation branch from 8dba52e to b8476f5 Compare November 21, 2025 12:34

andcarminati force-pushed the andreu.extend.2d3d.allocation branch from b8476f5 to 57f2573 Compare November 24, 2025 13:25

F-Stuckmann reviewed Nov 26, 2025

View reviewed changes

andcarminati and others added 9 commits November 26, 2025 23:43

[AIE2P] Extend allocation filter to include only Target MIs

85bd327

Now we filter by register class and usage. Basically, we exclude here instructions like copies and non-2D/3D ones. Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIEX] Fix LIs in SuperRegRewriter that are related to unallocated re…

ed90c67

…gisters Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIEX][NFC] Refator some AIESuperReg utilities

10114b2

Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIE2P] Add a base lit test with unallocated 2D/3D regs before main G…

02296ba

…reedy run Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIEX] Add a Pass to expand unallocated 2D/3D into individual ones

9b5534d

If we don't need a full register, we can expand to individual lanes. Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIE2P] Add base tests exposing cycles in copy bundles

acded9d

Co-Authored-By: Krishnam Tibrewala <[email protected]>

[AIEX] Expand copy bundles for unallocated dest VRegs

a08adb9

This avoids cycles in bundles that appear in VirtRegRewriter. We also update LIs related to src and dst operands of those expanded copies. Co-Authored-By: Krishnam Tibrewala <[email protected]>

andcarminati force-pushed the andreu.extend.2d3d.allocation branch from 57f2573 to a08adb9 Compare November 27, 2025 07:04

F-Stuckmann approved these changes Nov 27, 2025

View reviewed changes

andcarminati merged commit 74d4ba4 into aie-public Nov 27, 2025
7 checks passed

andcarminati deleted the andreu.extend.2d3d.allocation branch November 27, 2025 13:19

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

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

Conversation

andcarminati commented Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

andcarminati commented Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

krishnamtibrewala commented Oct 21, 2025

Uh oh!

Uh oh!

krishnamtibrewala commented Oct 21, 2025

Uh oh!

andcarminati commented Oct 22, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andcarminati Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

F-Stuckmann commented Nov 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andcarminati Nov 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati Nov 21, 2025 •

edited

Loading

andcarminati Nov 27, 2025 •

edited

Loading