[ConstantFold] Fold bitcasts on scalable vector splats.

[sdesmalen]

Summary

This patch also contains changes to InstCombine to avoid a regression for
scalable-const-fp-splat.ll which otherwise would no longer be folded
because the expression is not an explicit ConstExpr any more, but a folded
splat instead.

This patch additionally fixes

https://github.com/llvm/llvm-project/issues/55348.

where one of the operands to be sunk was a ConstExpr. Because the m_pattern()
functions also match ConstExpr's, the cast<Instruction> following the match
led to an assertion failure. Because we now fold the ConstExpr, there is no
longer a way to represent a sext/zext as a ConstExpr and the patterns no
longer match.

[frasercrmck]

From what I gather about the bug report, RISCV doesn't suffer from it as we always dyn_cast to Instruction. Is that right? Just wondering if we should add a similar test - I don't think we have constant expr coverage in sink-splat-operands.ll.

llvm/lib/IR/ConstantFold.cpp

406

DestVecTy == DestVTy, no?

409

cast<VectorType>(DestTy) == DestVTy.

llvm/test/Transforms/InstCombine/addrspacecast.ll

206 ↗

(On Diff #428865)

This one strikes me as peculiar. Is it all less canonical than it was before?

[sdesmalen]

• Addressed comments.
• Reverted the fixed-width change (leading to different addrspacecast constant)

[sdesmalen]

In D125449#3508487, @frasercrmck wrote:

From what I gather about the bug report, RISCV doesn't suffer from it as we always dyn_cast to Instruction. Is that right? Just wondering if we should add a similar test - I don't think we have constant expr coverage in sink-splat-operands.ll.

That's what I did initially, but I couldn't write a test with a ZExt/SExt as ConstExpr that wouldn't fold, so that would mean I'd rewrite the code without being able to test it.

llvm/lib/IR/ConstantFold.cpp

406

You're right, I moved this code around a few times and didn't spot this.

llvm/test/Transforms/InstCombine/addrspacecast.ll

206 ↗

(On Diff #428865)

You're right, that's probably not desired, because addrspacecast's don't fold any further. I've reverted that change now.

[craig.topper]

llvm/lib/IR/ConstantFold.cpp

402

Is it always a bitcast or any cast?

llvm/lib/IR/ConstantFold.cpp

if (CE->isCast()) {
           break;
         }
       if (isAllNull)
         // This is casting one pointer type to another, always BitCast
         return ConstantExpr::getPointerCast(CE->getOperand(0), DestTy);
     }
   }
 
+  // Handle scalable vector splats separately, because we know the
+  // bitcast can always be pushed to the splatted source element.
+  if (auto *SrcSVTy = dyn_cast<ScalableVectorType>(V->getType()))
+    if (auto *DestSVTy = dyn_cast<ScalableVectorType>(DestTy))
+      if (DestSVTy->getElementCount() == SrcSVTy->getElementCount())
+        if (Constant *Splat = V->getSplatValue())
+          return ConstantVector::getSplat(
+              DestSVTy->getElementCount(),
+              ConstantExpr::getCast(opc, Splat, DestSVTy->getElementType()));
+
   // 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()) {
     VectorType *DestVecTy = cast<VectorType>(DestTy);

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

