2 files changed, 19 insertions, 30 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index dcd86db036b..b90e4d846bc 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1715,40 +1715,29 @@ static Value *optimizeIntegerToVectorInsertions(BitCastInst &CI,
   return Result;
 }
 
-/// Given a bitcasted source operand fed into an extract element instruction and
-/// then bitcasted again to a scalar type, eliminate at least one bitcast by
-/// changing the vector type of the extractelement instruction.
-/// Example:
-///   bitcast (extractelement (bitcast <2 x float> %X to <2 x i32>), 1) to float
-///    --->
-///   extractelement <2 x float> %X, i32 1
-static Instruction *foldBitCastExtElt(BitCastInst &BitCast, InstCombiner &IC,
-                                      const DataLayout &DL) {
-  Type *DestType = BitCast.getType();
-  if (DestType->isVectorTy())
-    return nullptr;
-
+/// Canonicalize scalar bitcasts of extracted elements into a bitcast of the
+/// vector followed by extract element. The backend tends to handle bitcasts of
+/// vectors better than bitcasts of scalars because vector registers are
+/// usually not type-specific like scalar integer or scalar floating-point.
+static Instruction *canonicalizeBitCastExtElt(BitCastInst &BitCast,
+                                              InstCombiner &IC,
+                                              const DataLayout &DL) {
   // TODO: Create and use a pattern matcher for ExtractElementInst.
   auto *ExtElt = dyn_cast<ExtractElementInst>(BitCast.getOperand(0));
   if (!ExtElt || !ExtElt->hasOneUse())
     return nullptr;
 
-  Value *InnerBitCast = nullptr;
-  if (!match(ExtElt->getOperand(0), m_BitCast(m_Value(InnerBitCast))))
+  // The bitcast must be to a vectorizable type, otherwise we can't make a new
+  // type to extract from.
+  Type *DestType = BitCast.getType();
+  if (!VectorType::isValidElementType(DestType))
     return nullptr;
 
-  // If the source is not a vector or its element type doesn't match the result
-  // type, bitcast it to a vector type that we can extract from.
-  Type *SourceType = InnerBitCast->getType();
-  if (SourceType->getScalarType() != DestType) {
-    unsigned VecWidth = SourceType->getPrimitiveSizeInBits();
-    unsigned DestWidth = DestType->getPrimitiveSizeInBits();
-    unsigned NumElts = VecWidth / DestWidth;
-    SourceType = VectorType::get(DestType, NumElts);
-    InnerBitCast = IC.Builder->CreateBitCast(InnerBitCast, SourceType, "bc");
-  }
-
-  return ExtractElementInst::Create(InnerBitCast, ExtElt->getOperand(1));
+  unsigned NumElts = ExtElt->getVectorOperandType()->getNumElements();
+  auto *NewVecType = VectorType::get(DestType, NumElts);
+  auto *NewBC = IC.Builder->CreateBitCast(ExtElt->getVectorOperand(),
+                                          NewVecType, "bc");
+  return ExtractElementInst::Create(NewBC, ExtElt->getIndexOperand());
 }
 
 static Instruction *foldVecTruncToExtElt(Value *VecInput, Type *DestTy,
@@ -1922,7 +1911,7 @@ Instruction *InstCombiner::visitBitCast(BitCastInst &CI) {
     }
   }
 
-  if (Instruction *I = foldBitCastExtElt(CI, *this, DL))
+  if (Instruction *I = canonicalizeBitCastExtElt(CI, *this, DL))
     return I;
 
   if (SrcTy->isPointerTy())
diff --git a/llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll b/llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll
index e603036bfe8..97145221099 100644
--- a/llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll
+++ b/llvm/test/Transforms/InstCombine/bitcast-vec-canon.ll
@@ -5,8 +5,8 @@ define double @a(<1 x i64> %y) {
   ret double %c
  
 ; CHECK-LABEL: @a(
-; CHECK-NEXT:  extractelement <1 x i64> %y, i32 0
-; CHECK-NEXT:  bitcast i64 {{.*}} to double
+; CHECK-NEXT:  bitcast <1 x i64> %y to <1 x double>
+; CHECK-NEXT:  extractelement <1 x double> {{.*}}, i32 0
 ; CHECK-NEXT:  ret double
 }