[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

Author: RKSimon

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

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) {