[AMDGPU] Introduce RC flags for vector register classes

[cdevadas]

Summary

Configure and use the TSFlags in TargetRegisterClass to
have unique flags for VGPR and AGPR register classes.
The vector register class queries like hasVGPRs will
now become more efficient with just a bitwise operation.

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

314

Aren't TSFlags initialized to 0 by default? I think you can drop most of the changes to use SIRegisterClass and only keep it for VGPRs and AGPRs.

[cdevadas]

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

314

Yes, the bit vector is initialized to zero at its declaration statement in Target.td.

I thought having SIRegisterClass everywhere would make it simpler for any additional flags set in the future.
I could otherwise leave a comment where I define SIRegisterClass to indicate it and currently use it only for vector RCs.

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

314

No strong opinion, but I'd prefer a smaller diff.

[cdevadas]

SIRegisterClass interface only for the reg classes that use TSFlags.

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

853

This is suspicious. I wonder why SIRegisterInfo::isVGPR() works. Probably because VS classes are unallocatable and it is called only after selection when we do not have unallocatable classes anymore.

Anyway, using this bit on a combined classes seems misleading. Maybe it shall be named 'HasVGPRs' instead?

[cdevadas]

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

853

Anyway, using this bit on a combined classes seems misleading. Maybe it shall be named 'HasVGPRs' instead?

Yes, HasVGPRs flag makes more sense for the combined classes. I will update the patch.

853

This is suspicious. I wonder why SIRegisterInfo::isVGPR() works. Probably because VS classes are unallocatable and it is called only after selection when we do not have unallocatable classes anymore.

Right, it is possible that SIRegisterInfo::isVGPR() can identify it as an exclusive vgpr class. It is true even with the upstream compiler today. Maybe we don't have a test or a use case to catch it. The unallocatable classes do appear in the register-class query during selection. Do you think it is worth having another flag for SGPRs?

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2169

I think we need a comment here it only works with an allocatable class until we need to change it to work with combined classes.

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

853

It might be beneficial to add an SGPR flag, but not necessarily now. I think all of that will need a revisit when a first combined class will become allocatable. For now a rename of flags shall be OK.

[rampitec]

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

2169

Not here, actually. That is isVGPRClass would be broken if given a VS class. So it seems we will need a flag for SGPRs and check all 3 flags. Not in this patch.

[cdevadas]

Addressed the review comments.

[rampitec]

LGTM

