[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

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.

llvm/lib/Target/AIE/AIESuperRegUtils.h

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.

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

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 · 2025-11-10T10:48:21Z

llvm/test/CodeGen/AIE/aie2p/ra/staged-ra-check-undef.mir

+# RUN: llc -O2 -mtriple=aie2p -start-before=greedy \
+# RUN:     -stop-before=aie-unallocated-superreg-rewrite %s -o - | FileCheck %s
+
+# The goal of this test is to check if we properly insert undef flag on the def side


I could extend the verifier to cover partial definitions. It is far away from be trivial, but it could be done.

martien-de-jong · 2025-11-10T13:44:21Z

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

+static bool isRegUsedBy2DOr3DInstruction(const MachineRegisterInfo &MRI,
+                                         const Register &R) {
+
+  return std::any_of(


llvm::any_of on the range?

martien-de-jong · 2025-11-10T14:49:37Z

llvm/lib/Target/AIE/AIEUnallocatedSuperRegRewriter.cpp

+
+    for (MachineInstr &MI : MRI.def_instructions(VReg)) {
+
+      if (!MI.isBundledWithPred() || MI.isBundledWithSucc())


Can I have a comment explaining this directional asymmetry?

andcarminati · 2025-11-13T10:56:43Z

There is no failures across all benchmarks.

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

F-Stuckmann · 2025-11-14T13:40:41Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+
+    MachineInstr *PartCopy = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+                                     TII.get(TargetOpcode::COPY))
+                                 .addReg(DstReg, RegState::Define, SubRegIdx)


nit: we do not need undefs here because we fully define the subregisters that we copy to?

Yes, we start the definition on the first copy, and we consider that next copies will read the previous one. In practice, inserting undefs in all copies will not hurt in general, but as is, we have a more accurate model for incremental definitions.

F-Stuckmann · 2025-11-14T13:41:19Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+
+    // Only set undefined on the first partial copy. The first copy doesn't read
+    // other lanes, but subsequent copies do read the previously written lanes.
+    // Setting undefined on all copies breaks live interval tracking and causes


oh yeah i also found that out with subreg spilling XD

F-Stuckmann · 2025-11-14T13:55:56Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+
+  // Replace the original copy by the first one, so we automatically repair
+  // DstReg's LI.
+  LIS.ReplaceMachineInstrInMaps(MI, *FirstMI);


why aren't we inserting all the newly created machineinstructions into LIS?

Because, later on, each new lane copy will make dead the previous one if we try to fix the LI. In this case, we say that the LR will start on the first copy and then it will simply continue though the next ones, until the next last use. I this way we prevent false dead flags everywhere.

F-Stuckmann · 2025-11-14T13:56:45Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+}
+
+/// Rewrite a full copy into multiple copies using the subregs in \p CopySubRegs
+void rewriteFullCopy(MachineInstr &MI, const std::set<int> &CopySubRegs,


MI is a copy, right? Could you rename MI to CopyMI, so it becomes clearer?

F-Stuckmann · 2025-11-14T13:57:46Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  for (MachineInstr &MI : make_early_inc_range(MRI.reg_instructions(Reg))) {
+    if (MI.isFullCopy())
+      AIESuperRegUtils::rewriteFullCopy(
+          MI, TRI.getSubRegSplit(MRI.getRegClass(Reg)->getID()), LIS, *TII, TRI,


you are also passing TRI, I think you don't need to pass CopySubRegs here

You can get MRI like this:
MachineRegisterInfo &MRI = MI.getMF()->getRegInfo();

F-Stuckmann · 2025-11-14T14:32:39Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  VisitedVRegs.insert(Reg);
+  SmallSet<int, 8> UsedSubRegs;
+  for (MachineOperand &RegOp : MRI.reg_operands(Reg)) {
+    int SubReg = RegOp.getSubReg();


F-Stuckmann · 2025-11-14T14:33:30Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  SmallSet<int, 8> UsedSubRegs;
+  for (MachineOperand &RegOp : MRI.reg_operands(Reg)) {
+    int SubReg = RegOp.getSubReg();
+    if (SubReg && SubRegSplit.count(SubReg)) {


if we are a subreg, shouldn't we automatically be in the SubRegSplit?
I think we can assert it instead of checking it in the condition.

F-Stuckmann · 2025-11-14T14:40:12Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+
+  VisitedVRegs.insert(Reg);
+  SmallSet<int, 8> UsedSubRegs;
+  for (MachineOperand &RegOp : MRI.reg_operands(Reg)) {


Maybe add a comment that we are walking all defs and uses to find subreg uses and full copies.

F-Stuckmann · 2025-11-14T14:41:40Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+    int SubReg = RegOp.getSubReg();
+    if (SubReg && SubRegSplit.count(SubReg)) {
+      UsedSubRegs.insert(SubReg);
+    } else if (RegOp.getParent()->isFullCopy()) {


Why are we performing this check on uses?
Isn't this a check we can perform once on the DefI?

F-Stuckmann · 2025-11-14T14:42:30Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+      }
+      UsedSubRegs.insert(SubRegSplit.begin(), SubRegSplit.end());
+    } else {
+      LLVM_DEBUG(dbgs() << "  Cannot rewrite " << RegOp << " in "


Is this a regular use/Def of the full register?

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/AIESuperRegUtils.h 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 7, 2025

View reviewed changes

andcarminati commented Nov 10, 2025

View reviewed changes

martien-de-jong reviewed Nov 10, 2025

View reviewed changes

andcarminati force-pushed the andreu.extend.2d3d.allocation branch 4 times, most recently from 731a410 to eeedaa7 Compare November 11, 2025 14:24

andcarminati and others added 9 commits November 13, 2025 03:57

[AIE2P] Extend allocation filter to include only Target MIs

afdc992

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…

12da76d

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

[AIEX][NFC] Refator some AIESuperReg utilities

23c18b4

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

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

1968bf5

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

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

3527e24

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

fbaf155

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

[AIEX] Expand copy bundles for unallocated dest VRegs

925f9d9

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 eeedaa7 to 925f9d9 Compare November 13, 2025 11:07

F-Stuckmann reviewed Nov 14, 2025

View reviewed changes

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


		for (MachineInstr &MI : MRI.def_instructions(VReg)) {

		if (!MI.isBundledWithPred() \|\| MI.isBundledWithSucc())

[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!

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

andcarminati commented Nov 13, 2025

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati Nov 3, 2025 •

edited

Loading

andcarminati Nov 10, 2025 •

edited

Loading