[AMDGPU] Add missing AReg classes

[foad]

Summary

Add 96-bit, 160-bit and 256-bit AReg classes to match VReg and SReg.
NFC as far as I know, but it may avoid weird legalization problems.

[rampitec]

Missing changes to: AMDGPUAsmPrinter::analyzeResourceUsage(), AMDGPUAsmParser.cpp:getRegClass(), SITargetLowering::getRegForInlineAsmConstraint().

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

826

They all missing let Weight = ... I have somehow missed it.

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

1107

Missing AGPR* classes.

[foad]

Add weights for all AReg classes.
Update AMDGPUAsmPrinter, AMDGPUAsmParser, AMDGPUDisassembler,
SIMCCodeEmitter and SITargetLowering::getRegForInlineAsmConstraint.

[foad]

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

1107

True, but it is also missing most VGPR* classes and some SGPR* classes. Are they really all supposed to be included here?

[rampitec]

LGTM

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

1107

I think they are. If these will appear in MI operand definitions eventually it will assert.

[rampitec]

Actually SIMCCodeEmitter.cpp changes are an error.

llvm/lib/Target/AMDGPU/MCTargetDesc/SIMCCodeEmitter.cpp

422

Ugh. This should not be here. It never can be SrcA or SrcB.

[rampitec]

llvm/lib/Target/AMDGPU/MCTargetDesc/SIMCCodeEmitter.cpp

422

Ugh again. I did some testing and it turns to be right. Please disregard.

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

