[ConstantFold][SVE] Fix constant folding for bitcast.

[huihuiz]

Summary

Do not iterate on scalable vector type in BitCastConstantVector.
Continuation work of D70985, D71147.

Add support to fold bitcast into splat value for splat vector.

[huihuiz]

Current upstream crash with llvm/lib/IR/Value.cpp:404: void llvm::Value::doRAUW(llvm::Value *, llvm::Value::ReplaceMetadataUses): Assertion `New->getType() == getType() && "replaceAllUses of value with new value of different type!"' failed.

run: opt -S -constprop bitcast.ll -o -

bitcast.ll

define <vscale x 4 x float> @bitcast_scalable_constant() {
  %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
}

[spatel]

LGTM. I don't know how many of these patches are needed, but it would be slightly more efficient to create a "vscale.ll" test file and put all of the related tests in that 1 file.

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

52

Do we want to add support for a special case here for splat operations (like the one in the test)?

[huihuiz]

In D71389#1781620, @spatel wrote:

LGTM. I don't know how many of these patches are needed, but it would be slightly more efficient to create a "vscale.ll" test file and put all of the related tests in that 1 file.

Thank you Sanjay!
Yes, I will try my best to combine them. The current code base has a lot of assumption based on fixed length vector.
I am going through the IR passes, and try to fix them to the best I can.

llvm/lib/IR/ConstantFold.cpp

52

Actually we can not, we will end up in scalable vector constant.
Instead of <4 x i32> <i32 1, i32 1, i32 1, i32 1>, we will need a form like "<vscale x 4 x i32> [?]"

I don't know how to express scalable vector constant in current upstream, please let me know if I get anything wrong?

[efriedma]

llvm/lib/IR/ConstantFold.cpp

52

I think the suggestion is to special-case the specific sequence of insertelement+splatvector+bitcast, and bitcast the operand of the insertelement. Not sure how useful that would be in general, though.

[spatel]

llvm/lib/IR/ConstantFold.cpp

52

If we want to do that transform, it would go in instsimplify since we're doing a multi-instruction analysis/fold that becomes a constant.

[spatel]

llvm/lib/IR/ConstantFold.cpp

52

Sorry - might've misread that. If we're creating a new bitcast, that would have to go in instcombine (we don't create new instructions in instsimplify).

[huihuiz]

Yes, this transformation will need to go in InstCombine.
We should probably implement this when constant scalable vector support is ready, so that the effect will be most obvious.

[efriedma]

llvm/lib/IR/ConstantFold.cpp

52

By the time we get here, the shuffle and insertelement are ConstantExprs, not instructions.

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

52

In terms of usefulness, this would make it easier to determine whether the value is a splat when calling Constant::getSplatValue without having to look through the bitcast first.

[huihuiz]

Add special case handling:

For splat vector, fold bitcast to splat value.

[sdesmalen]

Thanks for adding the extra constant fold @huihuiz!

llvm/lib/IR/ConstantFold.cpp

580

Because we're matching for a zero mask, explicitly matching m_Undef here seems unnecessarily restrictive, m_Value() should be sufficient.

587

Can we match this value directly with m_Zero(Zero), rather than needing auto InsertElem and InsertElem->getOperand(2) ?

590

Can we match the value for CE->getOperand(1) value directly with something like m_Value(Vec2). I'd expect CE->getOperand(2) can reuse Zero as suggested above.

[bjope]

llvm/lib/IR/ConstantFold.cpp

572

Don't you need some check that both sides of the bitcast are vectors with matching number of elements?

[huihuiz]

Clean up code based on reviewer feedback.

llvm/lib/IR/ConstantFold.cpp

580

Yes, we don't need to restrict it to m_Undef().

I think we should restrict it with m_Constant(). I would expect operands to be constant at ConstantFoldCastInstruction().
Doing this can also help avoid unnecessary cast<constant> when calling ConstantExpr to get new InsertElement or cast operands into new types.

587

Current pattern matcher doesn't support m_Zero(Zero).
I use m_CombineAnd(m_Zero(), m_Constant(Zero)) instead.

llvm/lib/IR/ConstantFold.cpp

 /// Convert the specified vector Constant node to the specified vector type.
 /// At this point, we know that the elements of the input vector constant are
 /// all simple integer or FP values.
 static Constant *BitCastConstantVector(Constant *CV, VectorType *DstTy) {
 
   if (CV->isAllOnesValue()) return Constant::getAllOnesValue(DstTy);
   if (CV->isNullValue()) return Constant::getNullValue(DstTy);
 
+  // Do not iterate on scalable vector. The num of elements is unknown at
+  // compile-time.
+  if (DstTy->isScalable())
+    return nullptr;
+
   // If this cast changes element count then we can't handle it here:
   // doing so requires endianness information.  This should be handled by
   // Analysis/ConstantFolding.cpp
   unsigned NumElts = DstTy->getNumElements();
   if (NumElts != CV->getType()->getVectorNumElements())
     return nullptr;
 
   Type *DstEltTy = DstTy->getElementType();
if (CE->isCast()) {
       for (unsigned i = 1, e = CE->getNumOperands(); i != e; ++i)
         if (!CE->getOperand(i)->isNullValue()) {
           isAllNull = false;
           break;
         }
       if (isAllNull)
         // This is casting one pointer type to another, always BitCast
         return ConstantExpr::getPointerCast(CE->getOperand(0), DestTy);
+    } else if (CE->getType()->isVectorTy() && DestTy->isVectorTy() &&
+               CE->getType()->getVectorElementCount() ==
+                   DestTy->getVectorElementCount() &&
+               CE->getOpcode() == Instruction::ShuffleVector &&
+               opc == Instruction::BitCast) {
+      // For splat vector, fold bitcast to splat value.
+      // BitCast(ShuffleVector(InsertElement(C1, SplatV, Zero), C2, Zero)) to NewType
+      // into
+      // ShuffleVector(InsertElement(C1, BitCast(SplatV) to NewType, Zero), C2, Zero)
+      Constant *SplatV, *C1, *C2, *ZeroIdx, *ZeroMask;
+      if (match(CE, m_ShuffleVector(
+                        m_InsertElement(
+                            m_Constant(C1), m_Constant(SplatV),
+                            m_CombineAnd(m_Zero(), m_Constant(ZeroIdx))),
+                        m_Constant(C2),
+                        m_CombineAnd(m_Zero(), m_Constant(ZeroMask))))) {
+        auto *CastedSplatV =
+            ConstantExpr::getCast(opc, SplatV, DestTy->getScalarType());
+        return ConstantExpr::getShuffleVector(
+            ConstantExpr::getInsertElement(
+                ConstantExpr::getCast(opc, C1, DestTy), CastedSplatV, ZeroIdx),
+            ConstantExpr::getCast(opc, C2, DestTy), ZeroMask);
+      }
     }
   }
 
   // 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() &&

llvm/test/Analysis/ConstantFolding/bitcast.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -constprop -S -verify | FileCheck %s
+
+define <vscale x 4 x float> @bitcast_scalable_constant() {
+; CHECK-LABEL: @bitcast_scalable_constant(
+; CHECK-NEXT:    ret <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> undef, float 0x36A0000000000000, i32 0), <vscale x 4 x float> undef, <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
+}