[InstCombine] bitcast (extractelement <1 x elt>, dest) -> bitcast(<1 x elt>, dest)

[Chenbing.Zheng]

Summary

bitcast (extractelement <1 x elt>, dest) -> bitcast(<1 x elt>, dest)

Only solve dest type is vector to avoid inverse transform in visitBitCast.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2375

One-use should not be here - this transform replaces a value with an existing value, so it is always profitable.

llvm/test/Transforms/InstCombine/bitcast-inseltpoison.ll

361

This extractelement of a 1-element vector doesn't seem like the best canonicalization - it could just be a bitcast from vector to scalar. But we have this unusual inverse transform:
https://github.com/llvm/llvm-project/blob/cefe472c51fbcd1aed4d4a090709f25a12a8bc2c/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp#L2750

That potentially creates 2 instructions (extract + bitcast) from a single bitcast. I'm not sure why we do that, but there is some potential regression with the x86 MMX type (see a bit above that in the source code for a similar transform). @craig.topper or @RKSimon might have some ideas on that.

Ideally, we would transform any <1 x Elt> vector extract instruction in the other direction (extelt --> bitcast). Then we would not need this patch because all of the bitcasts would collapse with existing folds.

If that is not possible, then it would at least be better to have this patch moved over to InstSimplify - we are not creating any new instructions.

[Chenbing.Zheng]

address comments and add more tests

[spatel]

Please pre-commit the baseline tests.

But we also need tests with scalable vectors - the patch will crash on this example:

define <2 x i32> @bitcast_extelt5_scalable(<vscale x 1 x i64> %A) {
  %ext = extractelement <vscale x 1 x i64> %A, i32 0
  %bc = bitcast i64 %ext to <2 x i32>
  ret <2 x i32> %bc
}

[Chenbing.Zheng]

address comment and rebase tests

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2375

Is there an advantage to adding a new API for this?
Since we need the casted object anyway, this seems better:

auto *FixedVType = dyn_cast<FixedVectorType>(VecType);
if (DestType->isVectorTy() && FixedVType && FixedVType->getNumElements() == 1)

[Chenbing.Zheng]

address comments

[Chenbing.Zheng]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

2375

Done and thanks for your advice~

[spatel]

LGTM

llvm/include/llvm/IR/Type.h

public:
   /// Return true if this is a pointer type or a vector of pointer types.
   bool isPtrOrPtrVectorTy() const { return getScalarType()->isPointerTy(); }
 
   /// True if this is an instance of VectorType.
   inline bool isVectorTy() const {
     return getTypeID() == ScalableVectorTyID || getTypeID() == FixedVectorTyID;
   }
 
+  /// True if this is an instance of ScalableVectorType.
+  inline bool isScalableVectorTy() const {
+    return getTypeID() == ScalableVectorTyID;
+  }
+
   /// Return true if this type could be converted with a lossless BitCast to
   /// type 'Ty'. For example, i8* to i32*. BitCasts are valid for types of the
   /// same size only where no re-interpretation of the bits is done.
   /// Determine if this type could be losslessly bitcast to Ty
   bool canLosslesslyBitCastTo(Type *Ty) const;
 
   /// Return true if this type is empty, that is, it has no elements or all of
   /// its elements are empty.

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

static Instruction *canonicalizeBitCastExtElt(BitCastInst &BitCast,
   Value *VecOp, *Index;
   if (!match(BitCast.getOperand(0),
              m_OneUse(m_ExtractElt(m_Value(VecOp), m_Value(Index)))))
     return nullptr;
 
   // 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;
-
-  auto *NewVecType =
-      VectorType::get(DestType, cast<VectorType>(VecOp->getType()));
+  VectorType *VecType = cast<VectorType>(VecOp->getType());
+  if (VectorType::isValidElementType(DestType)) {
+    auto *NewVecType = VectorType::get(DestType, VecType);
     auto *NewBC = IC.Builder.CreateBitCast(VecOp, NewVecType, "bc");
     return ExtractElementInst::Create(NewBC, Index);
   }
 
+  // Only solve DestType is vector to avoid inverse transform in visitBitCast.
+  // bitcast (extractelement <1 x elt>, dest) -> bitcast(<1 x elt>, dest)
+  if (DestType->isVectorTy() && !VecType->isScalableVectorTy() &&
+      cast<FixedVectorType>(VecType)->getNumElements() == 1)
+    return CastInst::Create(Instruction::BitCast, VecOp, DestType);
+
+  return nullptr;
+}
+
 /// Change the type of a bitwise logic operation if we can eliminate a bitcast.
 static Instruction *foldBitCastBitwiseLogic(BitCastInst &BitCast,
                                             InstCombiner::BuilderTy &Builder) {
   Type *DestTy = BitCast.getType();
   BinaryOperator *BO;
 
   if (!match(BitCast.getOperand(0), m_OneUse(m_BinOp(BO))) ||
       !BO->isBitwiseLogicOp())

llvm/test/Transforms/InstCombine/bitcast-inseltpoison.ll

 ; CHECK-NEXT:    ret i64 [[BC2]]
 ;
   %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
 }
 