for (const MachineInstr &MI : MBB) {
           IsAGPR = true;
           Width = 2;
         } else if (AMDGPU::VReg_96RegClass.contains(Reg)) {
           IsSGPR = false;
           Width = 3;
         } else if (AMDGPU::SReg_96RegClass.contains(Reg)) {
           IsSGPR = true;
           Width = 3;
+        } else if (AMDGPU::AReg_96RegClass.contains(Reg)) {
+          IsSGPR = false;
+          IsAGPR = true;
+          Width = 3;
         } else if (AMDGPU::SReg_128RegClass.contains(Reg)) {
           assert(!AMDGPU::TTMP_128RegClass.contains(Reg) &&
             "trap handler registers should not be used");
           IsSGPR = true;
           Width = 4;
         } else if (AMDGPU::VReg_128RegClass.contains(Reg)) {
           IsSGPR = false;
           Width = 4;
         } else if (AMDGPU::AReg_128RegClass.contains(Reg)) {
           IsSGPR = false;
           IsAGPR = true;
           Width = 4;
         } else if (AMDGPU::VReg_160RegClass.contains(Reg)) {
           IsSGPR = false;
           Width = 5;
         } else if (AMDGPU::SReg_160RegClass.contains(Reg)) {
           IsSGPR = true;
           Width = 5;
+        } else if (AMDGPU::AReg_160RegClass.contains(Reg)) {
+          IsSGPR = false;
+          IsAGPR = true;
+          Width = 5;
         } else if (AMDGPU::SReg_256RegClass.contains(Reg)) {
           assert(!AMDGPU::TTMP_256RegClass.contains(Reg) &&
             "trap handler registers should not be used");
           IsSGPR = true;
           Width = 8;
         } else if (AMDGPU::VReg_256RegClass.contains(Reg)) {
           IsSGPR = false;
           Width = 8;
+        } else if (AMDGPU::AReg_256RegClass.contains(Reg)) {
+          IsSGPR = false;
+          IsAGPR = true;
+          Width = 8;
         } else if (AMDGPU::SReg_512RegClass.contains(Reg)) {
           assert(!AMDGPU::TTMP_512RegClass.contains(Reg) &&
             "trap handler registers should not be used");
           IsSGPR = true;
           Width = 16;
         } else if (AMDGPU::VReg_512RegClass.contains(Reg)) {
           IsSGPR = false;
           Width = 16;

llvm/lib/Target/AMDGPU/AMDGPURegisterBanks.td

 def VGPRRegBank : RegisterBank<"VGPR",
   [VGPR_LO16, VGPR_HI16, VGPR_32, VReg_64, VReg_96, VReg_128, VReg_256, VReg_512, VReg_1024]
 >;
 
 // It is helpful to distinguish conditions from ordinary SGPRs.
 def VCCRegBank : RegisterBank <"VCC", [SReg_1]>;
 
 def AGPRRegBank : RegisterBank <"AGPR",
-  [AGPR_32, AReg_64, AReg_128, AReg_512, AReg_1024]
+  [AGPR_32, AReg_64, AReg_96, AReg_128, AReg_160, AReg_256, AReg_512, AReg_1024]
 >;

llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp

 //===- AMDGPUAsmParser.cpp - Parse SI asm to MCInst instructions ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
switch (RegWidth) {
       case 8: return AMDGPU::SGPR_256RegClassID;
       case 16: return AMDGPU::SGPR_512RegClassID;
     }
   } else if (Is == IS_AGPR) {
     switch (RegWidth) {
       default: return -1;
       case 1: return AMDGPU::AGPR_32RegClassID;
       case 2: return AMDGPU::AReg_64RegClassID;
+      case 3: return AMDGPU::AReg_96RegClassID;
       case 4: return AMDGPU::AReg_128RegClassID;
+      case 5: return AMDGPU::AReg_160RegClassID;
+      case 8: return AMDGPU::AReg_256RegClassID;
       case 16: return AMDGPU::AReg_512RegClassID;
       case 32: return AMDGPU::AReg_1024RegClassID;
     }
   }
   return -1;
 }
 
 static unsigned getSpecialRegForName(StringRef RegName) {

llvm/lib/Target/AMDGPU/Disassembler/AMDGPUDisassembler.cpp

 //===- AMDGPUDisassembler.cpp - Disassembler for AMDGPU ISA ---------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//
unsigned AMDGPUDisassembler::getAgprClassId(const OpWidthTy Width) const {
   switch (Width) {
   default: // fall
   case OPW32:
   case OPW16:
   case OPWV216:
     return AGPR_32RegClassID;
   case OPW64: return AReg_64RegClassID;
   case OPW128: return AReg_128RegClassID;
+  case OPW256: return AReg_256RegClassID;
   case OPW512: return AReg_512RegClassID;
   case OPW1024: return AReg_1024RegClassID;
   }
 }
 
 
 unsigned AMDGPUDisassembler::getSgprClassId(const OpWidthTy Width) const {
   using namespace AMDGPU;

llvm/lib/Target/AMDGPU/MCTargetDesc/SIMCCodeEmitter.cpp

SIMCCodeEmitter::getAVOperandEncoding(const MCInst &MI, unsigned OpNo,
                                       const MCSubtargetInfo &STI) const {
   unsigned Reg = MI.getOperand(OpNo).getReg();
   uint64_t Enc = MRI.getEncodingValue(Reg);
 
   // VGPR and AGPR have the same encoding, but SrcA and SrcB operands of mfma
   // instructions use acc[0:1] modifier bits to distinguish. These bits are
   // encoded as a virtual 9th bit of the register for these operands.
   if (MRI.getRegClass(AMDGPU::AGPR_32RegClassID).contains(Reg) ||
-      MRI.getRegClass(AMDGPU::AReg_64RegClassID).contains(Reg))
+      MRI.getRegClass(AMDGPU::AReg_64RegClassID).contains(Reg) ||
+      MRI.getRegClass(AMDGPU::AReg_96RegClassID).contains(Reg) ||
+      MRI.getRegClass(AMDGPU::AReg_128RegClassID).contains(Reg) ||
+      MRI.getRegClass(AMDGPU::AReg_160RegClassID).contains(Reg) ||
+      MRI.getRegClass(AMDGPU::AReg_256RegClassID).contains(Reg))
     Enc |= 512;
 
   return Enc;
 }
 
 static bool needsPCRel(const MCExpr *Expr) {
   switch (Expr->getKind()) {
   case MCExpr::SymbolRef: {

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

case 'a':
         return std::make_pair(0U, nullptr);
       case 32:
       case 16:
         RC = &AMDGPU::AGPR_32RegClass;
         break;
       case 64:
         RC = &AMDGPU::AReg_64RegClass;
         break;
+      case 96:
+        RC = &AMDGPU::AReg_96RegClass;
+        break;
       case 128:
         RC = &AMDGPU::AReg_128RegClass;
         break;
+      case 160:
+        RC = &AMDGPU::AReg_160RegClass;
+        break;
+      case 256:
+        RC = &AMDGPU::AReg_256RegClass;
+        break;
       case 512:
         RC = &AMDGPU::AReg_512RegClass;
         break;
       case 1024:
         RC = &AMDGPU::AReg_1024RegClass;
         // v32 types are not legal but we support them here.
         return std::make_pair(0U, RC);
       }

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

 }
 
 static const TargetRegisterClass *getAGPRClassForBitWidth(unsigned BitWidth) {
   switch (BitWidth) {
   case 32:
     return &AMDGPU::AGPR_32RegClass;
   case 64:
     return &AMDGPU::AReg_64RegClass;
+  case 96:
+    return &AMDGPU::AReg_96RegClass;
   case 128:
     return &AMDGPU::AReg_128RegClass;
+  case 160:
+    return &AMDGPU::AReg_160RegClass;
+  case 256:
+    return &AMDGPU::AReg_256RegClass;
   case 512:
     return &AMDGPU::AReg_512RegClass;
   case 1024:
     return &AMDGPU::AReg_1024RegClass;
   default:
     return nullptr;
   }
 }
static const TargetRegisterClass *getSGPRClassForBitWidth(unsigned BitWidth) {
   }
 }
 
 // FIXME: This is very slow. It might be worth creating a map from physreg to
 // register class.
 const TargetRegisterClass *
 SIRegisterInfo::getPhysRegClass(MCRegister Reg) const {
   static const TargetRegisterClass *const BaseClasses[] = {
     &AMDGPU::VGPR_LO16RegClass,
     &AMDGPU::VGPR_HI16RegClass,
     &AMDGPU::SGPR_LO16RegClass,
     &AMDGPU::VGPR_32RegClass,
     &AMDGPU::SReg_32RegClass,
     &AMDGPU::AGPR_32RegClass,
     &AMDGPU::VReg_64RegClass,
     &AMDGPU::SReg_64RegClass,
     &AMDGPU::AReg_64RegClass,
     &AMDGPU::VReg_96RegClass,
     &AMDGPU::SReg_96RegClass,
+    &AMDGPU::AReg_96RegClass,
     &AMDGPU::VReg_128RegClass,
     &AMDGPU::SReg_128RegClass,
     &AMDGPU::AReg_128RegClass,
     &AMDGPU::VReg_160RegClass,
     &AMDGPU::SReg_160RegClass,
+    &AMDGPU::AReg_160RegClass,
     &AMDGPU::VReg_256RegClass,
     &AMDGPU::SReg_256RegClass,
+    &AMDGPU::AReg_256RegClass,
     &AMDGPU::VReg_512RegClass,
     &AMDGPU::SReg_512RegClass,
     &AMDGPU::AReg_512RegClass,
     &AMDGPU::SReg_1024RegClass,
     &AMDGPU::VReg_1024RegClass,
     &AMDGPU::AReg_1024RegClass,
     &AMDGPU::SCC_CLASSRegClass,
     &AMDGPU::Pseudo_SReg_32RegClass,

llvm/lib/Target/AMDGPU/SIRegisterInfo.td

 // VGPR 1024-bit registers
 def VGPR_1024 : SIRegisterTuples<getSubRegs<32>.ret, VGPR_32, 255, 1, 32, "v">;
 
 // AccVGPR 32-bit registers
 def AGPR_32 : RegisterClass<"AMDGPU", [i32, f32, i16, f16, v2i16, v2f16], 32,
                             (add (sequence "AGPR%u", 0, 255))> {
   let AllocationPriority = 1;
   let Size = 32;
+  let Weight = 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
 def AGPR_128 : SIRegisterTuples<getSubRegs<4>.ret, AGPR_32, 255, 1, 4, "a">;
 
+// AGPR 160-bit registers
+def AGPR_160 : SIRegisterTuples<getSubRegs<5>.ret, AGPR_32, 255, 1, 5, "a">;
+
+// AGPR 256-bit registers
+def AGPR_256 : SIRegisterTuples<getSubRegs<8>.ret, AGPR_32, 255, 1, 8, "a">;
+
 // AGPR 512-bit registers
 def AGPR_512 : SIRegisterTuples<getSubRegs<16>.ret, AGPR_32, 255, 1, 16, "a">;
 
 // AGPR 1024-bit registers
 def AGPR_1024 : SIRegisterTuples<getSubRegs<32>.ret, AGPR_32, 255, 1, 32, "a">;
 
 //===----------------------------------------------------------------------===//
 //  Register classes used as source and destination

 }
 
 def AReg_64 : RegisterClass<"AMDGPU", [i64, f64, v2i32, v2f32, v4f16, v4i16], 32,
                             (add AGPR_64)> {
   let Size = 64;
 
   let CopyCost = 5;
   let AllocationPriority = 2;
+  let Weight = 2;
+}
+
+def AReg_96 : RegisterClass<"AMDGPU", [v3i32, v3f32], 32,
+                             (add AGPR_96)> {
+  let Size = 96;
+
+  // Requires 3 v_accvgpr_write and 3 v_accvgpr_read to copy + burn 1 vgpr
+  let CopyCost = 7;
+  let AllocationPriority = 3;
+  let Weight = 3;
 }
 
 def AReg_128 : RegisterClass<"AMDGPU", [v4i32, v4f32, v2i64, v2f64], 32,
                              (add AGPR_128)> {
   let Size = 128;
 
   // Requires 4 v_accvgpr_write and 4 v_accvgpr_read to copy + burn 1 vgpr
   let CopyCost = 9;
   let AllocationPriority = 4;
+  let Weight = 4;
+}
+
+def AReg_160 : RegisterClass<"AMDGPU", [v5i32, v5f32], 32,
+                             (add AGPR_160)> {
+  let Size = 160;
+
+  // Requires 5 v_accvgpr_write and 5 v_accvgpr_read to copy + burn 1 vgpr
+  let CopyCost = 11;
+  let AllocationPriority = 5;
+  let Weight = 5;
+}
+
+def AReg_256 : RegisterClass<"AMDGPU", [v8i32, v8f32], 32,
+                             (add AGPR_256)> {
+  let Size = 256;
+
+  // Requires 8 v_accvgpr_write and 8 v_accvgpr_read to copy + burn 1 vgpr
+  let CopyCost = 17;
+  let AllocationPriority = 6;
+  let Weight = 8;
 }
 
 def AReg_512 : RegisterClass<"AMDGPU", [v16i32, v16f32], 32,
                              (add AGPR_512)> {
   let Size = 512;
   let CopyCost = 33;
   let AllocationPriority = 7;
+  let Weight = 16;
 }
 
 def AReg_1024 : RegisterClass<"AMDGPU", [v32i32, v32f32], 32,
                               (add AGPR_1024)> {
   let Size = 1024;
   let CopyCost = 65;
   let AllocationPriority = 8;
+  let Weight = 32;
 }
 } // End GeneratePressureSet = 0
 
 // This is not a real register. This is just to have a register to add
 // to VReg_1 that does not alias any real register that would
 // introduce inferred register classess.
 def ARTIFICIAL_VGPR : SIReg <"invalid vgpr", 0> {
   let isArtificial = 1;

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

unsigned getRegBitWidth(unsigned RCID) {
   case AMDGPU::SReg_64RegClassID:
   case AMDGPU::VReg_64RegClassID:
   case AMDGPU::AReg_64RegClassID:
   case AMDGPU::SReg_64_XEXECRegClassID:
     return 64;
   case AMDGPU::SGPR_96RegClassID:
   case AMDGPU::SReg_96RegClassID:
   case AMDGPU::VReg_96RegClassID:
+  case AMDGPU::AReg_96RegClassID:
     return 96;
   case AMDGPU::SGPR_128RegClassID:
   case AMDGPU::SReg_128RegClassID:
   case AMDGPU::VReg_128RegClassID:
   case AMDGPU::AReg_128RegClassID:
     return 128;
   case AMDGPU::SGPR_160RegClassID:
   case AMDGPU::SReg_160RegClassID:
   case AMDGPU::VReg_160RegClassID:
+  case AMDGPU::AReg_160RegClassID:
     return 160;
   case AMDGPU::SReg_256RegClassID:
   case AMDGPU::VReg_256RegClassID:
+  case AMDGPU::AReg_256RegClassID:
     return 256;
   case AMDGPU::SReg_512RegClassID:
   case AMDGPU::VReg_512RegClassID:
   case AMDGPU::AReg_512RegClassID:
     return 512;
   case AMDGPU::SReg_1024RegClassID:
   case AMDGPU::VReg_1024RegClassID:
   case AMDGPU::AReg_1024RegClassID: