From ade6fa46f94b31e89c8a488264ac79e319d1ccdb Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara@apple.com>
Date: Thu, 24 Sep 2020 13:06:03 -0700
Subject: [PATCH] [AArch64][GlobalISel] Make <8 x s16> for G_INSERT_VECTOR_ELT
 legal.

---
 .../AArch64/GISel/AArch64LegalizerInfo.cpp    |  6 +-
 .../GlobalISel/legalize-insert-vector-elt.mir | 79 +++++++++++++++++++
 .../GlobalISel/select-insert-vector-elt.mir   | 27 +++++++
 3 files changed, 107 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index 43f3c1d31da5..5274d643e624 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -597,11 +597,7 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .minScalarOrElt(0, s8); // Worst case, we need at least s8.
 
   getActionDefinitionsBuilder(G_INSERT_VECTOR_ELT)
-      .legalIf([=](const LegalityQuery &Query) {
-        const LLT &VecTy = Query.Types[0];
-        // TODO: Support s8 and s16
-        return VecTy == v2s32 || VecTy == v4s32 || VecTy == v2s64;
-      });
+      .legalIf(typeInSet(0, {v8s16, v2s32, v4s32, v2s64}));
 
   getActionDefinitionsBuilder(G_BUILD_VECTOR)
       .legalFor({{v8s8, s8},
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir
new file mode 100644
index 000000000000..328935b1a6ca
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir
@@ -0,0 +1,79 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-linux-gnu -O0 -run-pass=legalizer %s -o - -global-isel-abort=1 | FileCheck %s
+
+---
+name:            v8s16
+body: |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: v8s16
+    ; CHECK: [[COPY:%[0-9]+]]:_(<8 x s16>) = COPY $q0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %val:_(s16) = G_CONSTANT i16 42
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<8 x s16>) = G_INSERT_VECTOR_ELT [[COPY]], %val(s16), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<8 x s16>)
+    ; CHECK: RET_ReallyLR
+    %0:_(<8 x s16>) = COPY $q0
+    %1:_(s32) = G_CONSTANT i32 1
+    %val:_(s16) = G_CONSTANT i16 42
+    %2:_(<8 x s16>) = G_INSERT_VECTOR_ELT %0(<8 x s16>), %val(s16), %1(s32)
+    $q0 = COPY %2(<8 x s16>)
+    RET_ReallyLR
+...
+---
+name:            v2s32
+body: |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: v2s32
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %val:_(s32) = G_CONSTANT i32 42
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<2 x s32>) = G_INSERT_VECTOR_ELT [[COPY]], %val(s32), [[C]](s32)
+    ; CHECK: $d0 = COPY [[IVEC]](<2 x s32>)
+    ; CHECK: RET_ReallyLR
+    %0:_(<2 x s32>) = COPY $d0
+    %1:_(s32) = G_CONSTANT i32 1
+    %val:_(s32) = G_CONSTANT i32 42
+    %2:_(<2 x s32>) = G_INSERT_VECTOR_ELT %0(<2 x s32>), %val(s32), %1(s32)
+    $d0 = COPY %2(<2 x s32>)
+    RET_ReallyLR
+...
+---
+name:            v4s32
+body: |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: v4s32
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %val:_(s32) = G_CONSTANT i32 42
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<4 x s32>) = G_INSERT_VECTOR_ELT [[COPY]], %val(s32), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<4 x s32>)
+    ; CHECK: RET_ReallyLR
+    %0:_(<4 x s32>) = COPY $q0
+    %1:_(s32) = G_CONSTANT i32 1
+    %val:_(s32) = G_CONSTANT i32 42
+    %2:_(<4 x s32>) = G_INSERT_VECTOR_ELT %0(<4 x s32>), %val(s32), %1(s32)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR
+...
+---
+name:            v2s64
+body: |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: v2s64
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s64>) = COPY $q0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %val:_(s64) = G_CONSTANT i64 42
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<2 x s64>) = G_INSERT_VECTOR_ELT [[COPY]], %val(s64), [[C]](s32)
+    ; CHECK: $q0 = COPY [[IVEC]](<2 x s64>)
+    ; CHECK: RET_ReallyLR
+    %0:_(<2 x s64>) = COPY $q0
+    %1:_(s32) = G_CONSTANT i32 1
+    %val:_(s64) = G_CONSTANT i64 42
+    %2:_(<2 x s64>) = G_INSERT_VECTOR_ELT %0(<2 x s64>), %val(s64), %1(s32)
+    $q0 = COPY %2(<2 x s64>)
+    RET_ReallyLR
+...
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-insert-vector-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-insert-vector-elt.mir
index 7890a4c846ad..5c4a2e1c3544 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/select-insert-vector-elt.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-insert-vector-elt.mir
@@ -1,6 +1,33 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -verify-machineinstrs -mtriple aarch64-unknown-unknown -run-pass=instruction-select %s -o - | FileCheck %s
+---
+name:            v8s16_fpr
+alignment:       4
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $q1, $h0
 
+    ; CHECK-LABEL: name: v8s16_fpr
+    ; CHECK: liveins: $q1, $h0
+    ; CHECK: [[COPY:%[0-9]+]]:fpr16 = COPY $h0
+    ; CHECK: [[COPY1:%[0-9]+]]:fpr128 = COPY $q1
+    ; CHECK: [[DEF:%[0-9]+]]:fpr128 = IMPLICIT_DEF
+    ; CHECK: [[INSERT_SUBREG:%[0-9]+]]:fpr128 = INSERT_SUBREG [[DEF]], [[COPY]], %subreg.hsub
+    ; CHECK: [[INSvi16lane:%[0-9]+]]:fpr128 = INSvi16lane [[COPY1]], 1, [[INSERT_SUBREG]], 0
+    ; CHECK: $q0 = COPY [[INSvi16lane]]
+    ; CHECK: RET_ReallyLR implicit $q0
+    %0:fpr(s16) = COPY $h0
+    %1:fpr(<8 x s16>) = COPY $q1
+    %3:gpr(s32) = G_CONSTANT i32 1
+    %2:fpr(<8 x s16>) = G_INSERT_VECTOR_ELT %1, %0(s16), %3(s32)
+    $q0 = COPY %2(<8 x s16>)
+    RET_ReallyLR implicit $q0
+
+...
+---
 name:            v4s32_fpr
 alignment:       4
 legalized:       true