-; TODO: This should return %A.
-
 define <2 x i32> @bitcast_extelt3(<2 x i32> %A) {
 ; CHECK-LABEL: @bitcast_extelt3(
-; CHECK-NEXT:    [[BC1:%.*]] = bitcast <2 x i32> [[A:%.*]] to <1 x i64>
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <1 x i64> [[BC1]], i64 0
-; CHECK-NEXT:    [[BC2:%.*]] = bitcast i64 [[EXT]] to <2 x i32>
-; CHECK-NEXT:    ret <2 x i32> [[BC2]]
+; CHECK-NEXT:    ret <2 x i32> [[A:%.*]]
 ;
   %bc1 = bitcast <2 x i32> %A to <1 x i64>
   %ext = extractelement <1 x i64> %bc1, i32 0
   %bc2 = bitcast i64 %ext to <2 x i32>
   ret <2 x i32> %bc2
 }
 
 ; Handle the case where the input is not a vector.
 
 define double @bitcast_extelt4(i128 %A) {
 ; CHECK-LABEL: @bitcast_extelt4(

llvm/test/Transforms/InstCombine/bitcast.ll

 ; CHECK-NEXT:    ret i64 [[BC2]]
 ;
   %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
 }
 
-; TODO: This should return %A.
-
 define <2 x i32> @bitcast_extelt3(<2 x i32> %A) {
 ; CHECK-LABEL: @bitcast_extelt3(
-; CHECK-NEXT:    [[BC1:%.*]] = bitcast <2 x i32> [[A:%.*]] to <1 x i64>
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <1 x i64> [[BC1]], i64 0
-; CHECK-NEXT:    [[BC2:%.*]] = bitcast i64 [[EXT]] to <2 x i32>
-; CHECK-NEXT:    ret <2 x i32> [[BC2]]
+; CHECK-NEXT:    ret <2 x i32> [[A:%.*]]
 ;
   %bc1 = bitcast <2 x i32> %A to <1 x i64>
   %ext = extractelement <1 x i64> %bc1, i32 0
   %bc2 = bitcast i64 %ext to <2 x i32>
   ret <2 x i32> %bc2
 }
 
 ; Handle the case where the input is not a vector.
 
 define double @bitcast_extelt4(i128 %A) {
 ; CHECK-LABEL: @bitcast_extelt4(
 ; CHECK-NEXT:    [[BC:%.*]] = bitcast i128 [[A:%.*]] to <2 x double>
 ; CHECK-NEXT:    [[BC2:%.*]] = extractelement <2 x double> [[BC]], i64 0
 ; CHECK-NEXT:    ret double [[BC2]]
 ;
   %bc1 = bitcast i128 %A to <2 x i64>
   %ext = extractelement <2 x i64> %bc1, i32 0
   %bc2 = bitcast i64 %ext to double
   ret double %bc2
 }
 
 define <2 x i32> @bitcast_extelt5(<1 x i64> %A) {
 ; CHECK-LABEL: @bitcast_extelt5(
-; CHECK-NEXT:    [[EXT:%.*]] = extractelement <1 x i64> [[A:%.*]], i64 0
-; CHECK-NEXT:    [[BC:%.*]] = bitcast i64 [[EXT]] to <2 x i32>
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <1 x i64> [[A:%.*]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[BC]]
 ;
   %ext = extractelement <1 x i64> %A, i32 0
   %bc = bitcast i64 %ext to <2 x i32>
   ret <2 x i32> %bc
 }
 
 define <2 x i32> @bitcast_extelt5_scalable(<vscale x 1 x i64> %A) {