[llvm] [FIX] Use-Vector-Truncate Opt Needs Elt Type Check (PR #146003)

[Daniel Man via llvm-commits]

https://github.com/daniel-man updated https://github.com/llvm/llvm-project/pull/146003

>From 15e72aedcc1aa8e6ed44c540d97f2ee1bd735c3c Mon Sep 17 00:00:00 2001
From: "Man, Daniel" <daniel.man at intel.com>
Date: Thu, 26 Jun 2025 18:44:18 -0700
Subject: [PATCH 1/2] [FIX] Use-Vector-Truncate Opt Needs Elt Type Check

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.
---
 .../lib/CodeGen/GlobalISel/CombinerHelper.cpp |  7 +++
 ...legalizer-combiner-use-vector-truncate.mir | 50 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index b1e851183de0d..05dd269d48921 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/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());
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir b/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir
new file mode 100644
index 0000000000000..278f7e71f3d8b
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir
@@ -0,0 +1,50 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=aarch64-prelegalizer-combiner -global-isel -verify-machineinstrs %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

>From 5852305282e3d0f29c5c199b0f33bc16c922098c Mon Sep 17 00:00:00 2001
From: "Man, Daniel" <daniel.man at intel.com>
Date: Thu, 26 Jun 2025 23:01:26 -0700
Subject: [PATCH 2/2] PR Fixes

---
 .../GlobalISel/prelegalizer-combiner-use-vector-truncate.mir  | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir b/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir
index 278f7e71f3d8b..3d0d591bccc17 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir
@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=aarch64 -run-pass=aarch64-prelegalizer-combiner -global-isel -verify-machineinstrs %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64 -run-pass=aarch64-prelegalizer-combiner %s -o - | FileCheck %s
 ---
 name:            test_foldable_s32
 body:             |
@@ -48,3 +48,5 @@ body:             |
     %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
+
+...