[llvm] [GISel] Fix ShuffleVector assert (PR #139769)

[Alan Li via llvm-commits]

https://github.com/lialan updated https://github.com/llvm/llvm-project/pull/139769

>From 1704d82af42a9ae9e5d98524f6548767ba58b69b Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Tue, 13 May 2025 13:10:46 -0400
Subject: [PATCH] [GISel] Fix ShuffleVector assert

Fixes issue: https://github.com/llvm/llvm-project/issues/139752

When G_SHUFFLE_VECTOR has only 1 element then it is possible the vector
is decayed into a scalar.
---
 .../lib/CodeGen/GlobalISel/CombinerHelper.cpp |  8 +++++--
 .../prelegalizer-combiner-shuffle.mir         | 24 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
index 5191360c7718a..3abed4d062bfa 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp
@@ -420,8 +420,12 @@ void CombinerHelper::applyCombineShuffleToBuildVector(MachineInstr &MI) const {
     else
       Extracts.push_back(Unmerge2.getReg(Val - Width));
   }
-
-  Builder.buildBuildVector(MI.getOperand(0).getReg(), Extracts);
+  assert(Extracts.size() > 0 && "Expected at least one element in the shuffle");
+  if (Extracts.size() == 1) {
+    Builder.buildCopy(MI.getOperand(0).getReg(), Extracts[0]);
+  } else {
+    Builder.buildBuildVector(MI.getOperand(0).getReg(), Extracts);
+  }
   MI.eraseFromParent();
 }
 
diff --git a/llvm/test/CodeGen/AMDGPU/GlobalISel/prelegalizer-combiner-shuffle.mir b/llvm/test/CodeGen/AMDGPU/GlobalISel/prelegalizer-combiner-shuffle.mir
index bba608cceee19..e500cfe085110 100644
--- a/llvm/test/CodeGen/AMDGPU/GlobalISel/prelegalizer-combiner-shuffle.mir
+++ b/llvm/test/CodeGen/AMDGPU/GlobalISel/prelegalizer-combiner-shuffle.mir
@@ -135,3 +135,27 @@ body:             |
     SI_RETURN
 ...
 
+
+---
+name: shuffle_vector_to_copy
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0, $vgpr1
+    ; CHECK-LABEL: name: shuffle_vector_to_copy
+    ; CHECK: liveins: $vgpr0, $vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(p3) = COPY $vgpr0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(p3) = COPY $vgpr1
+    ; CHECK-NEXT: [[LOAD:%[0-9]+]]:_(<8 x s16>) = G_LOAD [[COPY]](p3) :: (load (<8 x s16>), align 8, addrspace 3)
+    ; CHECK-NEXT: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16), [[UV2:%[0-9]+]]:_(s16), [[UV3:%[0-9]+]]:_(s16), [[UV4:%[0-9]+]]:_(s16), [[UV5:%[0-9]+]]:_(s16), [[UV6:%[0-9]+]]:_(s16), [[UV7:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[LOAD]](<8 x s16>)
+    ; CHECK-NEXT: G_STORE [[UV4]](s16), [[COPY1]](p3) :: (store (s16), addrspace 3)
+    ; CHECK-NEXT: SI_RETURN
+    %0:_(p3) = COPY $vgpr0
+    %1:_(p3) = COPY $vgpr1
+    %12:_(<8 x s16>) = G_IMPLICIT_DEF
+    %10:_(<8 x s16>) = G_LOAD %0(p3) :: (load (<8 x s16>), align 8, addrspace 3)
+    %11:_(s16) = G_SHUFFLE_VECTOR %10(<8 x s16>), %12, shufflemask(4)
+    G_STORE %11(s16), %1(p3) :: (store (s16), addrspace 3)
+    SI_RETURN
+...