[AMDGPU][GlobalISel] Check exact width in get*ClassForBitWidth and wi…

…den if necessary

Instead of checking if the given bitwidth is less or equal to a bitwidth of an existing RegClass,
check if it has the exact same value.

For LLVM vector types that don't have a corresponding Register Class, widen them during legalization.
That goes for G_EXTRACT_VECTOR_ELT, G_INSERT_VECTOR_ELT and G_BUILD_VECTOR.

Differential revision: https://reviews.llvm.org/D148096
Reviewers: foad, arsenm

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

@@ -4835,6 +4835,7 @@ LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
     Observer.changedInstr(MI);
     return Legalized;
   }
+  case TargetOpcode::G_EXTRACT_VECTOR_ELT:
   case TargetOpcode::G_EXTRACT:
     if (TypeIdx != 1)
       return UnableToLegalize;
@@ -4843,6 +4844,7 @@ LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
     Observer.changedInstr(MI);
     return Legalized;
   case TargetOpcode::G_INSERT:
+  case TargetOpcode::G_INSERT_VECTOR_ELT:
   case TargetOpcode::G_FREEZE:
   case TargetOpcode::G_FNEG:
   case TargetOpcode::G_FABS:

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

@@ -131,6 +131,28 @@ static LegalizeMutation moreEltsToNext32Bit(unsigned TypeIdx) {
   };
 }
 
+// Increase the number of vector elements to reach the next legal RegClass.
+static LegalizeMutation moreElementsToNextExistingRegClass(unsigned TypeIdx) {
+  return [=](const LegalityQuery &Query) {
+    const LLT Ty = Query.Types[TypeIdx];
+    const unsigned NumElts = Ty.getNumElements();
+    const unsigned EltSize = Ty.getElementType().getSizeInBits();
+    const unsigned MaxNumElts = MaxRegisterSize / EltSize;
+
+    assert(EltSize == 32 || EltSize == 64);
+    assert(Ty.getSizeInBits() < MaxRegisterSize);
+
+    unsigned NewNumElts;
+    // Find the nearest legal RegClass that is larger than the current type.
+    for (NewNumElts = NumElts; NewNumElts < MaxNumElts; ++NewNumElts) {
+      if (SIRegisterInfo::getSGPRClassForBitWidth(NewNumElts * EltSize))
+        break;
+    }
+
+    return std::pair(TypeIdx, LLT::fixed_vector(NewNumElts, EltSize));
+  };
+}
+
 static LLT getBitcastRegisterType(const LLT Ty) {
   const unsigned Size = Ty.getSizeInBits();
 
@@ -215,6 +237,15 @@ static LegalityPredicate isRegisterType(unsigned TypeIdx) {
   };
 }
 
+// RegisterType that doesn't have a corresponding RegClass.
+static LegalityPredicate isIllegalRegisterType(unsigned TypeIdx) {
+  return [=](const LegalityQuery &Query) {
+    LLT Ty = Query.Types[TypeIdx];
+    return isRegisterType(Ty) &&
+           !SIRegisterInfo::getSGPRClassForBitWidth(Ty.getSizeInBits());
+  };
+}
+
 static LegalityPredicate elementTypeIsLegal(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {
     const LLT QueryTy = Query.Types[TypeIdx];
@@ -1455,10 +1486,13 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
           const LLT VecTy = Query.Types[VecTypeIdx];
           const LLT IdxTy = Query.Types[IdxTypeIdx];
           const unsigned EltSize = EltTy.getSizeInBits();
+          const bool isLegalVecType =
+              !!SIRegisterInfo::getSGPRClassForBitWidth(VecTy.getSizeInBits());
           return (EltSize == 32 || EltSize == 64) &&
                   VecTy.getSizeInBits() % 32 == 0 &&
                   VecTy.getSizeInBits() <= MaxRegisterSize &&
-                  IdxTy.getSizeInBits() == 32;
+                  IdxTy.getSizeInBits() == 32 &&
+                  isLegalVecType;
         })
       .bitcastIf(all(sizeIsMultipleOf32(VecTypeIdx), scalarOrEltNarrowerThan(VecTypeIdx, 32)),
                  bitcastToVectorElement32(VecTypeIdx))
@@ -1484,6 +1518,9 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
       .clampScalar(IdxTypeIdx, S32, S32)
       .clampMaxNumElements(VecTypeIdx, S32, 32)
       // TODO: Clamp elements for 64-bit vectors?
+      .moreElementsIf(
+        isIllegalRegisterType(VecTypeIdx),
+        moreElementsToNextExistingRegClass(VecTypeIdx))
       // It should only be necessary with variable indexes.
       // As a last resort, lower to the stack
       .lower();
@@ -1538,7 +1575,10 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
     .legalForCartesianProduct(AllS64Vectors, {S64})
     .clampNumElements(0, V16S32, V32S32)
     .clampNumElements(0, V2S64, V16S64)
-    .fewerElementsIf(isWideVec16(0), changeTo(0, V2S16));
+    .fewerElementsIf(isWideVec16(0), changeTo(0, V2S16))
+    .moreElementsIf(
+      isIllegalRegisterType(0),
+      moreElementsToNextExistingRegClass(0));
 
   if (ST.hasScalarPackInsts()) {
     BuildVector