[TableGen] Emit table mapping physical registers to base classes

[critson]

Summary

Allow targets to define a mapping from registers to register
classes such that each register has exactly one base class.
As registers may be in multiple register classes the base class
is determined by the container class with the lowest BaseClassOrder.

Only register classes with BaseClassOrder set are considered
when determining the base classes. By default BaseClassOrder is
unset in RegisterClass so no code is generated unless a target
explicit defines one or more base register classes.

[foad]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1202

Comparing with nullopt is weird. Can't you rely on implicit conversion to bool?

1607

You're not using Order here.

1611

< UINT8_MAX should be enough.

1621

This is return std::tie(LHS->getBaseClassOrder(), LHS->EnumValue) < std::tie(RHS->getBaseClassOrder(), RHS->EnumValue)

1628

Why + 1?

1648

Put nullptr as the first entry in the table to avoid this if?

[arsenm]

Thanks, it's nice to finally move some of this stuff into tablegen. Can you add a test, or at least the AMDGPU sample usage?

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

696

Should elaborate on what base means here.

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1628

Can do this directly in constructor

[critson]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1611

Not if we use first element for nullptr (in revised version).

1621

std::tie cannot handle std::optional argument it seems, so I have done this with std::pair instead.

1628

+ 1 is for NoRegister which is not part of Regs. I have added a comment to document this.

[critson]

• Address reviewer feedback
• Add lit test

[critson]

• Fix trivial off by one error in base class table

[foad]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1611

It would work fine to emit BaseClasses[256], so BaseClasses.size() can be 255, but you're restricting it to a max value of 253 here, so it's an off-by-2 error.

1621

Comparing the std::optionals themselves is a bit silly when you have already checked that they both have a value. I would suggest just adding a *: std::tie(*LHS->getBaseClassOrder(), LHS->EnumValue) ...

1644

Just assert Reg < sizeof(Mapping). Needs a small fix in IsCopyFromSGPR in D139422.

[critson]

• Use assertion instead of range check when resolving base classes

[critson]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1611

OK, yes, because of the comparison being less-than.

1621

I agree the optionals can be dereferenced, but we still cannot use std::tie as they resulting value is not an l-value. I really don't think it matters to use std::pair here (or std::tuple), this code is executed once during tablegen on a maximum of 254 values.

[critson]

• Address reviewer comments

[foad]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1612

This is still off by one. UINT8_MAX is 255 not 256.

[foad]

llvm/utils/TableGen/RegisterInfoEmitter.cpp

1621

Ack. I had misunderstood std::tie.

[critson]

• Extend available base class range by one

[foad]

LGTM, thanks!

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

if (!IdxA)
       return LaneMask;
     return reverseComposeSubRegIndexLaneMaskImpl(IdxA, LaneMask);
   }
 
   /// Debugging helper: dump register in human readable form to dbgs() stream.
   static void dumpReg(Register Reg, unsigned SubRegIndex = 0,
                       const TargetRegisterInfo *TRI = nullptr);
 
+  /// Return target defined base register class for a physical register.
+  /// This is the register class with the lowest BaseClassOrder containing the
+  /// register.
+  /// Will be nullptr if the register is not in any base register class.
+  virtual const TargetRegisterClass *getPhysRegBaseClass(MCRegister Reg) const {
+    return nullptr;
+  }
+
 protected:
   /// Overridden by TableGen in targets that have sub-registers.
   virtual unsigned composeSubRegIndicesImpl(unsigned, unsigned) const {
     llvm_unreachable("Target has no sub-registers");
   }
 
   /// Overridden by TableGen in targets that have sub-registers.
   virtual LaneBitmask

llvm/include/llvm/Target/Target.td

class RegisterClass<string namespace, list<ValueType> regTypes, int alignment,
   // register class when it is being used an an assembly operand. If this is
   // non-empty, an anonymous diagnostic type enum value will be generated, and
   // the assembly matcher will provide a function to map from diagnostic types
   // to message strings.
   string DiagnosticString = "";
 
   // Target-specific flags. This becomes the TSFlags field in TargetRegisterClass.
   bits<8> TSFlags = 0;
+
+  // If set then consider this register class to be the base class for registers in
+  // its MemberList.  The base class for registers present in multiple base register
+  // classes will be resolved in the order defined by this value, with lower values
+  // taking precedence over higher ones.  Ties are resolved by enumeration order.
+  int BaseClassOrder = ?;
 }
 
 // The memberList in a RegisterClass is a dag of set operations. TableGen
 // evaluates these set operations and expand them into register lists. These
 // are the most common operation, see test/TableGen/SetTheory.td for more
 // examples of what is possible:
 //
 // (add R0, R1, R2) - Set Union. Each argument can be an individual register, a

llvm/test/TableGen/RegisterInfoEmitter-BaseClassOrder.td

+// RUN: llvm-tblgen -gen-register-info -I %p/../../include -I %p/Common %s | FileCheck %s
+
+include "llvm/Target/Target.td"
+
+let Namespace = "MyTarget" in {
+  def R0 : Register<"r0">; // base class BaseA
+  def R1 : Register<"r1">; // base class BaseA
+  def R2 : Register<"r2">; // base class BaseC
+  def R3 : Register<"r3">; // base class BaseC
+  def R4 : Register<"r4">; // base class BaseB
+  def R5 : Register<"r5">; // base class BaseB
+  def R6 : Register<"r6">; // no base class
+} // Namespace = "MyTarget"
+
+
+// BaseA and BaseB are equal ordered so enumeration order determines base class for overlaps
+def BaseA : RegisterClass<"MyTarget", [i32], 32, (sequence "R%u", 0, 3)> {
+  let BaseClassOrder = 1;
+}
+def BaseB : RegisterClass<"MyTarget", [i32], 32, (sequence "R%u", 3, 5)> {
+  let BaseClassOrder = 1;
+}
+
+// BaseC defined order overrides BaseA and BaseB
+def BaseC : RegisterClass<"MyTarget", [i32], 32, (sequence "R%u", 2, 3)> {
+  let BaseClassOrder = 0;
+}
+
+def MyTarget : Target;
+
+// CHECK: static const TargetRegisterClass *BaseClasses[4] = {
+// CHECK-NEXT:   nullptr,
+// CHECK-NEXT:   &MyTarget::BaseCRegClass,
+// CHECK-NEXT:   &MyTarget::BaseARegClass,
+// CHECK-NEXT:   &MyTarget::BaseBRegClass,
+// CHECK-NEXT: }
+// CHECK-NEXT: static const uint8_t Mapping[8] = {
+// CHECK-NEXT:   0,2,2,1,1,3,3,0, };

llvm/utils/TableGen/CodeGenRegisters.h

struct Key {
       bool operator<(const Key&) const;
     };
 
     // Create a non-user defined register class.
     CodeGenRegisterClass(CodeGenRegBank&, StringRef Name, Key Props);
 
     // Called by CodeGenRegBank::CodeGenRegBank().
     static void computeSubClasses(CodeGenRegBank&);
