[GlobalISel] Use-Vector-Truncate Opt Needs Elt Type Check (#146003)

In the pre-legalizer combiner, there exists a bug with UseVectorTruncate
match-apply optimization. When the destinations' types do not match the
vector element type of the G_UNMERGE_VALUES instruction, the resulting
collapsed truncate does not preserve original functional behavior. This
commit introduces a simple type check to ensure that the destination
types match the vector element type.

Author: daniel-man

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -3487,6 +3487,13 @@ bool CombinerHelper::matchUseVectorTruncate(MachineInstr &MI,
   if (!DstTy.getElementCount().isKnownMultipleOf(UnmergeSrcTy.getNumElements()))
     return false;
 
+  // Check the unmerge source and destination element types match
+  LLT UnmergeSrcEltTy = UnmergeSrcTy.getElementType();
+  Register UnmergeDstReg = UnmergeMI->getOperand(0).getReg();
+  LLT UnmergeDstEltTy = MRI.getType(UnmergeDstReg);
+  if (UnmergeSrcEltTy != UnmergeDstEltTy)
+    return false;
+
   // Only generate legal instructions post-legalizer
   if (!IsPreLegalize) {
     LLT MidTy = DstTy.changeElementType(UnmergeSrcTy.getScalarType());

llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir

@@ -0,0 +1,52 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=aarch64-prelegalizer-combiner %s -o - | FileCheck %s
+---
+name:            test_foldable_s32
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: test_foldable_s32
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[SRC:%src]]:_(<2 x s32>) = COPY $x1
+    ; CHECK-NEXT: [[UNDEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[MID:%[0-9]+]]:_(<4 x s32>) = G_CONCAT_VECTORS [[SRC]](<2 x s32>), [[UNDEF]](<2 x s32>)
+    ; CHECK-NEXT: [[DST:%dst]]:_(<4 x s16>) = G_TRUNC [[MID]]
+    ; CHECK-NEXT: $x0 = COPY [[DST]]
+
+    %src:_(<2 x s32>) = COPY $x1
+    %a:_(s32), %b:_(s32) = G_UNMERGE_VALUES %src:_(<2 x s32>)
+    %T_a:_(s16) = G_TRUNC %a:_(s32)
+    %T_b:_(s16) = G_TRUNC %b:_(s32)
+    %Undef:_(s16) = G_IMPLICIT_DEF
+    %dst:_(<4 x s16>) = G_BUILD_VECTOR %T_a:_(s16), %T_b:_(s16), %Undef:_(s16), %Undef:_(s16)
+    $x0 = COPY %dst
+...
+---
+name:            test_unfoldable_s32
+body:             |
+  bb.0:
+    liveins: $w0, $x1
+
+    ; CHECK-LABEL: name: test_unfoldable_s32
+    ; CHECK: liveins: $w0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[SRC:%src]]:_(<2 x s32>) = COPY $x1
+    ; CHECK-NEXT: [[A:%a]]:_(s16), [[B:%b]]:_(s16), [[C:%c]]:_(s16), %d:_(s16) = G_UNMERGE_VALUES [[SRC]](<2 x s32>)
+    ; CHECK-NEXT: [[T_A:%T_a]]:_(s8) = G_TRUNC [[A]](s16)
+    ; CHECK-NEXT: [[T_B:%T_b]]:_(s8) = G_TRUNC [[B]](s16)
+    ; CHECK-NEXT: [[T_C:%T_c]]:_(s8) = G_TRUNC [[C]](s16)
+    ; CHECK-NEXT: [[UNDEF:%Undef]]:_(s8) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[DST:%dst]]:_(<4 x s8>) = G_BUILD_VECTOR [[T_A]](s8), [[T_B]](s8), [[T_C]](s8), [[UNDEF]](s8)
+    ; CHECK-NEXT: $w0 = COPY [[DST]](<4 x s8>)
+    %src:_(<2 x s32>) = COPY $x1
+    %a:_(s16), %b:_(s16), %c:_(s16), %d:_(s16) = G_UNMERGE_VALUES %src:_(<2 x s32>)
+    %T_a:_(s8) = G_TRUNC %a:_(s16)
+    %T_b:_(s8) = G_TRUNC %b:_(s16)
+    %T_c:_(s8) = G_TRUNC %c:_(s16)
+    %Undef:_(s8) = G_IMPLICIT_DEF
+    %dst:_(<4 x s8>) = G_BUILD_VECTOR %T_a:_(s8), %T_b:_(s8), %T_c:_(s8), %Undef:_(s8)
+    $w0 = COPY %dst
+
+...