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

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-13T10:56:43Z

There is no failures across all benchmarks.

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.

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

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.

I will keep as is, this code was just refactored.

Adding an assert crashes here because we cannot properly handle low and high parts of the 3d Registers, so a whole d0 or d4 register in the case of 3d_0.
They do not get returned from the TRI.getSubRegSplit()
Could you add a comment ?
Also do we know, why we did not add it to TRI.getSubRegSplit()?

These are not full copies but partial copies, that we currently don't handle

I think we are leaving something on the table here. not sure if it is needed for this PR though

Also do we know, why we did not add it to TRI.getSubRegSplit()?
We always rewrite into the smallest lanes as possible.

Here we are concerned with the things that we can rewrite after. Any update here should be aligned with the rewrite after, and it is not clear the benefits afterwards. I really don't want to complicate this PR even more. Any improvement could be done with an established baseline as follow up.

I included a TODO comment related to this investigation.

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?

This was just a refactor.

Is there a reason why we don't handle copies uniform in this method?
we also encounter copies in this case here if (SubReg && SubRegSplit.count(SubReg))

Because we should be able to full expand the copy after. At this moment we don't handle cases like %7.sub_hi_dim:eds = COPY %6:ed.

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?

This was just a refactor.

This will catch things like %7.sub_hi_dim:eds = COPY %6:ed, not handled after.

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp

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.

F-Stuckmann · 2025-11-21T08:09:44Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+  }
+
+  // Rewrite full copies into multiple copies using subregs
+  for (MachineInstr &MI : make_early_inc_range(MRI.reg_instructions(Reg))) {


Do we rewrite Full copies again when we AssignedPhysReg is not set?

I see that they may occur in case we have a VReg0_unallocated = COPY VReg1_unallocated. Considering Conv2D_bft16, we start with 5 3D virtual registers and end up creating 32 ones, so I this may be possible and I would like to keep this this possibility.

F-Stuckmann · 2025-11-21T08:24:35Z

llvm/lib/Target/AIE/AIESuperRegUtils.cpp

+      rewriteFullCopy(MI, LIS, *TII, TRI, VRM, LRM);
+  }
+
+  LLVM_DEBUG(dbgs() << "  Splitting range " << LIS.getInterval(Reg) << "\n");


Do we split ranges again when we AssignedPhysReg is not set?

Yes, in this part of the function we change the proper registers by others with SubRegisterClass (allocated) and LargestLegalSuperClass (unallocated). Also, we are not doing this again, each pass covesr one subset: allocated or unallocated.

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.

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-21T14:58:56Z

LGTM @martien-de-jong what do you say

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

andcarminati commented Nov 17, 2025

View reviewed changes

llvm/lib/Target/AIE/AIESuperRegRewriter.cpp 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 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/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

andcarminati and others added 9 commits November 21, 2025 03:32

[AIE2P] Extend allocation filter to include only Target MIs

47bfe10

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…

ad8d506

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

[AIEX][NFC] Refator some AIESuperReg utilities

97b0f9d

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

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

f5abcb9

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

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

d9e6d4d

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

e044750

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

[AIEX] Expand copy bundles for unallocated dest VRegs

b8476f5

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 8dba52e to b8476f5 Compare November 21, 2025 12:34

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

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!

Uh oh!

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!

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!

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!

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati commented Oct 21, 2025 •

edited

Loading

andcarminati Nov 21, 2025 •

edited

Loading

F-Stuckmann Nov 21, 2025 •

edited

Loading