[llvm] [AArch64][GlobalISel] Fix matchUseVectorTruncate to check element order before folding (PR #185834)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-globalisel

@llvm/pr-subscribers-backend-aarch64

Author: Igor Gorban (igorban-intel)

<details>
<summary>Changes</summary>

matchUseVectorTruncate incorrectly folded a
G_UNMERGE_VALUES + G_TRUNC + G_BUILD_VECTOR sequence into a single
vector G_TRUNC even when the BUILD_VECTOR operand order did not match
the UNMERGE_VALUES result order. This silently dropped element
reordering, producing a miscompile.

Add an element-ordering check: BUILD_VECTOR position I must use
UNMERGE_VALUES result I, otherwise the fold is rejected. Also
extract the repeated `SrcMI->getOperand(1).getReg()` into a local
`TruncSrcReg` variable for clarity.

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp (+9-2) 
- (added) llvm/test/CodeGen/AArch64/GlobalISel/combine-use-vector-truncate-order.mir (+58) 


``````````diff
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index f88b3f487e6cf..1e3ffb4b6e4ed 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -3596,21 +3596,28 @@ bool CombinerHelper::matchUseVectorTruncate(MachineInstr &MI,
   MachineInstr *UnmergeMI = nullptr;
 
   // Check all source TRUNCs come from the same UNMERGE instruction
+  // and that the element order matches (BUILD_VECTOR position I
+  // corresponds to UNMERGE result I)
   for (I = 0; I < NumOperands; ++I) {
     auto SrcMI = MRI.getVRegDef(BuildMI->getSourceReg(I));
     auto SrcMIOpc = SrcMI->getOpcode();
 
     // Check if the G_TRUNC instructions all come from the same MI
     if (SrcMIOpc == TargetOpcode::G_TRUNC) {
+      Register TruncSrcReg = SrcMI->getOperand(1).getReg();
       if (!UnmergeMI) {
-        UnmergeMI = MRI.getVRegDef(SrcMI->getOperand(1).getReg());
+        UnmergeMI = MRI.getVRegDef(TruncSrcReg);
         if (UnmergeMI->getOpcode() != TargetOpcode::G_UNMERGE_VALUES)
           return false;
       } else {
-        auto UnmergeSrcMI = MRI.getVRegDef(SrcMI->getOperand(1).getReg());
+        auto UnmergeSrcMI = MRI.getVRegDef(TruncSrcReg);
         if (UnmergeMI != UnmergeSrcMI)
           return false;
       }
+      // Verify element ordering: BUILD_VECTOR position I must use
+      // UNMERGE result I, otherwise the fold would lose element reordering
+      if (UnmergeMI->getOperand(I).getReg() != TruncSrcReg)
+        return false;
     } else {
       break;
     }
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-use-vector-truncate-order.mir b/llvm/test/CodeGen/AArch64/GlobalISel/combine-use-vector-truncate-order.mir
new file mode 100644
index 0000000000000..5b7748a247a95
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/combine-use-vector-truncate-order.mir
@@ -0,0 +1,58 @@
+# RUN: llc -mtriple=aarch64 -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+---
+name:            test_use_vector_truncate_reversed_order
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: test_use_vector_truncate_reversed_order
+    ; CHECK: liveins: $d0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<4 x s16>) = COPY $d0
+    ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16), [[UV2:%[0-9]+]]:_(s16), [[UV3:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[COPY]](<4 x s16>)
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s8) = G_TRUNC [[UV]](s16)
+    ; CHECK-NEXT: [[TRUNC1:%[0-9]+]]:_(s8) = G_TRUNC [[UV1]](s16)
+    ; CHECK-NEXT: [[TRUNC2:%[0-9]+]]:_(s8) = G_TRUNC [[UV2]](s16)
+    ; CHECK-NEXT: [[TRUNC3:%[0-9]+]]:_(s8) = G_TRUNC [[UV3]](s16)
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s8>) = G_BUILD_VECTOR [[TRUNC3]](s8), [[TRUNC2]](s8), [[TRUNC1]](s8), [[TRUNC]](s8)
+    ; CHECK-NEXT: $w0 = COPY [[BUILD_VECTOR]](<4 x s8>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:_(<4 x s16>) = COPY $d0
+    %1:_(s16), %2:_(s16), %3:_(s16), %4:_(s16) = G_UNMERGE_VALUES %0(<4 x s16>)
+    %5:_(s8) = G_TRUNC %1(s16)
+    %6:_(s8) = G_TRUNC %2(s16)
+    %7:_(s8) = G_TRUNC %3(s16)
+    %8:_(s8) = G_TRUNC %4(s16)
+    ; Reversed operand order
+    %9:_(<4 x s8>) = G_BUILD_VECTOR %8(s8), %7(s8), %6(s8), %5(s8)
+    $w0 = COPY %9(<4 x s8>)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            test_use_vector_truncate_normal_order
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0
+
+    ; CHECK-LABEL: name: test_use_vector_truncate_normal_order
+    ; CHECK: liveins: $d0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<4 x s16>) = COPY $d0
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(<4 x s8>) = G_TRUNC [[COPY]](<4 x s16>)
+    ; CHECK-NEXT: $w0 = COPY [[TRUNC]](<4 x s8>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:_(<4 x s16>) = COPY $d0
+    %1:_(s16), %2:_(s16), %3:_(s16), %4:_(s16) = G_UNMERGE_VALUES %0(<4 x s16>)
+    %5:_(s8) = G_TRUNC %1(s16)
+    %6:_(s8) = G_TRUNC %2(s16)
+    %7:_(s8) = G_TRUNC %3(s16)
+    %8:_(s8) = G_TRUNC %4(s16)
+    ; Normal operand order
+    %9:_(<4 x s8>) = G_BUILD_VECTOR %5(s8), %6(s8), %7(s8), %8(s8)
+    $w0 = COPY %9(<4 x s8>)
+    RET_ReallyLR implicit $w0
+...

``````````

</details>


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