[llvm] [GIsel][AArch64] Legalize <2 x i16> for G_INSERT_VECTOR_ELT (PR #65830)

[via llvm-commits]

https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/65830:

>From a263076456028bb7d3d014ab719ac45192c6e2b4 Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Thu, 7 Sep 2023 23:20:10 -0400
Subject: [PATCH] [GIsel][AArch64] Legalize <2 x i16> for G_INSERT_VECTOR_ELT

Widen the vector elements to 64 bits to make sure it legal instead by
clamping the number of elements. Depend on D153394.

Fixes https://github.com/llvm/llvm-project/issues/63826
---
 .../CodeGen/GlobalISel/LegalizerHelper.cpp    | 14 +++++++++
 .../AArch64/GISel/AArch64LegalizerInfo.cpp    |  5 +--
 .../GlobalISel/legalize-insert-vector-elt.mir | 31 ++++++++++---------
 .../GlobalISel/select-bitfield-insert.ll      | 12 +++++++
 4 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index cd08bf2c1a726b6..6ef56699379b4ce 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -2493,6 +2493,20 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
     return Legalized;
   }
   case TargetOpcode::G_INSERT_VECTOR_ELT: {
+    if (TypeIdx == 0) {
+      Observer.changingInstr(MI);
+
+      Register VecReg = MI.getOperand(1).getReg();
+      LLT VecTy = MRI.getType(VecReg);
+      const LLT WideEltTy = WideTy.getElementType();
+
+      widenScalarSrc(MI, WideTy, 1, TargetOpcode::G_ANYEXT);
+      widenScalarSrc(MI, WideEltTy, 2, TargetOpcode::G_ANYEXT);
+      widenScalarDst(MI, WideTy, 0);
+      Observer.changedInstr(MI);
+      return Legalized;
+    }
+
     if (TypeIdx == 1) {
       Observer.changingInstr(MI);
 
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index 00424f23e96d7d9..3b4a4c61ab56164 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -718,8 +718,9 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
 
   getActionDefinitionsBuilder(G_INSERT_VECTOR_ELT)
       .legalIf(typeInSet(0, {v16s8, v8s8, v8s16, v4s16, v4s32, v2s32, v2s64}))
-      .clampMinNumElements(0, s16, 4)
-      .clampMaxNumElements(0, s16, 8);
+      .minScalarOrEltIf(
+          [=](const LegalityQuery &Query) { return Query.Types[0] == v2s16; },
+          0, s32);
 
   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
index 26db18bd611a57c..a077b67f2069d13 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-insert-vector-elt.mir
@@ -4,21 +4,24 @@
 name:            pr63826
 body:             |
   bb.0:
+    liveins: $d0
     ; CHECK-LABEL: name: pr63826
-    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s16>) = COPY $w0
-    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s16) = G_CONSTANT i16 1
-    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 42
-    ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[COPY]](<2 x s16>)
-    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(s16) = G_IMPLICIT_DEF
-    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s16>) = G_BUILD_VECTOR [[UV]](s16), [[UV1]](s16), [[DEF]](s16), [[DEF]](s16)
-    ; CHECK-NEXT: [[IVEC:%[0-9]+]]:_(<4 x s16>) = G_INSERT_VECTOR_ELT [[BUILD_VECTOR]], [[C]](s16), [[C1]](s32)
-    ; CHECK-NEXT: [[UV2:%[0-9]+]]:_(<2 x s16>), [[UV3:%[0-9]+]]:_(<2 x s16>) = G_UNMERGE_VALUES [[IVEC]](<4 x s16>)
-    ; CHECK-NEXT: $w0 = COPY [[UV2]](<2 x s16>)
-    %0:_(<2 x s16>) = COPY $w0
-    %1:_(s16) = G_CONSTANT i16 1
-    %2:_(s32) = G_CONSTANT i32 42
-    %4:_(<2 x s16>) = G_INSERT_VECTOR_ELT %0(<2 x s16>), %1(s16), %2(s32)
-    $w0 = COPY %4(<2 x s16>)
+    ; CHECK: liveins: $d0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[IVEC:%[0-9]+]]:_(<2 x s32>) = G_INSERT_VECTOR_ELT [[COPY]], [[C1]](s32), [[C]](s32)
+    ; CHECK-NEXT: $d0 = COPY [[IVEC]](<2 x s32>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $d0
+    %1:_(<2 x s32>) = COPY $d0
+    %0:_(<2 x s16>) = G_TRUNC %1(<2 x s32>)
+    %4:_(s32) = G_CONSTANT i32 0
+    %3:_(s16) = G_CONSTANT i16 1
+    %2:_(<2 x s16>) = G_INSERT_VECTOR_ELT %0, %3(s16), %4(s32)
+    %5:_(<2 x s32>) = G_ANYEXT %2(<2 x s16>)
+    $d0 = COPY %5(<2 x s32>)
+    RET_ReallyLR implicit $d0
 ...
 ---
 name:            v8s8
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-bitfield-insert.ll b/llvm/test/CodeGen/AArch64/GlobalISel/select-bitfield-insert.ll
index e1df07c93ebf1f2..8ffa419c602eb31 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/select-bitfield-insert.ll
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-bitfield-insert.ll
@@ -155,3 +155,15 @@ bb:
   store i64 %tmp4, ptr %p
   ret i64 %out
 }
+
+define <2 x i16> @pr63826(<2 x i16> %vec) {
+; CHECK-LABEL: pr63826:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    ; kill: def $d0 killed $d0 def $q0
+; CHECK-NEXT:    mov w8, #1 ; =0x1
+; CHECK-NEXT:    mov.s v0[0], w8
+; CHECK-NEXT:    ; kill: def $d0 killed $d0 killed $q0
+; CHECK-NEXT:    ret
+  %vec1 = insertelement <2 x i16> %vec, i16 1, i32 0
+  ret <2 x i16> %vec1
+}