[AArch64][GlobalISel] Legalize <2 x s8> and <4 x s8> for G_BUILD_VECTOR

Refer to commit ccffc27, the remaining types <2 x s8> and <4 x s8> should
also be promoted to <2 x s32> and <4 x s16>.

Fixes #58274

Reviewed By: aemerson, tschuett, paquette, dmgreen
Differential Revision: https://reviews.llvm.org/D153394

Author: vfdff

llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

@@ -941,6 +941,28 @@ class LegalizeRuleSet {
                     changeElementTo(typeIdx(TypeIdx), Ty));
   }
 
+  /// Ensure the vector size is at least as wide as VectorSize by promoting the
+  /// element.
+  LegalizeRuleSet &widenVectorEltsToVectorMinSize(unsigned TypeIdx,
+                                                  unsigned VectorSize) {
+    using namespace LegalityPredicates;
+    using namespace LegalizeMutations;
+    return actionIf(
+        LegalizeAction::WidenScalar,
+        [=](const LegalityQuery &Query) {
+          const LLT VecTy = Query.Types[TypeIdx];
+          return VecTy.isVector() && !VecTy.isScalable() &&
+                 VecTy.getSizeInBits() < VectorSize;
+        },
+        [=](const LegalityQuery &Query) {
+          const LLT VecTy = Query.Types[TypeIdx];
+          unsigned NumElts = VecTy.getNumElements();
+          unsigned MinSize = VectorSize / NumElts;
+          LLT NewTy = LLT::fixed_vector(NumElts, LLT::scalar(MinSize));
+          return std::make_pair(TypeIdx, NewTy);
+        });
+  }
+
   /// Ensure the scalar is at least as wide as Ty.
   LegalizeRuleSet &minScalar(unsigned TypeIdx, const LLT Ty) {
     using namespace LegalityPredicates;

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

@@ -50,6 +50,7 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
   const LLT v16s8 = LLT::fixed_vector(16, 8);
   const LLT v8s8 = LLT::fixed_vector(8, 8);
   const LLT v4s8 = LLT::fixed_vector(4, 8);
+  const LLT v2s8 = LLT::fixed_vector(2, 8);
   const LLT v8s16 = LLT::fixed_vector(8, 16);
   const LLT v4s16 = LLT::fixed_vector(4, 16);
   const LLT v2s16 = LLT::fixed_vector(2, 16);
@@ -720,9 +721,7 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .clampNumElements(0, v4s32, v4s32)
       .clampNumElements(0, v2s64, v2s64)
       .minScalarOrElt(0, s8)
-      .minScalarOrEltIf(
-          [=](const LegalityQuery &Query) { return Query.Types[0] == v2s16; },
-          0, s32)
+      .widenVectorEltsToVectorMinSize(0, 64)
       .minScalarSameAs(1, 0);
 
   getActionDefinitionsBuilder(G_BUILD_VECTOR_TRUNC).lower();

llvm/test/CodeGen/AArch64/GlobalISel/legalize-build-vector.mir

@@ -133,3 +133,43 @@ body: |
     $d0 = COPY %ext(<2 x s32>)
     RET_ReallyLR
 ...
+
+---
+name:            widen_v2s8
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: widen_v2s8
+    ; CHECK: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY [[DEF]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[DEF]](s32)
+    ; CHECK-NEXT: $d0 = COPY [[BUILD_VECTOR]](<2 x s32>)
+    ; CHECK-NEXT: RET_ReallyLR
+    %0:_(s8) = G_IMPLICIT_DEF
+    %1:_(s8) = G_IMPLICIT_DEF
+    %2:_(<2 x s8>) = G_BUILD_VECTOR %0(s8), %1(s8)
+    %ext:_(<2 x s32>) = G_ANYEXT %2(<2 x s8>)
+    $d0 = COPY %ext(<2 x s32>)
+    RET_ReallyLR
+...
+
+---
+name:            widen_v4s8
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: widen_v4s8
+    ; CHECK: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s16) = COPY [[DEF]](s16)
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s16) = COPY [[DEF]](s16)
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:_(s16) = COPY [[DEF]](s16)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR [[COPY]](s16), [[COPY1]](s16), [[COPY2]](s16), [[DEF]](s16)
+    ; CHECK-NEXT: $d0 = COPY [[BUILD_VECTOR]](<4 x s16>)
+    ; CHECK-NEXT: RET_ReallyLR
+    %0:_(s8) = G_IMPLICIT_DEF
+    %1:_(s8) = G_IMPLICIT_DEF
+    %2:_(s8) = G_IMPLICIT_DEF
+    %3:_(s8) = G_IMPLICIT_DEF
+    %4:_(<4 x s8>) = G_BUILD_VECTOR %0(s8), %1(s8), %2(s8), %3(s8)
+    %ext:_(<4 x s16>) = G_ANYEXT %4(<4 x s8>)
+    $d0 = COPY %ext(<4 x s16>)
+    RET_ReallyLR
+...

llvm/test/CodeGen/AArch64/GlobalISel/legalize-itofp.mir

@@ -274,11 +274,12 @@ body: |
     ; CHECK-LABEL: name: test_uitofp_v2s64_v2i1
     ; CHECK: liveins: $q0
     ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s8) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s8>) = G_BUILD_VECTOR [[DEF]](s8), [[DEF]](s8)
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY [[DEF]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[DEF]](s32)
     ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
     ; CHECK-NEXT: [[BUILD_VECTOR1:%[0-9]+]]:_(<2 x s64>) = G_BUILD_VECTOR [[C]](s64), [[C]](s64)
-    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<2 x s64>) = G_ANYEXT [[BUILD_VECTOR]](<2 x s8>)
+    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<2 x s64>) = G_ANYEXT [[BUILD_VECTOR]](<2 x s32>)
     ; CHECK-NEXT: [[AND:%[0-9]+]]:_(<2 x s64>) = G_AND [[ANYEXT]], [[BUILD_VECTOR1]]
     ; CHECK-NEXT: [[UITOFP:%[0-9]+]]:_(<2 x s64>) = G_UITOFP [[AND]](<2 x s64>)
     ; CHECK-NEXT: $q0 = COPY [[UITOFP]](<2 x s64>)
@@ -295,9 +296,10 @@ body: |
     ; CHECK-LABEL: name: test_sitofp_v2s64_v2i1
     ; CHECK: liveins: $q0
     ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s8) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s8>) = G_BUILD_VECTOR [[DEF]](s8), [[DEF]](s8)
-    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<2 x s64>) = G_ANYEXT [[BUILD_VECTOR]](<2 x s8>)
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY [[DEF]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[DEF]](s32)
+    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<2 x s64>) = G_ANYEXT [[BUILD_VECTOR]](<2 x s32>)
     ; CHECK-NEXT: [[SEXT_INREG:%[0-9]+]]:_(<2 x s64>) = G_SEXT_INREG [[ANYEXT]], 1
     ; CHECK-NEXT: [[SITOFP:%[0-9]+]]:_(<2 x s64>) = G_SITOFP [[SEXT_INREG]](<2 x s64>)
     ; CHECK-NEXT: $q0 = COPY [[SITOFP]](<2 x s64>)