Register re-allocation tweaks #340

gbossu · 2025-02-06T16:27:10Z

This is a follow-up of the work from @F-Stuckmann and @martien-de-jong. It uses the least-recently-used approach for re-allocating registers. This limits WAW dependencies and gives more freedom to the post-pipeliner.

On top of this, there are two extra commits to allow:

Re-allocation of GPRs
Aggressive re-allocation: i.e. by de-allocation live-through registers.

QoR:

AIE2: 0.3% cycle count improvements on aie-public. There are a couple of regressions, but we get them back if forcing the post-pipeliner. This PR is also necessary to get the optimal schedule for GEMM_bf16 and Conv2D_int8
AIE2p: 4.2% cycle count improvement, because the register renamer was completely disabled.

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

martien-de-jong · 2025-02-07T11:24:42Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

+  };
+
+  // Least-Recently-Used list. Append new uses at the end, search from the start
+  std::list<MCPhysReg> LRURegisters;


RoundRobin?

I guess the declaration can be closer to its first use.

martien-de-jong · 2025-02-07T11:38:28Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

@@ -380,8 +396,7 @@ AIEWawRegRewriter::getReplacementPhysReg(const Register VReg,
    if (IK == LiveRegMatrix::IK_Free) {
      // Move it to the end of the list. We return, so don't have to
      // care about invalidation
-      LRURegisters.erase(It);
-      LRURegisters.emplace_back(PhysReg);
+      moveRegAndAliasesBack(It, LRURegisters, TRI);


The aliases are probably in another regclass. I'm not sure anymore why we filtered on regclass, but if we want to move all registers in the alias set of PhysReg back, it would be cheaper to do an 'isInAliasSet(PhysReg)' than to find all aliases separately.

LRURegisters can contain registers of multiple register classes, so we need to filter out the phys regs that are not compatible.

I'm not sure I understand the comment about isInAliasSet(PhysReg). Here, if we choose $cm0, I want to move bmh0 and bml0 to the back of the list as well to make sure they are not picked when we need to re-allocate a BM register.

It's just that we do a find over the full list for each alias. I'm assuming that sweeping it once to move all aliases back would be cheaper. No big deal.

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

martien-de-jong · 2025-02-07T12:52:12Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

@@ -248,7 +253,12 @@ bool AIEWawRegRewriter::renameMBBPhysRegs(const MachineBasicBlock *MBB) {
        continue;

      if (isWorthRenaming(Reg, VRegWithCopies)) {
-        Candidates.emplace_back(&MO, VRM->getPhys(Reg));
+        MCRegister AssignedPhysReg = VRM->getPhys(Reg);


llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

martien-de-jong · 2025-02-10T08:31:53Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

+              getKilledRegAtSingledDefPoint(VReg, *MRI);
+          KilledReg &&
+          any_of(TRI->regunits(PhysReg),
+                 [&UsedUnits](MCRegUnit RU) { return UsedUnits.test(RU); })) {


nit: this is far beyond my limits to reasonably nest statements in an if condition.

I will be happy if we can extract at least one lambda/predicate like WasUsedBefore(PhysReg) ;-)

llvm/test/CodeGen/AIE/aie2/ra/waw_reg_renaming_loop.mir

llvm/lib/Target/AIE/AIE2TargetMachine.cpp

martien-de-jong · 2025-02-10T08:42:13Z

llvm/test/CodeGen/AIE/aie2/ra/waw_reg_renaming_aggressive.mir

+#
+# (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+# unit test for the WAW register renaming pass and check edge cases
+# RUN: llc -mtriple=aie2 -verify-machineinstrs --start-before=greedy --stop-after=virtregrewriter \


nit: Do we have a motivating example for aie2p?

Not that I know of (so far)

I would be very happy if it could improve SigmoidTemplated_bf16_0

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

F-Stuckmann · 2025-02-10T15:59:40Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

+        LLVM_DEBUG(dbgs() << "Candidate " << printReg(Reg, TRI, 0, MRI) << ":"
+                          << TRI->getRegClassName(MRI->getRegClass(Reg)) << " ("
+                          << TRI->getName(AssignedPhysReg) << ")\n");
+        assert(AssignedPhysReg != VirtRegMap::NO_PHYS_REG);


why not check AssignedPhysReg.isVirtual()?

Do we want to rename StackSlots?

I'm not sure I get this, AssignedPhysReg is a physreg, not a virtual one. Here I just want to make sure that every candidate for renaming does have phys reg assigned.

It still looks funny that this value is defined in the VirtRegMap scope. How does it know it conflicts with the MCRegister encoding? Would VRM->hasPhys(Reg) be more suitable?

NO_PHYS_REG is the sentinel value used by VRM to say "nothing is assigned", so that scope makes sense to me. I'll use VRM->hasPhys(Reg) nonetheless :)

llvm/lib/Target/AIE/AIE2TargetMachine.cpp

F-Stuckmann · 2025-02-11T13:14:49Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

+  // The order in Registers represent the LeastRecentlyUsed, which will
+  // distribute register use evenly
+  for (auto It = LRURegisters.begin(); It != LRURegisters.end(); ++It) {
+    MCPhysReg PhysReg = *It;


Instead of using Physical Register as a whole for the round robin scheme, we could use RegUnits.
this way, we would differentiate between different registers, that partially populate the same RegUnits, i.e. wl/h, X and y registers.

Unsure if this is necessary at this point.

I think you are suggesting to use a list of RegUnits, instead of a list of MCRegister? I have to say I do prefer the latter, because it makes it really straightforward to retrieve the least-recently used register.

F-Stuckmann · 2025-02-11T13:25:43Z

llvm/lib/Target/AIE/AIE2TargetMachine.cpp

    addPass(createAIEWawRegRewriter());
+    addPass(createGreedyRegisterAllocator());


should greedy also run after createAIEWawRegRewriter in AIE2p?

Definitely, in general, the AIEWawRegRewriter pass does not even run for AIE2p. I assumed that the pass pipeline is shared, but that part is not. Good catch!

Added a new commit for AIE2p

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

andcarminati · 2025-02-11T14:28:37Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

+  };
+
+  // Least-Recently-Used list. Append new uses at the end, search from the start
+  std::list<MCPhysReg> LRURegisters;


I think this is closer to Least-Recently-Defined. LRU. Could we have any additional gain by implementing a full LRU approach? I mean, updating the list considering the uses as well?

I think this is meant as "least recently used for re-allocation". I could rename that to "least recently allocated" maybe?

Maybe for a future reference.

It's using the common acronym LRU from caching strategies. Indeed, it is referring to the allocation in a forward traversal, not as defining a register.

llvm/test/CodeGen/AIE/aie2/ra/waw_reg_renaming_loop.mir

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

martien-de-jong · 2025-02-12T11:10:39Z

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

@@ -326,6 +328,11 @@ bool AIEWawRegRewriter::renameMBBPhysRegs(const MachineBasicBlock *MBB) {
  // For each reg class, allocate the candidates in round-robin fashion.
  // If we fail, we fall back to the original allocation
  BitVector RegSet{TRI->getNumRegs()};


nit: more descriptive name?

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

gbossu · 2025-02-12T14:50:45Z

I think I addressed all the comments. :) I also gave some QoR numbers in the PR description.

martien-de-jong · 2025-02-13T11:33:07Z

llvm/test/CodeGen/AIE/aie2p/schedule/postpipeliner/end-to-end.ll

+; CHECK-NEXT:    vadd.16 x2, x0, x2
+; CHECK-NEXT:    vadd.16 x2, x0, x2
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    vmov x0, x2


nit: no net effect in the loop, but ABI constraint ignored.

This limits re-usal of registers and therefore limits the WAW dependencies as well which will hurt for SWP. Co-authored-by: Martien de Jong <[email protected]>

Some instructions represent a single-use and single-def point. For those, it is generally fine to allocate the same phys reg to the killed input, and the def.

Needs to be enabled with care, because live-through registers will get de-allocated, and re-allocated in a final run of greedy regalloc. This can yield a sub-optimal allocation.

gbossu requested review from abhinay-anubola, abnikant, andcarminati, F-Stuckmann, katerynamuts, khallouh, konstantinschwarz, martien-de-jong, niwinanto, SagarMaheshwari99 and stephenneuendorffer as code owners February 6, 2025 16:27

gbossu commented Feb 7, 2025

View reviewed changes

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

gbossu force-pushed the gaetan.register.realloc branch from c4d21bd to 47b6897 Compare February 7, 2025 10:59

martien-de-jong reviewed Feb 7, 2025

View reviewed changes

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

martien-de-jong reviewed Feb 7, 2025

View reviewed changes

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

martien-de-jong reviewed Feb 10, 2025

View reviewed changes

llvm/test/CodeGen/AIE/aie2/ra/waw_reg_renaming_loop.mir Show resolved Hide resolved

martien-de-jong reviewed Feb 10, 2025

View reviewed changes

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

martien-de-jong reviewed Feb 10, 2025

View reviewed changes

andcarminati reviewed Feb 10, 2025

View reviewed changes

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

F-Stuckmann reviewed Feb 10, 2025

View reviewed changes

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

F-Stuckmann reviewed Feb 11, 2025

View reviewed changes

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

andcarminati reviewed Feb 11, 2025

View reviewed changes

llvm/test/CodeGen/AIE/aie2/ra/waw_reg_renaming_loop.mir Show resolved Hide resolved

andcarminati reviewed Feb 11, 2025

View reviewed changes

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

martien-de-jong reviewed Feb 12, 2025

View reviewed changes

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

gbossu force-pushed the gaetan.register.realloc branch from 47b6897 to f708d1f Compare February 12, 2025 14:32

gbossu force-pushed the gaetan.register.realloc branch 2 times, most recently from 0e974e6 to 2553b39 Compare February 13, 2025 11:25

martien-de-jong reviewed Feb 13, 2025

View reviewed changes

martien-de-jong previously approved these changes Feb 13, 2025

View reviewed changes

gbossu and others added 7 commits February 13, 2025 13:19

[AIEX] Use a round-robin approach for register re-allocation.

a52c0af

This limits re-usal of registers and therefore limits the WAW dependencies as well which will hurt for SWP. Co-authored-by: Martien de Jong <[email protected]>

[AIEX][NFC] AIEWawRegRewriter: Harmonise logging

ac01b37

[AIEX] AIEWawRegRewriter: Track register aliases in LRU list

a0bc1af

[AIEX] AIEWawRegRewriter: Tweak re-allocation for "unary" instructions

927da71

Some instructions represent a single-use and single-def point. For those, it is generally fine to allocate the same phys reg to the killed input, and the def.

[AIEX] Aggressive re-allocation to favor loop-local vregs

fe6c322

Needs to be enabled with care, because live-through registers will get de-allocated, and re-allocated in a final run of greedy regalloc. This can yield a sub-optimal allocation.

[AIEX] AIEWawRegRewriter: re-allocation for GPRs

f5d5180

[AIE2p] Enable register re-allocation

1dde07c

gbossu dismissed martien-de-jong’s stale review via 1dde07c February 13, 2025 13:23

gbossu force-pushed the gaetan.register.realloc branch from 2553b39 to 1dde07c Compare February 13, 2025 13:23

martien-de-jong approved these changes Feb 13, 2025

View reviewed changes

gbossu enabled auto-merge (rebase) February 13, 2025 14:16

gbossu merged commit 3e75a3d into aie-public Feb 13, 2025
8 checks passed

gbossu deleted the gaetan.register.realloc branch February 13, 2025 14:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Register re-allocation tweaks #340

Register re-allocation tweaks #340

gbossu commented Feb 6, 2025 •

edited

Loading

martien-de-jong Feb 7, 2025

martien-de-jong Feb 7, 2025

martien-de-jong Feb 7, 2025

gbossu Feb 7, 2025

martien-de-jong Feb 10, 2025

martien-de-jong Feb 7, 2025

martien-de-jong Feb 10, 2025

andcarminati Feb 11, 2025

martien-de-jong Feb 10, 2025

gbossu Feb 11, 2025

martien-de-jong Feb 12, 2025

F-Stuckmann Feb 10, 2025

gbossu Feb 11, 2025

martien-de-jong Feb 12, 2025 •

edited

Loading

gbossu Feb 12, 2025

F-Stuckmann Feb 11, 2025 •

edited

Loading

gbossu Feb 11, 2025

F-Stuckmann Feb 11, 2025 •

edited

Loading

gbossu Feb 11, 2025

gbossu Feb 12, 2025

andcarminati Feb 11, 2025

gbossu Feb 11, 2025

andcarminati Feb 12, 2025

martien-de-jong Feb 13, 2025

martien-de-jong Feb 12, 2025

gbossu commented Feb 12, 2025

martien-de-jong Feb 13, 2025

		addPass(createAIEWawRegRewriter());
		addPass(createGreedyRegisterAllocator());

Register re-allocation tweaks #340

Register re-allocation tweaks #340

Conversation

gbossu commented Feb 6, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

martien-de-jong Feb 12, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

F-Stuckmann Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

F-Stuckmann Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gbossu commented Feb 12, 2025

Choose a reason for hiding this comment

gbossu commented Feb 6, 2025 •

edited

Loading

martien-de-jong Feb 12, 2025 •

edited

Loading

F-Stuckmann Feb 11, 2025 •

edited

Loading

F-Stuckmann Feb 11, 2025 •

edited

Loading