[MC] Cache MCRegAliasIterator (#93510)

AMDGPU has a lot of registers, almost 9000. Many of those registers have
aliases. For instance, SGPR0 has a ton of aliases due to the presence of
register tuples. It's even worse if you query the aliases of a register
tuple itself. A large register tuple can have hundreds of aliases
because it may include 16 registers, and each of those registers have
their own tuples as well.

The current implementation of MCRegAliasIterator is not good at this. In
some extreme cases it can iterate, 7000 more times than
necessary, just giving duplicates over and over again and using a lot of
expensive iterators.

This patch implements a cache system for MCRegAliasIterator. It does the
expensive part only once and then saves it for us so the next iterations
on that register's aliases are just a map lookup.

Furthermore, the cached data is uniqued (and sorted). Thus, this speeds
up code by both speeding up the iterator itself, but also by minimizing
the number of loop iterations users of the iterator do.

Author: Pierre-vh

llvm/include/llvm/MC/MCRegisterInfo.h

@@ -187,6 +187,9 @@ class MCRegisterInfo {
   DenseMap<MCRegister, int> L2SEHRegs;        // LLVM to SEH regs mapping
   DenseMap<MCRegister, int> L2CVRegs;         // LLVM to CV regs mapping
 
+  mutable std::vector<std::vector<MCPhysReg>> RegAliasesCache;
+  ArrayRef<MCPhysReg> getCachedAliasesOf(MCPhysReg R) const;
+
   /// Iterator class that can traverse the differentially encoded values in
   /// DiffLists. Don't use this class directly, use one of the adaptors below.
   class DiffListIterator
@@ -263,6 +266,7 @@ class MCRegisterInfo {
   friend class MCRegUnitIterator;
   friend class MCRegUnitMaskIterator;
   friend class MCRegUnitRootIterator;
+  friend class MCRegAliasIterator;
 
   /// Initialize MCRegisterInfo, called by TableGen
   /// auto-generated routines. *DO NOT USE*.
@@ -298,6 +302,8 @@ class MCRegisterInfo {
     EHDwarf2LRegsSize = 0;
     Dwarf2LRegs = nullptr;
     Dwarf2LRegsSize = 0;
+
+    RegAliasesCache.resize(NumRegs);
   }
 
   /// Used to initialize LLVM register to Dwarf
@@ -723,63 +729,30 @@ class MCRegUnitRootIterator {
   }
 };
 
-/// MCRegAliasIterator enumerates all registers aliasing Reg.  If IncludeSelf is
-/// set, Reg itself is included in the list.  This iterator does not guarantee
-/// any ordering or that entries are unique.
+/// MCRegAliasIterator enumerates all registers aliasing Reg.
 class MCRegAliasIterator {
 private:
-  MCRegister Reg;
-  const MCRegisterInfo *MCRI;
-  bool IncludeSelf;
-
-  MCRegUnitIterator RI;
-  MCRegUnitRootIterator RRI;
-  MCSuperRegIterator SI;
+  const MCPhysReg *It = nullptr;
+  const MCPhysReg *End = nullptr;
 
 public:
   MCRegAliasIterator(MCRegister Reg, const MCRegisterInfo *MCRI,
-                     bool IncludeSelf)
-    : Reg(Reg), MCRI(MCRI), IncludeSelf(IncludeSelf) {
-    // Initialize the iterators.
-    for (RI = MCRegUnitIterator(Reg, MCRI); RI.isValid(); ++RI) {
-      for (RRI = MCRegUnitRootIterator(*RI, MCRI); RRI.isValid(); ++RRI) {
-        for (SI = MCSuperRegIterator(*RRI, MCRI, true); SI.isValid(); ++SI) {
-          if (!(!IncludeSelf && Reg == *SI))
-            return;
-        }
-      }
-    }
-  }
-
-  bool isValid() const { return RI.isValid(); }
-
-  MCRegister operator*() const {
-    assert(SI.isValid() && "Cannot dereference an invalid iterator.");
-    return *SI;
+                     bool IncludeSelf) {
+    ArrayRef<MCPhysReg> Cache = MCRI->getCachedAliasesOf(Reg);
+    assert(Cache.back() == Reg);
+    It = Cache.begin();
+    End = Cache.end();
+    if (!IncludeSelf)
+      --End;
   }
 
-  void advance() {
-    // Assuming SI is valid.
-    ++SI;
-    if (SI.isValid()) return;
-
-    ++RRI;
-    if (RRI.isValid()) {
-      SI = MCSuperRegIterator(*RRI, MCRI, true);
-      return;
-    }
+  bool isValid() const { return It != End; }
 
-    ++RI;
-    if (RI.isValid()) {
-      RRI = MCRegUnitRootIterator(*RI, MCRI);
-      SI = MCSuperRegIterator(*RRI, MCRI, true);
-    }
-  }
+  MCRegister operator*() const { return *It; }
 
   MCRegAliasIterator &operator++() {
     assert(isValid() && "Cannot move off the end of the list.");
-    do advance();
-    while (!IncludeSelf && isValid() && *SI == Reg);
+    ++It;
     return *this;
   }
 };

llvm/lib/MC/MCRegisterInfo.cpp

@@ -20,6 +20,90 @@
 
 using namespace llvm;
 
+namespace {
+/// MCRegAliasIterator enumerates all registers aliasing Reg.  This iterator
+/// does not guarantee any ordering or that entries are unique.
+class MCRegAliasIteratorImpl {
+private:
+  MCRegister Reg;
+  const MCRegisterInfo *MCRI;
+
+  MCRegUnitIterator RI;
+  MCRegUnitRootIterator RRI;
+  MCSuperRegIterator SI;
+
+public:
+  MCRegAliasIteratorImpl(MCRegister Reg, const MCRegisterInfo *MCRI)
+      : Reg(Reg), MCRI(MCRI) {
+
+    // Initialize the iterators.
+    for (RI = MCRegUnitIterator(Reg, MCRI); RI.isValid(); ++RI) {
+      for (RRI = MCRegUnitRootIterator(*RI, MCRI); RRI.isValid(); ++RRI) {
+        for (SI = MCSuperRegIterator(*RRI, MCRI, true); SI.isValid(); ++SI) {
+          if (Reg != *SI)
+            return;
+        }
+      }
+    }
+  }
+
+  bool isValid() const { return RI.isValid(); }
+
+  MCRegister operator*() const {
+    assert(SI.isValid() && "Cannot dereference an invalid iterator.");
+    return *SI;
+  }
+
+  void advance() {
+    // Assuming SI is valid.
+    ++SI;
+    if (SI.isValid())
+      return;
+
+    ++RRI;
+    if (RRI.isValid()) {
+      SI = MCSuperRegIterator(*RRI, MCRI, true);
+      return;
+    }
+
+    ++RI;
+    if (RI.isValid()) {
+      RRI = MCRegUnitRootIterator(*RI, MCRI);
+      SI = MCSuperRegIterator(*RRI, MCRI, true);
+    }
+  }
+
+  MCRegAliasIteratorImpl &operator++() {
+    assert(isValid() && "Cannot move off the end of the list.");
+    do
+      advance();
+    while (isValid() && *SI == Reg);
+    return *this;
+  }
+};
+} // namespace
+
+ArrayRef<MCPhysReg> MCRegisterInfo::getCachedAliasesOf(MCPhysReg R) const {
+  auto &Aliases = RegAliasesCache[R];
+  if (!Aliases.empty())
+    return Aliases;
+
+  for (MCRegAliasIteratorImpl It(R, this); It.isValid(); ++It)
+    Aliases.push_back(*It);
+
+  sort(Aliases);
+  Aliases.erase(unique(Aliases), Aliases.end());
+  assert(none_of(Aliases, [&](auto &Cur) { return R == Cur; }) &&
+         "MCRegAliasIteratorImpl includes Self!");
+
+  // Always put "self" at the end, so the iterator can choose to ignore it.
+  // For registers without aliases, it also serves as a sentinel value that
+  // tells us to not recompute the alias set.
+  Aliases.push_back(R);
+  Aliases.shrink_to_fit();
+  return Aliases;
+}
+
 MCRegister
 MCRegisterInfo::getMatchingSuperReg(MCRegister Reg, unsigned SubIdx,
                                     const MCRegisterClass *RC) const {

llvm/test/CodeGen/ARM/constant-island-movwt.mir

@@ -898,13 +898,14 @@ body:             |
 # CHECK-NEXT:    CONSTPOOL_ENTRY 1, %const.0, 4
 # CHECK-NEXT: {{^  $}}
 # CHECK-NEXT:   bb.2.entry (align 2):
-# CHECK-NEXT:   liveins: $d13, $s27, $r10, $r9, $r8, $s26, $d12, $s25, $s24,
-# CHECK-SAME:            $d15, $s30, $s31, $d14, $s28, $s29, $lr, $r0, $d21,
-# CHECK-SAME:            $r3, $q10, $d20, $d17, $r2, $d25, $q11, $d22, $d23,
-# CHECK-SAME:            $r1, $q8, $d16, $s3, $q14, $d28, $d29, $d19, $s17,
-# CHECK-SAME:            $d8, $s16, $r6, $r7, $r4, $q12, $q9, $d18, $s0, $q15,
-# CHECK-SAME:            $d30, $d31, $r12, $s1, $d0, $d24, $s2, $d1, $q0, $s6,
-# CHECK-SAME:            $d3, $d2, $s4, $q1, $s7, $s5, $d9, $s18, $s19, $q4
+# CHECK-NEXT:   liveins: $s26, $s27, $r10, $r9, $r8, $d13, $s24, $s25,
+# CHECK-SAME:            $d12, $d15, $s30, $s31, $d14, $s28, $s29, $lr,
+# CHECK-SAME:            $d21, $q10, $r7, $r0, $d20, $d17, $r2, $q12,
+# CHECK-SAME:            $q11, $d22, $d23, $r1, $q8, $d16, $d30, $q14,
+# CHECK-SAME:            $d28, $d29, $d19, $s17, $r4, $d8, $r6, $r3,
+# CHECK-SAME:            $s16, $d25, $q9, $d18, $s0, $d31, $s3, $q15,
+# CHECK-SAME:            $r12, $d0, $s1, $d24, $d1, $s2, $q0, $s5, $d2,
+# CHECK-SAME:            $q1, $s4, $s7, $d3, $s6, $d9, $s18, $s19, $q4
 # CHECK-NEXT: {{^  $}}
 # CHECK-NEXT:     $r5 = t2MOVi16 target-flags(arm-lo16) @.str.84, 14 /* CC::al */, $noreg
 # CHECK-NEXT:     $r5 = t2MOVTi16 $r5, target-flags(arm-hi16) @.str.84, 14 /* CC::al */, $noreg