[AIEX] Register re-allocation for GPRs

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

@@ -44,6 +44,8 @@ static cl::opt<bool> AggressiveReAlloc(
     "aie-aggressive-realloc", cl::Hidden, cl::init(false),
     cl::desc("Aggressively de-allocate live-through registers to favor "
              "loop-local registers"));
+static cl::opt<bool> GPRRealloc("aie-gpr-realloc", cl::Hidden, cl::init(false),
+                                cl::desc("Re-allocate GPRs as well"));
 
 namespace {
 
@@ -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()};
+
+  // Exclude CSRs
+  for (const MCPhysReg *CSR = MRI->getCalleeSavedRegs(); CSR && *CSR; ++CSR)
+    RegSet[*CSR] = true;
+
   for (const auto *RC : RegClasses) {
 
     LLVM_DEBUG(dbgs() << "Allowed registers in RC=" << TRI->getRegClassName(RC)
@@ -355,7 +362,10 @@ bool AIEWawRegRewriter::isWorthRenaming(const Register &Reg,
   if (!VRM->hasPhys(Reg))
     return false;
 
-  if (!TRI->isVecOrAccRegClass(*(MRI->getRegClass(Reg))))
+  // Only consider vec/acc registers as candidates, and optionally GPRs.
+  if (!TRI->isVecOrAccRegClass(*(MRI->getRegClass(Reg))) &&
+      (!GPRRealloc ||
+       !TRI->getGPRRegClass(*MF)->hasSubClassEq(MRI->getRegClass(Reg))))
     return false;
 
   return !VRegWithCopies[Reg.virtRegIndex()];

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

@@ -0,0 +1,55 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# NOTE: Example file for Write After Write Register Renaming in Loop test
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+
+# RUN: llc -mtriple=aie2 -verify-machineinstrs --start-before=greedy --stop-after=virtregrewriter \
+# RUN:   --aie-gpr-realloc %s -o - | FileCheck %s
+
+
+# Check general purpose registers can also be renamed.
+---
+name:            gpr_renaming
+alignment:       16
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: gpr_renaming
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $r0, $r1, $r2, $r8
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   LoopStart $r0, 0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $r1, $r2, $r8
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   renamable $r0 = AND $r1, $r2
+  ; CHECK-NEXT:   renamable $r3 = AND $r1, $r8
+  ; CHECK-NEXT:   renamable $r4 = AND killed renamable $r0, renamable $r3
+  ; CHECK-NEXT:   dead renamable $r5 = AND killed renamable $r3, killed renamable $r4
+  ; CHECK-NEXT:   PseudoLoopEnd <mcsymbol .L_1120>, %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   PseudoRET implicit $lr
+  bb.0.entry:
+  successors: %bb.1
+  liveins: $r0, $r1, $r2, $r8
+  LoopStart $r0, 0
+  bb.1:
+    successors: %bb.1, %bb.2
+    liveins: $r1, $r2, $r8
+    %0:er = AND $r1, $r2
+    %1:er = AND $r1, $r8
+    %2:er = AND %0, %1
+    %3:er = AND %1, %2
+    PseudoLoopEnd <mcsymbol .L_1120>, %bb.1
+  bb.2:
+    PseudoRET implicit $lr
+...
+