static Type *shrinkFPConstant(ConstantFP *CFP) {
   return nullptr;
 }
 
 // Determine if this is a vector of ConstantFPs and if so, return the minimal
 // type we can safely truncate all elements to.
 // TODO: Make these support undef elements.
 static Type *shrinkFPConstantVector(Value *V) {
   auto *CV = dyn_cast<Constant>(V);
-  auto *CVVTy = dyn_cast<FixedVectorType>(V->getType());
-  if (!CV || !CVVTy)
+  auto *CVT = dyn_cast<VectorType>(V->getType());
+  if (!CV || !CVT)
     return nullptr;
 
-  Type *MinType = nullptr;
+  // Handle the splat case separately, because that works for scalable
+  // vectors, whereas the rest of this function does not.
+  if (Constant *SplatVal = CV->getSplatValue())
+    if (auto *CFP = dyn_cast<ConstantFP>(SplatVal))
+      if (Type *T = shrinkFPConstant(CFP))
+        return VectorType::get(T, CVT->getElementCount());
+
+  if (!isa<FixedVectorType>(CVT))
+    return nullptr;
 
-  unsigned NumElts = CVVTy->getNumElements();
+  Type *MinType = nullptr;
+  unsigned NumElts = cast<FixedVectorType>(CVT)->getNumElements();
 
   // For fixed-width vectors we find the minimal type by looking
   // through the constant values of the vector.
   for (unsigned i = 0; i != NumElts; ++i) {
     auto *CFP = dyn_cast_or_null<ConstantFP>(CV->getAggregateElement(i));
     if (!CFP)
       return nullptr;
 
static Type *getMinimumFPType(Value *V) {
 
   // If this value is a constant, return the constant in the smallest FP type
   // that can accurately represent it.  This allows us to turn
   // (float)((double)X+2.0) into x+2.0f.
   if (auto *CFP = dyn_cast<ConstantFP>(V))
     if (Type *T = shrinkFPConstant(CFP))
       return T;
 
-  // We can only correctly find a minimum type for a scalable vector when it is
-  // a splat. For splats of constant values the fpext is wrapped up as a
-  // ConstantExpr.
-  if (auto *FPCExt = dyn_cast<ConstantExpr>(V))
-    if (FPCExt->getOpcode() == Instruction::FPExt)
-      return FPCExt->getOperand(0)->getType();
-
-  // Try to shrink a vector of FP constants. This returns nullptr on scalable
-  // vectors
   if (Type *T = shrinkFPConstantVector(V))
     return T;
 
   return V->getType();
 }
 
 /// Return true if the cast from integer to FP can be proven to be exact for all
 /// possible inputs (the conversion does not lose any precision).

llvm/test/Transforms/CodeGenPrepare/AArch64/sink-free-instructions.ll

if.then:
   %vmull0 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s1, <8 x i8> %s2)
   ret <8 x i16> %vmull0
 
 if.else:
   %s4 = shufflevector <16 x i8> %b, <16 x i8> poison, <8 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
   %vmull1 = tail call <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8> %s3, <8 x i8> %s4)
   ret <8 x i16> %vmull1
 }