llvm/lib/Target/AMDGPU/SIDefines.h

 //===-- SIDefines.h - SI Helper Macros ----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, 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
 //
 /// \file
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AMDGPU_SIDEFINES_H
 #define LLVM_LIB_TARGET_AMDGPU_SIDEFINES_H
 
 #include "llvm/MC/MCInstrDesc.h"
 
 namespace llvm {
 
+namespace SIRCFlags {
+// This needs to be kept in sync with the field bits in SIRegisterClass.
+enum : uint8_t {
+  // For vector registers.
+  IsVGPR = 1 << 0,
+  IsAGPR = 1 << 1
+}; // enum
+} // namespace SIRCFlags
+
 namespace SIInstrFlags {
 // This needs to be kept in sync with the field bits in InstSI.
 enum : uint64_t {
   // Low bits - basic encoding information.
   SALU = 1 << 0,
   VALU = 1 << 1,
 
   // SALU instruction formats.

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

public:
 
   /// \returns true if this class ID contains only SGPR registers
   bool isSGPRClassID(unsigned RCID) const {
     return isSGPRClass(getRegClass(RCID));
   }
 
   bool isSGPRReg(const MachineRegisterInfo &MRI, Register Reg) const;
 
+  /// \returns true if this class contains only VGPR registers
+  bool isVGPRClass(const TargetRegisterClass *RC) const {
+    return hasVGPRs(RC) && !hasAGPRs(RC);
+  }
+
   /// \returns true if this class contains only AGPR registers
   bool isAGPRClass(const TargetRegisterClass *RC) const {
     return hasAGPRs(RC) && !hasVGPRs(RC);
   }
 
   /// \returns true if this class contains VGPR registers.
   bool hasVGPRs(const TargetRegisterClass *RC) const;
 

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

bool SIRegisterInfo::isSGPRReg(const MachineRegisterInfo &MRI,
   const TargetRegisterClass *RC;
   if (Reg.isVirtual())
     RC = MRI.getRegClass(Reg);
   else
     RC = getPhysRegClass(Reg);
   return isSGPRClass(RC);
 }
 
-// TODO: It might be helpful to have some target specific flags in
-// TargetRegisterClass to mark which classes are VGPRs to make this trivial.
 bool SIRegisterInfo::hasVGPRs(const TargetRegisterClass *RC) const {
-  unsigned Size = getRegSizeInBits(*RC);
-  if (Size == 16) {
-    return getCommonSubClass(&AMDGPU::VGPR_LO16RegClass, RC) != nullptr ||
-           getCommonSubClass(&AMDGPU::VGPR_HI16RegClass, RC) != nullptr;
-  }
-  const TargetRegisterClass *VRC = getVGPRClassForBitWidth(Size);
-  if (!VRC) {
-    assert(Size < 32 && "Invalid register class size");
-    return false;
-  }
-  return getCommonSubClass(VRC, RC) != nullptr;
+  return RC->TSFlags & SIRCFlags::IsVGPR;
 }
 
 bool SIRegisterInfo::hasAGPRs(const TargetRegisterClass *RC) const {
-  unsigned Size = getRegSizeInBits(*RC);
-  if (Size < 16)
-    return false;
-  const TargetRegisterClass *ARC = getAGPRClassForBitWidth(Size);
-  if (!ARC) {
-    assert(getVGPRClassForBitWidth(Size) && "Invalid register class size");
-    return false;
-  }
-  return getCommonSubClass(ARC, RC) != nullptr;
+  return RC->TSFlags & SIRCFlags::IsAGPR;
 }
 
 const TargetRegisterClass *
 SIRegisterInfo::getEquivalentVGPRClass(const TargetRegisterClass *SRC) const {
   unsigned Size = getRegSizeInBits(*SRC);
   const TargetRegisterClass *VRC = getVGPRClassForBitWidth(Size);
   assert(VRC && "Invalid register class size");
   return VRC;
SIRegisterInfo::getRegClassForReg(const MachineRegisterInfo &MRI,
                                   Register Reg) const {
   return Reg.isVirtual() ? MRI.getRegClass(Reg) : getPhysRegClass(Reg);
 }
 
 bool SIRegisterInfo::isVGPR(const MachineRegisterInfo &MRI,
                             Register Reg) const {
   const TargetRegisterClass *RC = getRegClassForReg(MRI, Reg);
   // Registers without classes are unaddressable, SGPR-like registers.
-  return RC && hasVGPRs(RC);
+  return RC && isVGPRClass(RC);
 }
 
 bool SIRegisterInfo::isAGPR(const MachineRegisterInfo &MRI,
                             Register Reg) const {
   const TargetRegisterClass *RC = getRegClassForReg(MRI, Reg);
 
   // Registers without classes are unaddressable, SGPR-like registers.
-  return RC && hasAGPRs(RC);
+  return RC && isAGPRClass(RC);
 }
 
 bool SIRegisterInfo::shouldCoalesce(MachineInstr *MI,
                                     const TargetRegisterClass *SrcRC,
                                     unsigned SubReg,
                                     const TargetRegisterClass *DstRC,
                                     unsigned DstSubReg,
                                     const TargetRegisterClass *NewRC,

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

class SIReg <string n, bits<16> regIdx = 0> :
   let Namespace = "AMDGPU";
   let HWEncoding = regIdx;
 }
 
 class SIRegWithSubRegs <string n, list<Register> subregs, bits<16> regIdx> :
   RegisterWithSubRegs<n, subregs> {
 }
 
+// For register classes that use TSFlags.
+class SIRegisterClass <string n, list<ValueType> rTypes, int Align, dag rList>
+  : RegisterClass <n, rTypes, Align, rList> {
+  // For vector register classes.
+  field bit IsVGPR = 0;
+  field bit IsAGPR = 0;
+
+  // These need to be kept in sync with the enum SIRCFlags.
+  let TSFlags{0} = IsVGPR;
+  let TSFlags{1} = IsAGPR;
+}
+
 multiclass SIRegLoHi16 <string n, bits<16> regIdx, bit ArtificialHigh = 1,
                         bit HWEncodingHigh = 0> {
   // There is no special encoding for 16 bit subregs, these are not real
   // registers but rather operands for instructions preserving other 16 bits
   // of the result or reading just 16 bits of a 32 bit VGPR.
   // It is encoded as a corresponding 32 bit register.
   // Non-VGPR register classes use it as we need to have matching subregisters
   // to move instructions and data between ALUs.
defm AGPR#Index :
       SIRegLoHi16 <"a"#Index, Index, 1, 1>,
       DwarfRegNum<[!add(Index, 3072), !add(Index, 2048)]>;
 }
 
 //===----------------------------------------------------------------------===//
 //  Groupings using register classes and tuples
 //===----------------------------------------------------------------------===//
 
 def SCC_CLASS : RegisterClass<"AMDGPU", [i1], 1, (add SCC)> {
   let CopyCost = -1;
   let isAllocatable = 0;
 }
 
 def M0_CLASS : RegisterClass<"AMDGPU", [i32], 32, (add M0)> {
   let CopyCost = 1;
   let isAllocatable = 0;
 }

 
 class RegisterTypes<list<ValueType> reg_types> {
   list<ValueType> types = reg_types;
 }
 
 def Reg16Types : RegisterTypes<[i16, f16]>;
 def Reg32Types : RegisterTypes<[i32, f32, v2i16, v2f16, p2, p3, p5, p6]>;
 
-def VGPR_LO16 : RegisterClass<"AMDGPU", Reg16Types.types, 16,
+let IsVGPR = 1 in {
+def VGPR_LO16 : SIRegisterClass<"AMDGPU", Reg16Types.types, 16,
                               (add (sequence "VGPR%u_LO16", 0, 255))> {
   let AllocationPriority = 1;
   let Size = 16;
   let GeneratePressureSet = 0;
 }
 
-def VGPR_HI16 : RegisterClass<"AMDGPU", Reg16Types.types, 16,
+def VGPR_HI16 : SIRegisterClass<"AMDGPU", Reg16Types.types, 16,
                               (add (sequence "VGPR%u_HI16", 0, 255))> {
   let AllocationPriority = 1;
   let Size = 16;
   let GeneratePressureSet = 0;
 }
 
 // VGPR 32-bit registers
 // i16/f16 only on VI+
-def VGPR_32 : RegisterClass<"AMDGPU", !listconcat(Reg32Types.types, Reg16Types.types), 32,
+def VGPR_32 : SIRegisterClass<"AMDGPU", !listconcat(Reg32Types.types, Reg16Types.types), 32,
                             (add (sequence "VGPR%u", 0, 255))> {
   let AllocationPriority = 1;
   let Size = 32;
   let Weight = 1;
 }
+} // End IsVGPR = 1
 
 // VGPR 64-bit registers
 def VGPR_64 : SIRegisterTuples<getSubRegs<2>.ret, VGPR_32, 255, 1, 2, "v">;
 
 // VGPR 96-bit registers
 def VGPR_96 : SIRegisterTuples<getSubRegs<3>.ret, VGPR_32, 255, 1, 3, "v">;
 
 // VGPR 128-bit registers

 def VGPR_256 : SIRegisterTuples<getSubRegs<8>.ret, VGPR_32, 255, 1, 8, "v">;
 
 // VGPR 512-bit registers
 def VGPR_512 : SIRegisterTuples<getSubRegs<16>.ret, VGPR_32, 255, 1, 16, "v">;
 
 // VGPR 1024-bit registers
 def VGPR_1024 : SIRegisterTuples<getSubRegs<32>.ret, VGPR_32, 255, 1, 32, "v">;
 
-def AGPR_LO16 : RegisterClass<"AMDGPU", Reg16Types.types, 16,
+let IsAGPR = 1 in {
+def AGPR_LO16 : SIRegisterClass<"AMDGPU", Reg16Types.types, 16,
                               (add (sequence "AGPR%u_LO16", 0, 255))> {
   let isAllocatable = 0;
   let Size = 16;
   let GeneratePressureSet = 0;
 }
 
 // AccVGPR 32-bit registers
-def AGPR_32 : RegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
+def AGPR_32 : SIRegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
                             (add (sequence "AGPR%u", 0, 255))> {
   let AllocationPriority = 1;
   let Size = 32;
   let Weight = 1;
 }
+} // End IsAGPR = 1
 
 // AGPR 64-bit registers
 def AGPR_64 : SIRegisterTuples<getSubRegs<2>.ret, AGPR_32, 255, 1, 2, "a">;
 
 // AGPR 96-bit registers
 def AGPR_96 : SIRegisterTuples<getSubRegs<3>.ret, AGPR_32, 255, 1, 3, "a">;
 
 // AGPR 128-bit registers

 defm "" : SRegClass<4, 15, [v4i32, v4f32, v2i64], SGPR_128Regs, TTMP_128Regs>;
 defm "" : SRegClass<5, 16, [v5i32, v5f32], SGPR_160Regs, TTMP_160Regs>;
 defm "" : SRegClass<6, 17, [v6i32, v6f32, v3i64, v3f64], SGPR_192Regs, TTMP_192Regs>;
 defm "" : SRegClass<7, 18, [v7i32, v7f32], SGPR_224Regs, TTMP_224Regs>;
 defm "" : SRegClass<8, 19, [v8i32, v8f32, v4i64, v4f64], SGPR_256Regs, TTMP_256Regs>;
 defm "" : SRegClass<16, 20, [v16i32, v16f32, v8i64, v8f64], SGPR_512Regs, TTMP_512Regs>;
 defm "" : SRegClass<32, 21, [v32i32, v32f32, v16i64, v16f64], SGPR_1024Regs>;
 
-def VRegOrLds_32 : RegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
+def VRegOrLds_32 : SIRegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
                                  (add VGPR_32, LDS_DIRECT_CLASS)> {
   let isAllocatable = 0;
+  let IsVGPR = 1;
 }
 
 // Register class for all vector registers (VGPRs + Interpolation Registers)
 class VRegClassBase<int numRegs, list<ValueType> regTypes, dag regList> :
-    RegisterClass<"AMDGPU", regTypes, 32, regList> {
+    SIRegisterClass<"AMDGPU", regTypes, 32, regList> {
   let Size = !mul(numRegs, 32);
 
   // Requires n v_mov_b32 to copy
   let CopyCost = numRegs;
   let AllocationPriority = numRegs;
   let Weight = numRegs;
 }
 
 // Define a register tuple class, along with one requiring an even
 // aligned base register.
 multiclass VRegClass<int numRegs, list<ValueType> regTypes, dag regList> {
+  let IsVGPR = 1 in {
     // Define the regular class.
     def "" : VRegClassBase<numRegs, regTypes, regList>;
 
     // Define 2-aligned variant
     def _Align2 : VRegClassBase<numRegs, regTypes, (decimate regList, 2)>;
   }
+}
 
 defm VReg_64 : VRegClass<2, [i64, f64, v2i32, v2f32, v4f16, v4i16, p0, p1, p4],
                                 (add VGPR_64)>;
 defm VReg_96 : VRegClass<3, [v3i32, v3f32], (add VGPR_96)>;
 defm VReg_128 : VRegClass<4, [v4i32, v4f32, v2i64, v2f64], (add VGPR_128)>;
 defm VReg_160 : VRegClass<5, [v5i32, v5f32], (add VGPR_160)>;
 
 defm VReg_192 : VRegClass<6, [v6i32, v6f32, v3i64, v3f64], (add VGPR_192)>;
 defm VReg_224 : VRegClass<7, [v7i32, v7f32], (add VGPR_224)>;
 defm VReg_256 : VRegClass<8, [v8i32, v8f32, v4i64, v4f64], (add VGPR_256)>;
 defm VReg_512 : VRegClass<16, [v16i32, v16f32, v8i64, v8f64], (add VGPR_512)>;
 defm VReg_1024 : VRegClass<32, [v32i32, v32f32, v16i64, v16f64], (add VGPR_1024)>;
 
 multiclass ARegClass<int numRegs, list<ValueType> regTypes, dag regList> {
-  let CopyCost = !add(numRegs, numRegs, 1) in {
+  let CopyCost = !add(numRegs, numRegs, 1), IsAGPR = 1 in {
     // Define the regular class.
     def "" : VRegClassBase<numRegs, regTypes, regList>;
 
     // Define 2-aligned variant
     def _Align2 : VRegClassBase<numRegs, regTypes, (decimate regList, 2)>;
   }
 }
 

 
 let GeneratePressureSet = 0 in {
 // FIXME: Should specify an empty set for this. No register should
 // ever be allocated using VReg_1. This is a hack for SelectionDAG
 // that should always be lowered by SILowerI1Copies. TableGen crashes
 // on an empty register set, but also sorts register classes based on
 // the number of registerss in them. Add only one register so this is
 // sorted to the end and not preferred over VGPR_32.
-def VReg_1 : RegisterClass<"AMDGPU", [i1], 32, (add ARTIFICIAL_VGPR)> {
+def VReg_1 : SIRegisterClass<"AMDGPU", [i1], 32, (add ARTIFICIAL_VGPR)> {
   let Size = 1;
+  let IsVGPR = 1;
 }
 
-def VS_32 : RegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
+def VS_32 : SIRegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
                           (add VGPR_32, SReg_32, LDS_DIRECT_CLASS)> {
   let isAllocatable = 0;
+  let IsVGPR = 1;
 }
 
-def VS_64 : RegisterClass<"AMDGPU", [i64, f64, v2f32], 32, (add VReg_64, SReg_64)> {
+def VS_64 : SIRegisterClass<"AMDGPU", [i64, f64, v2f32], 32, (add VReg_64, SReg_64)> {
   let isAllocatable = 0;
+  let IsVGPR = 1;
 }
 
-def AV_32 : RegisterClass<"AMDGPU", VGPR_32.RegTypes, 32,
+def AV_32 : SIRegisterClass<"AMDGPU", VGPR_32.RegTypes, 32,
                           (add AGPR_32, VGPR_32)> {
   let isAllocatable = 0;
+  let IsVGPR = 1;
+  let IsAGPR = 1;
 }
 
-def AV_64 : RegisterClass<"AMDGPU", VReg_64.RegTypes, 32,
+def AV_64 : SIRegisterClass<"AMDGPU", VReg_64.RegTypes, 32,
                           (add AReg_64, VReg_64)> {
   let isAllocatable = 0;
+  let IsVGPR = 1;
+  let IsAGPR = 1;
 }
 } // End GeneratePressureSet = 0
 
-def AV_96 : RegisterClass<"AMDGPU", VReg_96.RegTypes, 32,
+let IsVGPR = 1, IsAGPR = 1 in {
+def AV_96 : SIRegisterClass<"AMDGPU", VReg_96.RegTypes, 32,
                           (add AReg_96, VReg_96)> {
   let isAllocatable = 0;
 }
 
-def AV_128 : RegisterClass<"AMDGPU", VReg_128.RegTypes, 32,
+def AV_128 : SIRegisterClass<"AMDGPU", VReg_128.RegTypes, 32,
                           (add AReg_128, VReg_128)> {
   let isAllocatable = 0;
 }
 
-def AV_160 : RegisterClass<"AMDGPU", VReg_160.RegTypes, 32,
+def AV_160 : SIRegisterClass<"AMDGPU", VReg_160.RegTypes, 32,
                           (add AReg_160, VReg_160)> {
   let isAllocatable = 0;
 }
+} // End IsVGPR = 1, IsAGPR = 1
 
 //===----------------------------------------------------------------------===//
 //  Register operands
 //===----------------------------------------------------------------------===//
 
 class RegImmMatcher<string name> : AsmOperandClass {
   let Name = name;
   let RenderMethod = "addRegOrImmOperands";