[llvm] 717d3f3 - [VectorCombine] foldShuffleOfCastops - add initial shuffle(bitcast(x),bitcast(y)) -> bitcast(shuffle(x,y)) support

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-04-11T11:43:11+01:00
New Revision: 717d3f3974f43d90c1b8829a4077bbc2a2413c83

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

LOG: [VectorCombine] foldShuffleOfCastops - add initial shuffle(bitcast(x),bitcast(y)) -> bitcast(shuffle(x,y)) support

Just handle cases where the bitcast src/dst element counts are the same (future patches will add shuffle mask scaling)

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 2f9767538e6cb1..b74fdf27d213a1 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -1459,7 +1459,7 @@ bool VectorCombine::foldShuffleOfCastops(Instruction &I) {
     return false;
 
   Instruction::CastOps Opcode = C0->getOpcode();
-  if (Opcode == Instruction::BitCast || C0->getSrcTy() != C1->getSrcTy())
+  if (C0->getSrcTy() != C1->getSrcTy())
     return false;
 
   // Handle shuffle(zext_nneg(x), sext(y)) -> sext(shuffle(x,y)) folds.
@@ -1473,10 +1473,9 @@ bool VectorCombine::foldShuffleOfCastops(Instruction &I) {
   auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
   auto *CastDstTy = dyn_cast<FixedVectorType>(C0->getDestTy());
   auto *CastSrcTy = dyn_cast<FixedVectorType>(C0->getSrcTy());
-  if (!ShuffleDstTy || !CastDstTy || !CastSrcTy)
+  if (!ShuffleDstTy || !CastDstTy || !CastSrcTy ||
+      CastDstTy->getElementCount() != CastSrcTy->getElementCount())
     return false;
-  assert(CastDstTy->getElementCount() == CastSrcTy->getElementCount() &&
-         "Unexpected src/dst element counts");
 
   auto *NewShuffleDstTy =
       FixedVectorType::get(CastSrcTy->getScalarType(), Mask.size());

diff  --git a/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll b/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
index 60c6ff97c58b52..97fceacd827580 100644
--- a/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
+++ b/llvm/test/Transforms/VectorCombine/X86/shuffle-of-casts.ll
@@ -165,13 +165,12 @@ define <8 x double> @interleave_fpext_v4f32_v8f64(<4 x float> %a0, <4 x float> %
   ret <8 x double> %r
 }
 
-; TODO - bitcasts (same element count)
+; bitcasts (same element count)
 
 define <8 x float> @concat_bitcast_v4i32_v8f32(<4 x i32> %a0, <4 x i32> %a1) {
 ; CHECK-LABEL: @concat_bitcast_v4i32_v8f32(
-; CHECK-NEXT:    [[X0:%.*]] = bitcast <4 x i32> [[A0:%.*]] to <4 x float>
-; CHECK-NEXT:    [[X1:%.*]] = bitcast <4 x i32> [[A1:%.*]] to <4 x float>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x float> [[X0]], <4 x float> [[X1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[A0:%.*]], <4 x i32> [[A1:%.*]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[R:%.*]] = bitcast <8 x i32> [[TMP1]] to <8 x float>
 ; CHECK-NEXT:    ret <8 x float> [[R]]
 ;
   %x0 = bitcast <4 x i32> %a0 to <4 x float>