[AArch64][GlobalISel] Legalize G_STORE for v4s8 Vector Type

Author: chuongg3

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

@@ -52,6 +52,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);
@@ -422,8 +423,14 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .clampMaxNumElements(0, s64, 2)
       .clampMaxNumElements(0, p0, 2)
       .lowerIfMemSizeNotPow2()
+      // TODO: Use BITCAST for v2i8, v2i16 after G_TRUNC gets sorted out
+      .bitcastIf(typeInSet(0, {v4s8}),
+                 [=](const LegalityQuery &Query) {
+                   const LLT VecTy = Query.Types[0];
+                   return std::pair(0, LLT::scalar(VecTy.getSizeInBits()));
+                 })
       .customIf(IsPtrVecPred)
-      .scalarizeIf(typeIs(0, v2s16), 0);
+      .scalarizeIf(typeInSet(0, {v2s16, v2s8}), 0);
 
   getActionDefinitionsBuilder(G_INDEXED_STORE)
       // Idx 0 == Ptr, Idx 1 == Val

llvm/test/CodeGen/AArch64/store.ll

@@ -2,9 +2,6 @@
 ; RUN: llc -mtriple=aarch64 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-SD
 ; RUN: llc -mtriple=aarch64 -global-isel -global-isel-abort=2 %s -o - 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-GI
 
-; CHECK-GI:         warning: Instruction selection used fallback path for store_v2i8
-; CHECK-GI-NEXT:    warning: Instruction selection used fallback path for store_v4i8
-
 ; ===== Legal Scalars =====
 define void @store_i8(i8 %a, ptr %ptr){
 ; CHECK-LABEL: store_i8:
@@ -110,14 +107,22 @@ define void @store_v2i64(<2 x i64> %a, ptr %ptr){
 ; ===== Smaller/Larger Width Vectors with Legal Element Sizes =====
 
 define void @store_v2i8(<2 x i8> %a, ptr %ptr){
-; CHECK-LABEL: store_v2i8:
-; CHECK:       // %bb.0:
-; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
-; CHECK-NEXT:    mov w8, v0.s[1]
-; CHECK-NEXT:    fmov w9, s0
-; CHECK-NEXT:    strb w9, [x0]
-; CHECK-NEXT:    strb w8, [x0, #1]
-; CHECK-NEXT:    ret
+; CHECK-SD-LABEL: store_v2i8:
+; CHECK-SD:       // %bb.0:
+; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 def $q0
+; CHECK-SD-NEXT:    mov w8, v0.s[1]
+; CHECK-SD-NEXT:    fmov w9, s0
+; CHECK-SD-NEXT:    strb w9, [x0]
+; CHECK-SD-NEXT:    strb w8, [x0, #1]
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: store_v2i8:
+; CHECK-GI:       // %bb.0:
+; CHECK-GI-NEXT:    // kill: def $d0 killed $d0 def $q0
+; CHECK-GI-NEXT:    mov s1, v0.s[1]
+; CHECK-GI-NEXT:    str b0, [x0]
+; CHECK-GI-NEXT:    str b1, [x0, #1]
+; CHECK-GI-NEXT:    ret
     store <2 x i8> %a, ptr %ptr
     ret void
 }