[llvm] ea3d0db - [VectorCombine] foldShuffleOfCastops - ensure we can scale shuffle masks between bitcasted vector types

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-04-12T13:53:02+01:00
New Revision: ea3d0db130b9a6c4678c11dd8bc48e5630624b62

URL: https://github.com/llvm/llvm-project/commit/ea3d0db130b9a6c4678c11dd8bc48e5630624b62
DIFF: https://github.com/llvm/llvm-project/commit/ea3d0db130b9a6c4678c11dd8bc48e5630624b62.diff

LOG: [VectorCombine] foldShuffleOfCastops - ensure we can scale shuffle masks between bitcasted vector types

Don't just assert that the src/dst vector element counts are multiples of one another - in general IR this can actually happen.

Reported by @mikaelholmen

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VectorCombine.cpp
    llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 44cba60013afa3..653e0ff7a3ece4 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1483,6 +1483,12 @@ bool VectorCombine::foldShuffleOfCastops(Instruction &I) {
   assert((NumDstElts == NumSrcElts || Opcode == Instruction::BitCast) &&
          "Only bitcasts expected to alter src/dst element counts");
 
+  // Check for bitcasting of unscalable vector types.
+  // e.g. <32 x i40> -> <40 x i32>
+  if (NumDstElts != NumSrcElts && (NumSrcElts % NumDstElts) != 0 &&
+      (NumDstElts % NumSrcElts) != 0)
+    return false;
+
   SmallVector<int, 16> NewMask;
   if (NumSrcElts >= NumDstElts) {
     // The bitcast is from wide to narrow/equal elements. The shuffle mask can

diff  --git a/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll b/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
index 4f8cfea146eade..7d9f7e390b9c04 100644
--- a/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
@@ -239,6 +239,21 @@ define <16 x i16> @revpair_bitcast_v4i32_v16i16(<4 x i32> %a0, <4 x i32> %a1) {
   ret <16 x i16> %r
 }
 
+; negative - bitcasts (unscalable element counts)
+
+define <4 x i32> @shuffle_bitcast_v32i40_v4i32(<32 x i40> %a0, <32 x i40> %a1) {
+; CHECK-LABEL: @shuffle_bitcast_v32i40_v4i32(
+; CHECK-NEXT:    [[X0:%.*]] = bitcast <32 x i40> [[A0:%.*]] to <40 x i32>
+; CHECK-NEXT:    [[X1:%.*]] = bitcast <32 x i40> [[A1:%.*]] to <40 x i32>
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <40 x i32> [[X0]], <40 x i32> [[X1]], <4 x i32> <i32 0, i32 42, i32 poison, i32 poison>
+; CHECK-NEXT:    ret <4 x i32> [[R]]
+;
+  %x0 = bitcast <32 x i40> %a0 to <40 x i32>
+  %x1 = bitcast <32 x i40> %a1 to <40 x i32>
+  %r = shufflevector <40 x i32> %x0, <40 x i32> %x1, <4 x i32> <i32 0, i32 42, i32 poison, i32 poison>
+  ret <4 x i32> %r
+}
+
 ; negative - src type mismatch
 
 define <8 x i32> @concat_sext_v4i8_v4i16_v8i32(<4 x i8> %a0, <4 x i16> %a1) {