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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-globalisel

Author: Daniel Man (daniel-man)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/146003.diff


2 Files Affected:

- (modified) llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp (+7) 
- (added) llvm/test/CodeGen/AArch64/GlobalISel/prelegalizer-combiner-use-vector-truncate.mir (+50) 


``````````diff
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/146003