[InstCombine] fold bitcasts around an extractelement

[spatel]

Summary

This is part of fixing PR25543:
https://llvm.org/bugs/show_bug.cgi?id=25543

If this looks good, the next step will be to generalize this fold:
trunc ( lshr ( bitcast X) ) -> extractelement (X)

Ie, I'm hoping to replace the existing transform of:
bitcast ( trunc ( lshr ( bitcast X)))
added by:
http://reviews.llvm.org/rL112232

with 2 more general transforms to catch the case in the bug report.

[spatel]

Ping.

[spatel]

Ping * 2.

[hfinkel]

LGTM.

[spatel]

Thanks, Hal! I'll post the 2nd half of this shortly.

llvm/trunk/lib/Transforms/InstCombine/InstCombineCasts.cpp

for (unsigned i = 0, e = Elements.size(); i != e; ++i) {
 
     Result = IC.Builder->CreateInsertElement(Result, Elements[i],
                                              IC.Builder->getInt32(i));
   }
 
   return Result;
 }
 
+/// Given a bitcasted vector fed into an extract element instruction and then
+/// bitcasted again, 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) {
+  // 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))))
+    return nullptr;
+
+  VectorType *VecType = cast<VectorType>(InnerBitCast->getType());
+  Type *DestType = BitCast.getType();
+
+  // If the element type of the vector doesn't match the result type,
+  // bitcast it to a vector type that we can extract from.
+  if (VecType->getElementType() != DestType) {
+    unsigned VecWidth = VecType->getPrimitiveSizeInBits();
+    unsigned DestWidth = DestType->getPrimitiveSizeInBits();
+    unsigned NumElts = VecWidth / DestWidth;
+    VecType = VectorType::get(DestType, NumElts);
+    InnerBitCast = IC.Builder->CreateBitCast(InnerBitCast, VecType, "bc");
+  }
+
+  return ExtractElementInst::Create(InnerBitCast, ExtElt->getOperand(1));
+}
+
 static Instruction *foldVecTruncToExtElt(Value *VecInput, Type *DestTy,
                                          unsigned ShiftAmt, InstCombiner &IC,
                                          const DataLayout &DL) {
   VectorType *VecTy = cast<VectorType>(VecInput->getType());
   unsigned DestWidth = DestTy->getPrimitiveSizeInBits();
   unsigned VecWidth = VecTy->getPrimitiveSizeInBits();
 
   if ((VecWidth % DestWidth != 0) || (ShiftAmt % DestWidth != 0))
if (SVI->hasOneUse() && DestTy->isVectorTy() &&
         Value *RHS = Builder->CreateBitCast(SVI->getOperand(1), DestTy);
         // Return a new shuffle vector.  Use the same element ID's, as we
         // know the vector types match #elts.
         return new ShuffleVectorInst(LHS, RHS, SVI->getOperand(2));
       }
     }
   }
 
+  if (Instruction *I = foldBitCastExtElt(CI, *this, DL))
+    return I;
+
   if (SrcTy->isPointerTy())
     return commonPointerCastTransforms(CI);
   return commonCastTransforms(CI);
 }
 
 Instruction *InstCombiner::visitAddrSpaceCast(AddrSpaceCastInst &CI) {
   // If the destination pointer element type is not the same as the source's
   // first do a bitcast to the destination type, and then the addrspacecast.

llvm/trunk/test/Transforms/InstCombine/bitcast.ll

 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:  %tmp24 = extractelement <2 x float> %A, i32 1
 ; CHECK-NEXT:  bitcast <2 x i64> %B to <4 x float>
 ; CHECK-NEXT:  %tmp4 = extractelement <4 x float> {{.*}}, i32 2
 ; CHECK-NEXT:  %add = fadd float %tmp24, %tmp4
 ; CHECK-NEXT:  ret float %add
 }
 
-; TODO: Both bitcasts are unnecessary; change the extractelement.
+; Both bitcasts are unnecessary; change the extractelement.
 
 define float @bitcast_extelt1(<2 x float> %A) {
   %bc1 = bitcast <2 x float> %A to <2 x i32>
   %ext = extractelement <2 x i32> %bc1, i32 0
   %bc2 = bitcast i32 %ext to float
   ret float %bc2
 
 ; CHECK-LABEL: @bitcast_extelt1(
-; CHECK-NEXT:  %bc1 = bitcast <2 x float> %A to <2 x i32>
-; CHECK-NEXT:  %ext = extractelement <2 x i32> %bc1, i32 0
-; CHECK-NEXT:  %bc2 = bitcast i32 %ext to float
+; CHECK-NEXT:  %bc2 = extractelement <2 x float> %A, i32 0
 ; CHECK-NEXT:  ret float %bc2
 }
 
-; TODO: Second bitcast can be folded into the first.
+; Second bitcast can be folded into the first.
 
 define i64 @bitcast_extelt2(<4 x float> %A) {
   %bc1 = bitcast <4 x float> %A to <2 x double>
   %ext = extractelement <2 x double> %bc1, i32 1
   %bc2 = bitcast double %ext to i64
   ret i64 %bc2
 
 ; CHECK-LABEL: @bitcast_extelt2(
-; CHECK-NEXT:  %bc1 = bitcast <4 x float> %A to <2 x double>
-; CHECK-NEXT:  %ext = extractelement <2 x double> %bc1, i32 1
-; CHECK-NEXT:  %bc2 = bitcast double %ext to i64
+; CHECK-NEXT:  %bc = bitcast <4 x float> %A to <2 x i64>
+; CHECK-NEXT:  %bc2 = extractelement <2 x i64> %bc, i32 1
 ; CHECK-NEXT:  ret i64 %bc2
 }
 
 define <2 x i32> @test4(i32 %A, i32 %B){
   %tmp38 = zext i32 %A to i64
   %tmp32 = zext i32 %B to i64
   %tmp33 = shl i64 %tmp32, 32
   %ins35 = or i64 %tmp33, %tmp38