+
+; Ensure sinking add/sub operands doesn't crash when one of them is a Const or ConstExpr.
+define <vscale x 16 x i16> @zext_scalable_constexpr(<vscale x 16 x i8> %a, i1 %p) {
+; CHECK-LABEL: @zext_scalable_constexpr(
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <vscale x 16 x i8> [[A:%.*]] to <vscale x 16 x i16>
+; CHECK-NEXT:    br i1 [[P:%.*]], label [[ZEXT_THEN:%.*]], label [[ZEXT_EXIT:%.*]]
+; CHECK:       zext.then:
+; CHECK-NEXT:    [[TMP2:%.*]] = add <vscale x 16 x i16> [[TMP1]], shufflevector (<vscale x 16 x i16> insertelement (<vscale x 16 x i16> poison, i16 2, i32 0), <vscale x 16 x i16> poison, <vscale x 16 x i32> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 16 x i16> [[TMP2]]
+; CHECK:       zext.exit:
+; CHECK-NEXT:    ret <vscale x 16 x i16> [[TMP1]]
+;
+  %1 = zext <vscale x 16 x i8> %a to <vscale x 16 x i16>
+  br i1 %p, label %zext.then, label %zext.exit
+zext.then:
+  %2 = add <vscale x 16 x i16> %1, zext (<vscale x 16 x i8> shufflevector (<vscale x 16 x i8> insertelement (<vscale x 16 x i8> undef, i8 2, i32 0), <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer) to <vscale x 16 x i16>)
+  ret <vscale x 16 x i16> %2
+
+zext.exit:
+  ret <vscale x 16 x i16> %1
+}

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

 ; RUN: opt -passes=instcombine -S -o - < %s | FileCheck %s
 
 define <vscale x 2 x float> @shrink_splat_scalable_extend(<vscale x 2 x float> %a) {
   ; CHECK-LABEL: @shrink_splat_scalable_extend
-  ; CHECK-NEXT:  %[[FADD:.*]] = fadd <vscale x 2 x float> %a, shufflevector (<vscale x 2 x float> insertelement (<vscale x 2 x float> undef, float -1.000000e+00, i32 0), <vscale x 2 x float> undef, <vscale x 2 x i32> zeroinitializer)
+  ; CHECK-NEXT:  %[[FADD:.*]] = fadd <vscale x 2 x float> %a, shufflevector (<vscale x 2 x float> insertelement (<vscale x 2 x float> poison, float -1.000000e+00, i32 0), <vscale x 2 x float> poison, <vscale x 2 x i32> zeroinitializer)
   ; CHECK-NEXT:  ret <vscale x 2 x float> %[[FADD]]
   %1 = shufflevector <vscale x 2 x float> insertelement (<vscale x 2 x float> undef, float -1.000000e+00, i32 0), <vscale x 2 x float> undef, <vscale x 2 x i32> zeroinitializer
   %2 = fpext <vscale x 2 x float> %a to <vscale x 2 x double>
   %3 = fpext <vscale x 2 x float> %1 to <vscale x 2 x double>
   %4 = fadd <vscale x 2 x double> %2, %3
   %5 = fptrunc <vscale x 2 x double> %4 to <vscale x 2 x float>
   ret <vscale x 2 x float> %5
 }

llvm/test/Transforms/InstSimplify/ConstProp/vscale-inseltpoison.ll

 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Conversion Operations
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 define <vscale x 4 x float> @bitcast() {
 ; CHECK-LABEL: @bitcast(
-; CHECK-NEXT:    ret <vscale x 4 x float> bitcast (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i32 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer) to <vscale x 4 x float>)
+; CHECK-NEXT:    ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0x36A0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ;
   %i1 = insertelement <vscale x 4 x i32> poison, i32 1, i32 0
   %i2 = shufflevector <vscale x 4 x i32> %i1, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %i3 = bitcast <vscale x 4 x i32> %i2 to <vscale x 4 x float>
   ret <vscale x 4 x float> %i3
 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

llvm/test/Transforms/InstSimplify/ConstProp/vscale.ll

 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Conversion Operations
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 define <vscale x 4 x float> @bitcast() {
 ; CHECK-LABEL: @bitcast(
-; CHECK-NEXT:    ret <vscale x 4 x float> bitcast (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer) to <vscale x 4 x float>)
+; CHECK-NEXT:    ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0x36A0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ;
   %i1 = insertelement <vscale x 4 x i32> undef, i32 1, i32 0
   %i2 = shufflevector <vscale x 4 x i32> %i1, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
   %i3 = bitcast <vscale x 4 x i32> %i2 to <vscale x 4 x float>
   ret <vscale x 4 x float> %i3
 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

llvm/test/Transforms/InstSimplify/vscale-inseltpoison.ll

 ; CHECK-NEXT:    ret <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer)
 ;
   %cmp = icmp sle <vscale x 2 x i64> %x, shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 9223372036854775807, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
   ret <vscale x 2 x i1> %cmp
 }
 
 define <vscale x 4 x float> @bitcast() {
 ; CHECK-LABEL: @bitcast(
-; CHECK-NEXT:    ret <vscale x 4 x float> bitcast (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> poison, i32 1, i32 0), <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer) to <vscale x 4 x float>)
+; CHECK-NEXT:    ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0x36A0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ;
   %i1 = insertelement <vscale x 4 x i32> poison, i32 1, i32 0
   %i2 = shufflevector <vscale x 4 x i32> %i1, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %i3 = bitcast <vscale x 4 x i32> %i2 to <vscale x 4 x float>
   ret <vscale x 4 x float> %i3
 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

llvm/test/Transforms/InstSimplify/vscale.ll

 ; CHECK-NEXT:    ret <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer)
 ;
   %cmp = icmp sle <vscale x 2 x i64> %x, shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 9223372036854775807, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
   ret <vscale x 2 x i1> %cmp
 }
 
 define <vscale x 4 x float> @bitcast() {
 ; CHECK-LABEL: @bitcast(
-; CHECK-NEXT:    ret <vscale x 4 x float> bitcast (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer) to <vscale x 4 x float>)
+; CHECK-NEXT:    ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float 0x36A0000000000000, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
 ;
   %i1 = insertelement <vscale x 4 x i32> undef, i32 1, i32 0
   %i2 = shufflevector <vscale x 4 x i32> %i1, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
   %i3 = bitcast <vscale x 4 x i32> %i2 to <vscale x 4 x float>
   ret <vscale x 4 x float> %i3
 }
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;