[ConstantFold] Support scalable constant splats in ConstantFoldCastInstruction

Stacked on #132960 to prevent a regression

Previously only fixed vector splats were handled. This adds supports for scalable vectors too by allowing ConstantExpr splats.

We need to add the extra V->getType()->isVectorTy() check because a ConstantExpr might be a scalar to vector bitcast.

I believe this will also allow casts of fixed vector ConstantExprs to be folded but I couldn't come up with a test case for this, the ConstantExprs seem to be folded away before reaching InstCombine.

Fixes #132922

Author: lukel97

llvm/lib/IR/ConstantFold.cpp

@@ -160,10 +160,10 @@ Constant *llvm::ConstantFoldCastInstruction(unsigned opc, Constant *V,
   // If the cast operand is a constant vector, perform the cast by
   // operating on each element. In the cast of bitcasts, the element
   // count may be mismatched; don't attempt to handle that here.
-  if ((isa<ConstantVector>(V) || isa<ConstantDataVector>(V)) &&
-      DestTy->isVectorTy() &&
-      cast<FixedVectorType>(DestTy)->getNumElements() ==
-          cast<FixedVectorType>(V->getType())->getNumElements()) {
+  if ((isa<ConstantVector, ConstantDataVector, ConstantExpr>(V)) &&
+      DestTy->isVectorTy() && V->getType()->isVectorTy() &&
+      cast<VectorType>(DestTy)->getElementCount() ==
+          cast<VectorType>(V->getType())->getElementCount()) {
     VectorType *DestVecTy = cast<VectorType>(DestTy);
     Type *DstEltTy = DestVecTy->getElementType();
     // Fast path for splatted constants.
@@ -174,6 +174,8 @@ Constant *llvm::ConstantFoldCastInstruction(unsigned opc, Constant *V,
       return ConstantVector::getSplat(
           cast<VectorType>(DestTy)->getElementCount(), Res);
     }
+    if (isa<ScalableVectorType>(DestTy))
+      return nullptr;
     SmallVector<Constant *, 16> res;
     Type *Ty = IntegerType::get(V->getContext(), 32);
     for (unsigned i = 0,

llvm/test/Transforms/InstCombine/scalable-const-fp-splat.ll

@@ -17,8 +17,7 @@ define <vscale x 2 x float> @shrink_splat_scalable_extend(<vscale x 2 x float> %
 define <vscale x 2 x float> @shrink_splat_scalable_extend_rhs_constexpr(<vscale x 2 x float> %a) {
 ; CHECK-LABEL: define <vscale x 2 x float> @shrink_splat_scalable_extend_rhs_constexpr(
 ; CHECK-SAME: <vscale x 2 x float> [[A:%.*]]) {
-; CHECK-NEXT:    [[TMP1:%.*]] = fptrunc <vscale x 2 x double> splat (double -1.000000e+00) to <vscale x 2 x float>
-; CHECK-NEXT:    [[TMP3:%.*]] = fadd <vscale x 2 x float> [[A]], [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = fadd <vscale x 2 x float> [[A]], splat (float -1.000000e+00)
 ; CHECK-NEXT:    ret <vscale x 2 x float> [[TMP3]]
 ;
   %2 = fpext <vscale x 2 x float> %a to <vscale x 2 x double>

llvm/test/Transforms/InstCombine/scalable-trunc.ll

@@ -22,15 +22,15 @@ entry:
 
 define <vscale x 1 x i8> @constant_splat_trunc() {
 ; CHECK-LABEL: @constant_splat_trunc(
-; CHECK-NEXT:    ret <vscale x 1 x i8> trunc (<vscale x 1 x i64> splat (i64 1) to <vscale x 1 x i8>)
+; CHECK-NEXT:    ret <vscale x 1 x i8> splat (i8 1)
 ;
   %1 = trunc <vscale x 1 x i64> splat (i64 1) to <vscale x 1 x i8>
   ret <vscale x 1 x i8> %1
 }
 
 define <vscale x 1 x i8> @constant_splat_trunc_constantexpr() {
 ; CHECK-LABEL: @constant_splat_trunc_constantexpr(
-; CHECK-NEXT:    ret <vscale x 1 x i8> trunc (<vscale x 1 x i64> splat (i64 1) to <vscale x 1 x i8>)
+; CHECK-NEXT:    ret <vscale x 1 x i8> splat (i8 1)
 ;
   ret <vscale x 1 x i8> trunc (<vscale x 1 x i64> splat (i64 1) to <vscale x 1 x i8>)
 }