[GlobalISel][AArch64] Legalize G_INSERT_VECTOR_ELT for SVE

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

@@ -273,6 +273,11 @@ inline LegalityPredicate typeIsNot(unsigned TypeIdx, LLT Type) {
 LegalityPredicate
 typePairInSet(unsigned TypeIdx0, unsigned TypeIdx1,
               std::initializer_list<std::pair<LLT, LLT>> TypesInit);
+/// True iff the given types for the given tuple of type indexes is one of the
+/// specified type tuple.
+LegalityPredicate
+typeTupleInSet(unsigned TypeIdx0, unsigned TypeIdx1, unsigned Type2,
+               std::initializer_list<std::tuple<LLT, LLT, LLT>> TypesInit);
 /// True iff the given types for the given pair of type indexes is one of the
 /// specified type pairs.
 LegalityPredicate typePairAndMemDescInSet(
@@ -504,6 +509,15 @@ class LegalizeRuleSet {
     using namespace LegalityPredicates;
     return actionIf(Action, typePairInSet(typeIdx(0), typeIdx(1), Types));
   }
+
+  LegalizeRuleSet &
+  actionFor(LegalizeAction Action,
+            std::initializer_list<std::tuple<LLT, LLT, LLT>> Types) {
+    using namespace LegalityPredicates;
+    return actionIf(Action,
+                    typeTupleInSet(typeIdx(0), typeIdx(1), typeIdx(2), Types));
+  }
+
   /// Use the given action when type indexes 0 and 1 is any type pair in the
   /// given list.
   /// Action should be an action that requires mutation.
@@ -615,6 +629,12 @@ class LegalizeRuleSet {
       return *this;
     return actionFor(LegalizeAction::Legal, Types);
   }
+  LegalizeRuleSet &
+  legalFor(bool Pred, std::initializer_list<std::tuple<LLT, LLT, LLT>> Types) {
+    if (!Pred)
+      return *this;
+    return actionFor(LegalizeAction::Legal, Types);
+  }
   /// The instruction is legal when type index 0 is any type in the given list
   /// and imm index 0 is anything.
   LegalizeRuleSet &legalForTypeWithAnyImm(std::initializer_list<LLT> Types) {

llvm/lib/CodeGen/GlobalISel/LegalityPredicates.cpp

@@ -49,6 +49,17 @@ LegalityPredicate LegalityPredicates::typePairInSet(
   };
 }
 
+LegalityPredicate LegalityPredicates::typeTupleInSet(
+    unsigned TypeIdx0, unsigned TypeIdx1, unsigned TypeIdx2,
+    std::initializer_list<std::tuple<LLT, LLT, LLT>> TypesInit) {
+  SmallVector<std::tuple<LLT, LLT, LLT>, 4> Types = TypesInit;
+  return [=](const LegalityQuery &Query) {
+    std::tuple<LLT, LLT, LLT> Match = {
+        Query.Types[TypeIdx0], Query.Types[TypeIdx1], Query.Types[TypeIdx2]};
+    return llvm::is_contained(Types, Match);
+  };
+}
+
 LegalityPredicate LegalityPredicates::typePairAndMemDescInSet(
     unsigned TypeIdx0, unsigned TypeIdx1, unsigned MMOIdx,
     std::initializer_list<TypePairAndMemDesc> TypesAndMemDescInit) {
@@ -202,7 +213,7 @@ LegalityPredicate LegalityPredicates::memSizeNotByteSizePow2(unsigned MMOIdx) {
 LegalityPredicate LegalityPredicates::numElementsNotPow2(unsigned TypeIdx) {
   return [=](const LegalityQuery &Query) {
     const LLT QueryTy = Query.Types[TypeIdx];
-    return QueryTy.isVector() && !isPowerOf2_32(QueryTy.getNumElements());
+    return QueryTy.isFixedVector() && !isPowerOf2_32(QueryTy.getNumElements());
   };
 }
 

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

@@ -978,6 +978,10 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
   getActionDefinitionsBuilder(G_INSERT_VECTOR_ELT)
       .legalIf(
           typeInSet(0, {v16s8, v8s8, v8s16, v4s16, v4s32, v2s32, v2s64, v2p0}))
+      .legalFor(HasSVE, {{nxv16s8, s32, s64},
+                         {nxv8s16, s32, s64},
+                         {nxv4s32, s32, s64},
+                         {nxv2s64, s64, s64}})
       .moreElementsToNextPow2(0)
       .widenVectorEltsToVectorMinSize(0, 64)
       .clampNumElements(0, v8s8, v16s8)

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

@@ -426,6 +426,9 @@ void applyNonConstInsert(MachineInstr &MI, MachineRegisterInfo &MRI,
   LLT EltTy = MRI.getType(Insert.getElementReg());
   LLT IdxTy = MRI.getType(Insert.getIndexReg());
 
+  if (VecTy.isScalableVector())
+    return;
+
   // Create a stack slot and store the vector into it
   MachineFunction &MF = Builder.getMF();
   Align Alignment(

llvm/test/CodeGen/AArch64/extract-vector-elt-sve.ll

@@ -0,0 +1,123 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc -mtriple=aarch64 -mattr=+sve -verify-machineinstrs %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-SD
+; RUN: llc -mtriple=aarch64 -mattr=+sve -aarch64-enable-gisel-sve=1 -global-isel -global-isel-abort=2 -verify-machineinstrs %s -o - 2>&1 | FileCheck %s --check-prefixes=CHECK,CHECK-GI
+
+define <vscale x 2 x i64> @insert_vscale_2_i64_zero(<vscale x 2 x i64> %vec, i64 %elt) {
+; CHECK-SD-LABEL: insert_vscale_2_i64_zero:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    ptrue p0.d, vl1
+; CHECK-SD-NEXT:    mov z0.d, p0/m, x0
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: insert_vscale_2_i64_zero:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    mov x8, xzr
+; CHECK-GI-NEXT:    index z1.d, #0, #1
+; CHECK-GI-NEXT:    ptrue p0.d
+; CHECK-GI-NEXT:    mov z2.d, x8
+; CHECK-GI-NEXT:    cmpeq p0.d, p0/z, z1.d, z2.d
+; CHECK-GI-NEXT:    mov z0.d, p0/m, x0
+; CHECK-GI-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 2 x i64> %vec, i64 %elt, i64 0
+  ret <vscale x 2 x i64> %d
+}
+
+define <vscale x 2 x i64> @insert_vscale_2_i64(<vscale x 2 x i64> %vec, i64 %elt, i64 %idx) {
+; CHECK-LABEL: insert_vscale_2_i64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z1.d, #0, #1
+; CHECK-NEXT:    mov z2.d, x1
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    cmpeq p0.d, p0/z, z1.d, z2.d
+; CHECK-NEXT:    mov z0.d, p0/m, x0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x  2 x i64> %vec, i64 %elt, i64 %idx
+  ret <vscale x 2 x i64> %d
+}
+
+define <vscale x 4 x i32> @insert_vscale_4_i32_zero(<vscale x 4 x i32> %vec, i32 %elt) {
+; CHECK-SD-LABEL: insert_vscale_4_i32_zero:
+; CHECK-SD:       // %bb.0: // %entry
+; CHECK-SD-NEXT:    ptrue p0.s, vl1
+; CHECK-SD-NEXT:    mov z0.s, p0/m, w0
+; CHECK-SD-NEXT:    ret
+;
+; CHECK-GI-LABEL: insert_vscale_4_i32_zero:
+; CHECK-GI:       // %bb.0: // %entry
+; CHECK-GI-NEXT:    mov w8, wzr
+; CHECK-GI-NEXT:    index z1.s, #0, #1
+; CHECK-GI-NEXT:    ptrue p0.s
+; CHECK-GI-NEXT:    mov z2.s, w8
+; CHECK-GI-NEXT:    cmpeq p0.s, p0/z, z1.s, z2.s
+; CHECK-GI-NEXT:    mov z0.s, p0/m, w0
+; CHECK-GI-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 4 x i32> %vec, i32 %elt, i64 0
+  ret <vscale x 4 x i32> %d
+}
+
+define <vscale x 4 x i32> @insert_vscale_4_i32(<vscale x 4 x i32> %vec, i32 %elt, i64 %idx) {
+; CHECK-LABEL: insert_vscale_4_i32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z1.s, #0, #1
+; CHECK-NEXT:    mov z2.s, w1
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmpeq p0.s, p0/z, z1.s, z2.s
+; CHECK-NEXT:    mov z0.s, p0/m, w0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 4 x i32> %vec, i32 %elt, i64 %idx
+  ret <vscale x 4 x i32> %d
+}
+
+define <vscale x 8 x i16> @insert_vscale_8_i16_zero(<vscale x 8 x i16> %vec, i16 %elt) {
+; CHECK-LABEL: insert_vscale_8_i16_zero:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.h, vl1
+; CHECK-NEXT:    mov z0.h, p0/m, w0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 8 x i16> %vec, i16 %elt, i64 0
+  ret <vscale x 8 x i16> %d
+}
+
+define <vscale x 8 x i16> @insert_vscale_8_i16(<vscale x 8 x i16> %vec, i16 %elt, i64 %idx) {
+; CHECK-LABEL: insert_vscale_8_i16:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z1.h, #0, #1
+; CHECK-NEXT:    mov z2.h, w1
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    cmpeq p0.h, p0/z, z1.h, z2.h
+; CHECK-NEXT:    mov z0.h, p0/m, w0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 8 x i16> %vec, i16 %elt, i64 %idx
+  ret <vscale x 8 x i16> %d
+}
+
+define <vscale x 16 x i8> @insert_vscale_16_i8_zero(<vscale x 16 x i8> %vec, i8 %elt) {
+; CHECK-LABEL: insert_vscale_16_i8_zero:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.b, vl1
+; CHECK-NEXT:    mov z0.b, p0/m, w0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 16 x i8> %vec, i8 %elt, i64 0
+  ret <vscale x 16 x i8> %d
+}
+
+define <vscale x 16 x i8> @insert_vscale_16_i8(<vscale x 16 x i8> %vec, i8 %elt, i64 %idx) {
+; CHECK-LABEL: insert_vscale_16_i8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z1.b, #0, #1
+; CHECK-NEXT:    mov z2.b, w1
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    cmpeq p0.b, p0/z, z1.b, z2.b
+; CHECK-NEXT:    mov z0.b, p0/m, w0
+; CHECK-NEXT:    ret
+entry:
+  %d = insertelement <vscale x 16 x i8> %vec, i8 %elt, i64 %idx
+  ret <vscale x 16 x i8> %d
+}

llvm/test/CodeGen/AArch64/extract-vector-elt.ll

@@ -1099,5 +1099,3 @@ loop:
 ret:
   ret i32 %3
 }
-
-