Skip to content

Commit

Permalink
[VectorCombine] foldShuffleOfCastops - ensure we can scale shuffle ma…
Browse files Browse the repository at this point in the history
…sks 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
  • Loading branch information
RKSimon committed Apr 12, 2024
1 parent 7e7468c commit ea3d0db
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 0 deletions.
6 changes: 6 additions & 0 deletions llvm/lib/Transforms/Vectorize/VectorCombine.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 15 additions & 0 deletions llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down

0 comments on commit ea3d0db

Please sign in to comment.