[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

F-Stuckmann · 2025-11-03T13:28:52Z

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

+            MI.getMF()->getSubtarget().getInstrInfo());
+
+        // We should recognize both cases, with and without splitting. A 2D/3D
+        // instruction will always be split os splittable.


Repaired as a fixup in the refactor commit.

F-Stuckmann · 2025-11-03T13:30:21Z

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp

  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,


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

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.

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.

F-Stuckmann · 2025-11-03T14:54:26Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+      }
+      Register SrcReg = RegOp.getParent()->getOperand(1).getReg();
+      if (!VisitedVRegs.count(SrcReg) &&
+          getRewritableSubRegs(SrcReg, MRI, TRI, VisitedVRegs).empty()) {


Does it ever happen, that SrcReg has no SubRegs, but DstReg has them or vis versa?
Can they also have different SubRegs?

This part was a refactor, but as we are handling a full copy here, we can expect subregs on both sides.

As a full copy, I mean a 2d to 2d copy or a 3d to 3d copy.

F-Stuckmann · 2025-11-03T14:56:53Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  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,


F-Stuckmann · 2025-11-03T15:00:30Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  // We should make sure that all LiveIntervals are correct.
+  // If we dont't repair, MI will compose the LIs of some registers,
+  // what is not correct because MI was deleted.
+  updateLRMAndLIS(RegistersToRepair, VRM, LRM, LIS);


nit: could we get a better name that does not have the abbreviations in it?
updateLiveIntervals. Also since LIS is the most relevant Datastructure, have it first and then followed by helper structures LRM and VRM.

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]>

…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]>

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. In the test, MachineVerifier is still off, because some lanes redefinition are consirered full register redefinition, causing some non-accurate expactations aroud dead flags. The affected test is a extreme corner case where individual lanes are handled apart of original 2D register. Co-Authored-By: Krishnam Tibrewala <[email protected]>

andcarminati · 2025-11-07T13:49:03Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

-
+  MachineInstr *FirstMI = nullptr;
  SmallSet<Register, 8> RegistersToRepair;
+  bool IsFirstCopy = true;


nit: We can use FirstMI as a replacement for IsFirstCopy.

andcarminati · 2025-11-07T13:50:54Z

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

@@ -0,0 +1,172 @@
+//===----- AIEUnallocatedSuperRegRewriter.cpp - Constrain tied sub-registers


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

andcarminati and others added 2 commits November 7, 2025 01:23

[AIE2P] Extend allocation filter to include only Target MIs

6211046

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…

0d00bc6

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

andcarminati and others added 7 commits November 7, 2025 03:07

[AIEX][NFC] Refator some AIESuperReg utilities

82cafcc

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

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

ad061d1

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

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

d470b25

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

[AIE2P] Add base test exposing cycles in copy bundles

70c2adc

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

[AIE2P] Add base lit test with copy bundles that should be expanded

9401b37

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

[AIEX] Expand copy bundles for unallocated dest VRegs

b8d75ba

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 4c47705 to a88a541 Compare November 7, 2025 12:36

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

andcarminati commented Nov 7, 2025

View reviewed changes

		@@ -0,0 +1,172 @@
		//===----- AIEUnallocatedSuperRegRewriter.cpp - Constrain tied sub-registers

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

Are you sure you want to change the base?

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

Uh oh!

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!

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!

Choose a reason for hiding this comment

Uh oh!

andcarminati Nov 3, 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!

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati Nov 3, 2025 •

edited

Loading