+
+    // Get ordering value among register base classes.
+    std::optional<int> getBaseClassOrder() const {
+      if (TheDef && !TheDef->isValueUnset("BaseClassOrder"))
+        return TheDef->getValueAsInt("BaseClassOrder");
+      return {};
+    }
   };
 
   // Register categories are used when we need to deterine the category a
   // register falls into (GPR, vector, fixed, etc.) without having to know
   // specific information about the target architecture.
   class CodeGenRegisterCategory {
     Record *TheDef;
     std::string Name;

llvm/utils/TableGen/RegisterInfoEmitter.cpp

OS << "  const RegClassWeight &getRegClassWeight("
      << "MCRegister) const override;\n"
      << "  bool isFixedRegister(const MachineFunction &, "
      << "MCRegister) const override;\n"
      << "  bool isArgumentRegister(const MachineFunction &, "
      << "MCRegister) const override;\n"
      << "  bool isConstantPhysReg(MCRegister PhysReg) const override final;\n"
      << "  /// Devirtualized TargetFrameLowering.\n"
      << "  static const " << TargetName << "FrameLowering *getFrameLowering(\n"
-     << "      const MachineFunction &MF);\n"
-     << "};\n\n";
+     << "      const MachineFunction &MF);\n";
 
   const auto &RegisterClasses = RegBank.getRegClasses();
+  if (llvm::any_of(RegisterClasses, [](const auto &RC) { return RC.getBaseClassOrder(); })) {
+    OS << "  const TargetRegisterClass *getPhysRegBaseClass(MCRegister Reg) const override;\n";
+  }
+
+  OS << "};\n\n";
 
   if (!RegisterClasses.empty()) {
     OS << "namespace " << RegisterClasses.front().Namespace
        << " { // Register classes\n";
 
     for (const auto &RC : RegisterClasses) {
       const std::string &Name = RC.getName();
 
OS << "  };\n  assert(RC && \"Missing regclass\");\n"
        << "  if (!Idx) return RC;\n  --Idx;\n"
        << "  assert(Idx < " << SubRegIndicesSize << " && \"Bad subreg\");\n"
        << "  unsigned TV = Table[RC->getID()][Idx];\n"
        << "  return TV ? getRegClass(TV - 1) : nullptr;\n}\n\n";
   }
 
   EmitRegUnitPressure(OS, RegBank, ClassName);
 
+  // Emit register base class mapper
+  if (!RegisterClasses.empty()) {
+    // Collect base classes
+    SmallVector<const CodeGenRegisterClass*> BaseClasses;
+    for (const auto &RC : RegisterClasses) {
+      if (RC.getBaseClassOrder())
+        BaseClasses.push_back(&RC);
+    }
+    if (!BaseClasses.empty()) {
+      // Represent class indexes with uint8_t and allocate one index for nullptr
+      assert(BaseClasses.size() <= UINT8_MAX && "Too many base register classes");
+
+      // Apply order
+      struct BaseClassOrdering {
+        bool operator()(const CodeGenRegisterClass *LHS, const CodeGenRegisterClass *RHS) const {
+          return std::pair(*LHS->getBaseClassOrder(), LHS->EnumValue)
+               < std::pair(*RHS->getBaseClassOrder(), RHS->EnumValue);
+        }
+      };
+      llvm::stable_sort(BaseClasses, BaseClassOrdering());
+
+      // Build mapping for Regs (+1 for NoRegister)
+      std::vector<uint8_t> Mapping(Regs.size() + 1, 0);
+      for (int RCIdx = BaseClasses.size() - 1; RCIdx >= 0; --RCIdx) {
+        for (const auto Reg : BaseClasses[RCIdx]->getMembers())
+          Mapping[Reg->EnumValue] = RCIdx + 1;
+      }
+
+      OS << "\n// Register to base register class mapping\n\n";
+      OS << "\n";
+      OS << "const TargetRegisterClass *" << ClassName
+         << "::getPhysRegBaseClass(MCRegister Reg)"
+         << " const {\n";
+      OS << "  static const TargetRegisterClass *BaseClasses[" << (BaseClasses.size() + 1) << "] = {\n";
+      OS << "    nullptr,\n";
+      for (const auto RC : BaseClasses)
+        OS << "    &" << RC->getQualifiedName() << "RegClass,\n";
+      OS << "  };\n";
+      OS << "  static const uint8_t Mapping[" << Mapping.size() << "] = {\n    ";
+      for (const uint8_t Value : Mapping)
+        OS << (unsigned)Value << ",";
+      OS << "  };\n\n";
+      OS << "  assert(Reg < sizeof(Mapping));\n";
+      OS << "  return BaseClasses[Mapping[Reg]];\n";
+      OS << "}\n";
+    }
+  }
+
   // Emit the constructor of the class...
   OS << "extern const MCRegisterDesc " << TargetName << "RegDesc[];\n";
   OS << "extern const MCPhysReg " << TargetName << "RegDiffLists[];\n";
   OS << "extern const LaneBitmask " << TargetName << "LaneMaskLists[];\n";
   OS << "extern const char " << TargetName << "RegStrings[];\n";
   OS << "extern const char " << TargetName << "RegClassStrings[];\n";
   OS << "extern const MCPhysReg " << TargetName << "RegUnitRoots[][2];\n";
   OS << "extern const uint16_t " << TargetName << "SubRegIdxLists[];\n";