[llvm] [CodeGen] Use RegUnits in RegisterClassInfo::getLastCalleeSavedAlias (PR #79996)

[Jay Foad via llvm-commits]

https://github.com/jayfoad created https://github.com/llvm/llvm-project/pull/79996

Change the implementation of getLastCalleeSavedAlias to use RegUnits
instead of register aliases. This is much faster on targets like AMDGPU
which define a very large number of overlapping register tuples.

No functional change intended. If PhysReg overlaps multiple CSRs then
getLastCalleeSavedAlias(PhysReg) could conceivably return a different
arbitrary one, but currently it is only used for some debug printing
anyway.

Differential Revision: https://reviews.llvm.org/D146734


>From 964188e87fb2f74582f6fe24e3a893b031ffac56 Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Thu, 23 Mar 2023 15:33:23 +0000
Subject: [PATCH] [CodeGen] Use RegUnits in
 RegisterClassInfo::getLastCalleeSavedAlias

Change the implementation of getLastCalleeSavedAlias to use RegUnits
instead of register aliases. This is much faster on targets like AMDGPU
which define a very large number of overlapping register tuples.

No functional change intended. If PhysReg overlaps multiple CSRs then
getLastCalleeSavedAlias(PhysReg) could conceivably return a different
arbitrary one, but currently it is only used for some debug printing
anyway.

Differential Revision: https://reviews.llvm.org/D146734
---
 llvm/include/llvm/CodeGen/RegisterClassInfo.h | 16 ++++++++++------
 llvm/lib/CodeGen/RegisterClassInfo.cpp        |  8 ++++----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/RegisterClassInfo.h b/llvm/include/llvm/CodeGen/RegisterClassInfo.h
index 0e50d2feb9b2..800bebea0ddd 100644
--- a/llvm/include/llvm/CodeGen/RegisterClassInfo.h
+++ b/llvm/include/llvm/CodeGen/RegisterClassInfo.h
@@ -56,8 +56,8 @@ class RegisterClassInfo {
   // Used only to determine if an update for CalleeSavedAliases is necessary.
   SmallVector<MCPhysReg, 16> LastCalleeSavedRegs;
 
-  // Map register alias to the callee saved Register.
-  SmallVector<MCPhysReg, 4> CalleeSavedAliases;
+  // Map regunit to the callee saved Register.
+  SmallVector<MCPhysReg> CalleeSavedAliases;
 
   // Indicate if a specified callee saved register be in the allocation order
   // exactly as written in the tablegen descriptions or listed later.
@@ -113,12 +113,16 @@ class RegisterClassInfo {
   }
 
   /// getLastCalleeSavedAlias - Returns the last callee saved register that
-  /// overlaps PhysReg, or NoRegister if Reg doesn't overlap a
+  /// overlaps PhysReg, or NoRegister if PhysReg doesn't overlap a
   /// CalleeSavedAliases.
   MCRegister getLastCalleeSavedAlias(MCRegister PhysReg) const {
-    if (PhysReg.id() < CalleeSavedAliases.size())
-      return CalleeSavedAliases[PhysReg];
-    return MCRegister::NoRegister;
+    MCRegister CSR;
+    for (MCRegUnitIterator UI(PhysReg, TRI); UI.isValid(); ++UI) {
+      CSR = CalleeSavedAliases[*UI];
+      if (CSR)
+        break;
+    }
+    return CSR;
   }
 
   /// Get the minimum register cost in RC's allocation order.
diff --git a/llvm/lib/CodeGen/RegisterClassInfo.cpp b/llvm/lib/CodeGen/RegisterClassInfo.cpp
index 17a9f55cccc0..8869a861de06 100644
--- a/llvm/lib/CodeGen/RegisterClassInfo.cpp
+++ b/llvm/lib/CodeGen/RegisterClassInfo.cpp
@@ -80,10 +80,10 @@ void RegisterClassInfo::runOnMachineFunction(const MachineFunction &mf) {
     LastCalleeSavedRegs.clear();
     // Build a CSRAlias map. Every CSR alias saves the last
     // overlapping CSR.
-    CalleeSavedAliases.assign(TRI->getNumRegs(), 0);
+    CalleeSavedAliases.assign(TRI->getNumRegUnits(), 0);
     for (const MCPhysReg *I = CSR; *I; ++I) {
-      for (MCRegAliasIterator AI(*I, TRI, true); AI.isValid(); ++AI)
-        CalleeSavedAliases[*AI] = *I;
+      for (MCRegUnitIterator UI(*I, TRI); UI.isValid(); ++UI)
+        CalleeSavedAliases[*UI] = *I;
       LastCalleeSavedRegs.push_back(*I);
     }
 
@@ -150,7 +150,7 @@ void RegisterClassInfo::compute(const TargetRegisterClass *RC) const {
     uint8_t Cost = RegCosts[PhysReg];
     MinCost = std::min(MinCost, Cost);
 
-    if (CalleeSavedAliases[PhysReg] &&
+    if (getLastCalleeSavedAlias(PhysReg) &&
         !STI.ignoreCSRForAllocationOrder(*MF, PhysReg))
       // PhysReg aliases a CSR, save it for later.
       CSRAlias.push_back(PhysReg);