[InstCombine][SVE] Fix visitInsertElementInst for scalable type.

[huihuiz]

Summary

This patch fixes the following issues in visitInsertElementInst:

1. Bail out for scalable type when analysis requires fixed size number of vector elements.
2. Use cast<FixedVectorType> to get vector number of elements. This ensure assertion on scalable vector type.
3. For scalable type, avoid folding a chain of insertelement into splat: insertelt(insertelt(insertelt(insertelt X, %k, 0), %k, 1), %k, 2) ... -> shufflevector(insertelt(X, %k, 0), undef, zero) The length of scalable vector is unknown at compile-time, therefore we don't know if given insertelement sequence is valid for splat.

[huihuiz]

1. Current upstream crash at:

llvm::ShuffleVectorInst::ShuffleVectorInst(llvm::Value *, llvm::Value *, ArrayRef<int>, const llvm::Twine &, llvm::Instruction *): Assertion `isValidOperands(V1, V2, Mask) && "Invalid shuffle vector instruction operands!"' failed.
take test.ll , run opt -S -instcombine test.ll

define <vscale x 4 x i32> @insertelement_extractelement(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
  %t0 = extractelement <vscale x 4 x i32> %a, i32 1
  %t1 = insertelement <vscale x 4 x i32> %b, i32 %t0, i32 0
  ret <vscale x 4 x i32> %t1
}

2. Current upstream is doing wrong fold, trying to fold a chain of insertelement into splat for scalable type. The vector length of scalable type is unknown at compile-time. Therefore a valid sequence of insertelement for fixed-length vector, may not be valid for scalable type to fold into splat.

take test.ll, run opt -S -instcombine test.ll

define <vscale x 4 x float> @insertelement_sequene_may_not_be_splat(float %x) {
  %t0 = insertelement <vscale x 4 x float> undef, float %x, i32 0
  %t1 = insertelement <vscale x 4 x float> %t0, float %x, i32 1
  %t2 = insertelement <vscale x 4 x float> %t1, float %x, i32 2
  %t3 = insertelement <vscale x 4 x float> %t2, float %x, i32 3
  ret <vscale x 4 x float> %t3
}

currently generating:

define <vscale x 4 x float> @insertelement_sequene_may_not_be_splat(float %x) {
  %t0 = insertelement <vscale x 4 x float> undef, float %x, i32 0
  %t3 = shufflevector <vscale x 4 x float> %t0, <vscale x 4 x float> undef, <vscale x 4 x i32> zeroinitializer
  ret <vscale x 4 x float> %t3
}

[sdesmalen]

Thanks for working on this @huihuiz! I've left some comments on the patch.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1069–1070

nit: !isa<ScalableVectorType> -> isa<FixedVectorType>

1074

This is missing a && isa<FixedVectorType>(ExtVecOp->getType()), because the vector that the value is extracted from could be scalable.
(which should be guarded by an extra test case)

1119

This probably also needs a change in SimplifyDemandedVectorElts to bail out early on scalable vectors.

llvm/test/Transforms/InstCombine/vscale_insertelement.ll

5

nit: Can you add a comment what this is testing <=> what it is trying to prevent?

18

nit: This comment doesn't explain what causes the invalid isValidOperands to occur. Can you explain what code-path/optimization you're trying to prevent in this test?

31

Same here.

[huihuiz]

Thanks Sander for the feedback!

Addressed review comments.

[huihuiz]

Rebase.

[efriedma]

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

852

I'd prefer to bail out more aggressively here; I'm not sure this makes sense for scalable vectors.

890

I'd prefer to bail out more aggressively here; I'm not sure this makes sense for scalable vectors.

[huihuiz]

Address review comments.

[efriedma]

LGTM with one minor nit

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1193

Cannot

[sdesmalen]

LGTM as well. Thanks for adding the extra comments to the tests!

[huihuiz]

Thanks! Fixed in the committed patch.

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

 ///
 /// If the information about demanded elements can be used to simplify the
 /// operation, the operation is simplified, then the resultant value is
 /// returned.  This returns null if no change was made.
 Value *InstCombiner::SimplifyDemandedVectorElts(Value *V, APInt DemandedElts,
                                                 APInt &UndefElts,
                                                 unsigned Depth,
                                                 bool AllowMultipleUsers) {
-  unsigned VWidth = cast<VectorType>(V->getType())->getNumElements();
+  // Cannot analyze scalable type. The number of vector elements is not a
+  // compile-time constant.
+  if (isa<ScalableVectorType>(V->getType()))
+    return nullptr;
+
+  unsigned VWidth = cast<FixedVectorType>(V->getType())->getNumElements();
   APInt EltMask(APInt::getAllOnesValue(VWidth));
   assert((DemandedElts & ~EltMask) == 0 && "Invalid DemandedElts!");
 
   if (isa<UndefValue>(V)) {
     // If the entire vector is undefined, just return this info.
     UndefElts = EltMask;
     return nullptr;
   }

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

 /// Note: we intentionally don't try to fold earlier shuffles since they have
 /// often been chosen carefully to be efficiently implementable on the target.
 using ShuffleOps = std::pair<Value *, Value *>;
 
 static ShuffleOps collectShuffleElements(Value *V, SmallVectorImpl<int> &Mask,
                                          Value *PermittedRHS,
                                          InstCombiner &IC) {
   assert(V->getType()->isVectorTy() && "Invalid shuffle!");
-  unsigned NumElts = cast<VectorType>(V->getType())->getNumElements();
+  unsigned NumElts = cast<FixedVectorType>(V->getType())->getNumElements();
 
   if (isa<UndefValue>(V)) {
     Mask.assign(NumElts, -1);
     return std::make_pair(
         PermittedRHS ? UndefValue::get(PermittedRHS->getType()) : V, nullptr);
   }
 
   if (isa<ConstantAggregateZero>(V)) {
Instruction *InstCombiner::visitInsertValueInst(InsertValueInst &I) {
   }
 
   if (IsRedundant)
     return replaceInstUsesWith(I, I.getOperand(0));
   return nullptr;
 }
 
 static bool isShuffleEquivalentToSelect(ShuffleVectorInst &Shuf) {
+  // Can not analyze scalable type, the number of elements is not a compile-time
+  // constant.
+  if (isa<ScalableVectorType>(Shuf.getOperand(0)->getType()))
+    return false;
+
   int MaskSize = Shuf.getShuffleMask().size();
   int VecSize =
-      cast<VectorType>(Shuf.getOperand(0)->getType())->getNumElements();
+      cast<FixedVectorType>(Shuf.getOperand(0)->getType())->getNumElements();
 
   // A vector select does not change the size of the operands.
   if (MaskSize != VecSize)
     return false;
 
   // Each mask element must be undefined or choose a vector element from one of
   // the source operands without crossing vector lanes.
   for (int i = 0; i != MaskSize; ++i) {

 /// insertelt(insertelt(insertelt(insertelt X, %k, 0), %k, 1), %k, 2) ... ->
 /// shufflevector(insertelt(X, %k, 0), undef, zero)
 static Instruction *foldInsSequenceIntoSplat(InsertElementInst &InsElt) {
   // We are interested in the last insert in a chain. So if this insert has a
   // single user and that user is an insert, bail.
   if (InsElt.hasOneUse() && isa<InsertElementInst>(InsElt.user_back()))
     return nullptr;
 
-  auto *VecTy = cast<VectorType>(InsElt.getType());
-  unsigned NumElements = VecTy->getNumElements();
+  VectorType *VecTy = InsElt.getType();
+  // Can not handle scalable type, the number of elements is not a compile-time
+  // constant.
+  if (isa<ScalableVectorType>(VecTy))
+    return nullptr;
+  unsigned NumElements = cast<FixedVectorType>(VecTy)->getNumElements();
 
   // Do not try to do this for a one-element vector, since that's a nop,
   // and will cause an inf-loop.
   if (NumElements == 1)
     return nullptr;
 
   Value *SplatVal = InsElt.getOperand(1);
   InsertElementInst *CurrIE = &InsElt;

 /// Try to fold an insert element into an existing splat shuffle by changing
 /// the shuffle's mask to include the index of this insert element.
 static Instruction *foldInsEltIntoSplat(InsertElementInst &InsElt) {
   // Check if the vector operand of this insert is a canonical splat shuffle.
   auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0));
   if (!Shuf || !Shuf->isZeroEltSplat())
     return nullptr;
 
+  // Bail out early if shuffle is scalable type. The number of elements in
+  // shuffle mask is unknown at compile-time.
+  if (isa<ScalableVectorType>(Shuf->getType()))
+    return nullptr;
+
   // Check for a constant insertion index.
   uint64_t IdxC;
   if (!match(InsElt.getOperand(2), m_ConstantInt(IdxC)))
     return nullptr;
 
   // Check if the splat shuffle's input is the same as this insert's scalar op.
   Value *X = InsElt.getOperand(1);
   Value *Op0 = Shuf->getOperand(0);
   if (!match(Op0, m_InsertElement(m_Undef(), m_Specific(X), m_ZeroInt())))
     return nullptr;
 
   // Replace the shuffle mask element at the index of this insert with a zero.
   // For example:
   // inselt (shuf (inselt undef, X, 0), undef, <0,undef,0,undef>), X, 1
   //   --> shuf (inselt undef, X, 0), undef, <0,0,0,undef>
   unsigned NumMaskElts = Shuf->getType()->getNumElements();
   SmallVector<int, 16> NewMask(NumMaskElts);
   for (unsigned i = 0; i != NumMaskElts; ++i)
     NewMask[i] = i == IdxC ? 0 : Shuf->getMaskValue(i);
 
   return new ShuffleVectorInst(Op0, UndefValue::get(Op0->getType()), NewMask);
 }
 
 /// Try to fold an extract+insert element into an existing identity shuffle by
 /// changing the shuffle's mask to include the index of this insert element.
 static Instruction *foldInsEltIntoIdentityShuffle(InsertElementInst &InsElt) {
   // Check if the vector operand of this insert is an identity shuffle.
   auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0));
   if (!Shuf || !isa<UndefValue>(Shuf->getOperand(1)) ||
       !(Shuf->isIdentityWithExtract() || Shuf->isIdentityWithPadding()))
     return nullptr;
 
+  // Bail out early if shuffle is scalable type. The number of elements in
+  // shuffle mask is unknown at compile-time.
+  if (isa<ScalableVectorType>(Shuf->getType()))
+    return nullptr;
+
   // Check for a constant insertion index.
   uint64_t IdxC;
   if (!match(InsElt.getOperand(2), m_ConstantInt(IdxC)))
     return nullptr;
 
   // Check if this insert's scalar op is extracted from the identity shuffle's
   // input vector.
   Value *Scalar = InsElt.getOperand(1);
   Value *X = Shuf->getOperand(0);
   if (!match(Scalar, m_ExtractElement(m_Specific(X), m_SpecificInt(IdxC))))
     return nullptr;
 
   // Replace the shuffle mask element at the index of this extract+insert with
   // that same index value.
   // For example:
   // inselt (shuf X, IdMask), (extelt X, IdxC), IdxC --> shuf X, IdMask'
   unsigned NumMaskElts = Shuf->getType()->getNumElements();
   SmallVector<int, 16> NewMask(NumMaskElts);
   ArrayRef<int> OldMask = Shuf->getShuffleMask();
   for (unsigned i = 0; i != NumMaskElts; ++i) {
     if (i != IdxC) {
       // All mask elements besides the inserted element remain the same.
       NewMask[i] = OldMask[i];
     } else if (OldMask[i] == (int)IdxC) {
       // If the mask element was already set, there's nothing to do
if (auto *Shuf = dyn_cast<ShuffleVectorInst>(InsElt.getOperand(0))) {
 
     // Create new operands for a shuffle that includes the constant of the
     // original insertelt. The old shuffle will be dead now.
     return new ShuffleVectorInst(Shuf->getOperand(0),
                                  ConstantVector::get(NewShufElts), NewMaskElts);
   } else if (auto *IEI = dyn_cast<InsertElementInst>(Inst)) {
     // Transform sequences of insertelements ops with constant data/indexes into
     // a single shuffle op.
-    unsigned NumElts = InsElt.getType()->getNumElements();
+    // Can not handle scalable type, the number of elements needed to create
+    // shuffle mask is not a compile-time constant.
+    if (isa<ScalableVectorType>(InsElt.getType()))
+      return nullptr;
+    unsigned NumElts =
+        cast<FixedVectorType>(InsElt.getType())->getNumElements();
 
     uint64_t InsertIdx[2];
     Constant *Val[2];
     if (!match(InsElt.getOperand(2), m_ConstantInt(InsertIdx[0])) ||
         !match(InsElt.getOperand(1), m_Constant(Val[0])) ||
         !match(IEI->getOperand(2), m_ConstantInt(InsertIdx[1])) ||
         !match(IEI->getOperand(1), m_Constant(Val[1])))
       return nullptr;
if (match(VecOp, m_BitCast(m_Value(VecSrc))) &&
       cast<VectorType>(VecSrc->getType())->getElementType() ==
           ScalarSrc->getType()) {
     // inselt (bitcast VecSrc), (bitcast ScalarSrc), IdxOp -->
     //   bitcast (inselt VecSrc, ScalarSrc, IdxOp)
     Value *NewInsElt = Builder.CreateInsertElement(VecSrc, ScalarSrc, IdxOp);
     return new BitCastInst(NewInsElt, IE.getType());
   }
 
-  // If the inserted element was extracted from some other vector and both
-  // indexes are valid constants, try to turn this into a shuffle.
+  // If the inserted element was extracted from some other fixed-length vector
+  // and both indexes are valid constants, try to turn this into a shuffle.
+  // Can not handle scalable vector type, the number of elements needed to
+  // create shuffle mask is not a compile-time constant.
   uint64_t InsertedIdx, ExtractedIdx;
   Value *ExtVecOp;
-  if (match(IdxOp, m_ConstantInt(InsertedIdx)) &&
+  if (isa<FixedVectorType>(IE.getType()) &&
+      match(IdxOp, m_ConstantInt(InsertedIdx)) &&
       match(ScalarOp,
             m_ExtractElement(m_Value(ExtVecOp), m_ConstantInt(ExtractedIdx))) &&
-      ExtractedIdx < cast<VectorType>(ExtVecOp->getType())->getNumElements()) {
+      isa<FixedVectorType>(ExtVecOp->getType()) &&
+      ExtractedIdx <
+          cast<FixedVectorType>(ExtVecOp->getType())->getNumElements()) {
     // TODO: Looking at the user(s) to determine if this insert is a
     // fold-to-shuffle opportunity does not match the usual instcombine
     // constraints. We should decide if the transform is worthy based only
     // on this instruction and its operands, but that may not work currently.
     //
     // Here, we are trying to avoid creating shuffles before reaching
     // the end of a chain of extract-insert pairs. This is complicated because
     // we do not generally form arbitrary shuffle masks in instcombine
if (isShuffleRootCandidate(IE)) {
         // We now have a shuffle of LHS, RHS, Mask.
         if (LR.second == nullptr)
           LR.second = UndefValue::get(LR.first->getType());
         return new ShuffleVectorInst(LR.first, LR.second, Mask);
       }
     }
   }
 
-  unsigned VWidth = cast<VectorType>(VecOp->getType())->getNumElements();
+  if (auto VecTy = dyn_cast<FixedVectorType>(VecOp->getType())) {
+    unsigned VWidth = VecTy->getNumElements();
     APInt UndefElts(VWidth, 0);
     APInt AllOnesEltMask(APInt::getAllOnesValue(VWidth));
     if (Value *V = SimplifyDemandedVectorElts(&IE, AllOnesEltMask, UndefElts)) {
       if (V != &IE)
         return replaceInstUsesWith(IE, V);
       return &IE;
     }
+  }
 
   if (Instruction *Shuf = foldConstantInsEltIntoShuffle(IE))
     return Shuf;
 
   if (Instruction *NewInsElt = hoistInsEltConst(IE, Builder))
     return NewInsElt;
 
   if (Instruction *Broadcast = foldInsSequenceIntoSplat(IE))

llvm/test/Transforms/InstCombine/vscale_insertelement.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+; This test checks that bitcast is moved after insertelement when both vector and scalar are
+; bitcast from the same element type.
+; inselt (bitcast VecSrc), (bitcast ScalarSrc), IdxOp
+;  --> bitcast (inselt VecSrc, ScalarSrc, IdxOp)
+define <vscale x 4 x float> @insertelement_bitcast(<vscale x 4 x i32> %vec, i32 %x) {
+; CHECK-LABEL: @insertelement_bitcast(
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <vscale x 4 x i32> [[VEC:%.*]], i32 [[X:%.*]], i32 0
+; CHECK-NEXT:    [[R:%.*]] = bitcast <vscale x 4 x i32> [[TMP1]] to <vscale x 4 x float>
+; CHECK-NEXT:    ret <vscale x 4 x float> [[R]]
+;
+  %x_cast = bitcast i32 %x to float
+  %vec_cast = bitcast <vscale x 4 x i32> %vec to <vscale x 4 x float>
+  %r = insertelement <vscale x 4 x float> %vec_cast, float %x_cast, i32 0
+  ret <vscale x 4 x float> %r
+}
+
+; This test checks that code-path "Try to form a shuffle from a chain of extract-insert ops" is
+; not taken when both extract and insert are scalable type.
+; For scalable type, the vector length needed to create shuffle mask is not a compile-time constant.
+; Meanwhile, for scalable type shuffle mask only support splat and undef in the current code base.
+; Otherwise we crash at:
+; "Assertion `isValidOperands(V1, V2, Mask) && "Invalid shuffle vector instruction operands!"' failed."
+define <vscale x 4 x i32> @insertelement_extractelement(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
+; CHECK-LABEL: @insertelement_extractelement(
+; CHECK-NEXT:    [[T0:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i32 1
+; CHECK-NEXT:    [[T1:%.*]] = insertelement <vscale x 4 x i32> [[B:%.*]], i32 [[T0]], i32 0
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[T1]]
+;
+  %t0 = extractelement <vscale x 4 x i32> %a, i32 1
+  %t1 = insertelement <vscale x 4 x i32> %b, i32 %t0, i32 0
+  ret <vscale x 4 x i32> %t1
+}
+
+; This test checks that we are not attempting to create a shuffle from extract/insert chain,
+; when extract is from a scalable type, and the insert vector is fixed-length.
+define <4 x i32> @insertelement_extractelement_fixed_vec_extract_from_scalable(<vscale x 4 x i32> %a, <4 x i32> %b) {
+; CHECK-LABEL: @insertelement_extractelement_fixed_vec_extract_from_scalable(
+; CHECK-NEXT:    [[T0:%.*]] = extractelement <vscale x 4 x i32> [[A:%.*]], i32 1
+; CHECK-NEXT:    [[T1:%.*]] = insertelement <4 x i32> [[B:%.*]], i32 [[T0]], i32 0
+; CHECK-NEXT:    ret <4 x i32> [[T1]]
+;
+  %t0 = extractelement <vscale x 4 x i32> %a, i32 1
+  %t1 = insertelement <4 x i32> %b, i32 %t0, i32 0
+  ret <4 x i32> %t1
+}
+
+; This test checks that the optimization "foldConstantInsEltInfoShuffle" is not taken for scalable type.
+; Particularly the fold:
+; insertelt (insertelt X, C1, CIndex1), C, CIndex
+;  --> shufflevector X, CVec', Mask'
+; For scalable type, the vector length needed to create shuffle mask is not a compile-time constant.
+; Meanwhile, for scalable type shuffle mask only support splat and undef in the current code base.
+; Otherwise we crash at:
+; "Assertion `isValidOperands(V1, V2, Mask) && "Invalid shuffle vector instruction operands!"' failed."
+define <vscale x 4 x i32> @insertelement_insertelement(<vscale x 4 x i32> %vec) {
+; CHECK-LABEL: @insertelement_insertelement(
+; CHECK-NEXT:    [[T0:%.*]] = insertelement <vscale x 4 x i32> [[VEC:%.*]], i32 1, i32 1
+; CHECK-NEXT:    [[T1:%.*]] = insertelement <vscale x 4 x i32> [[T0]], i32 2, i32 2
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[T1]]
+;
+  %t0 = insertelement <vscale x 4 x i32> %vec, i32 1, i32 1
+  %t1 = insertelement <vscale x 4 x i32> %t0, i32 2, i32 2
+  ret <vscale x 4 x i32> %t1
+}
+
+; This test checks that the following insertelement sequence is not folded into shuffle splat.
+; The length of scalable vector is unknown at compile-time. Therefore the following insertelements
+; may not form a valid splat.
+define <vscale x 4 x float> @insertelement_sequene_may_not_be_splat(float %x) {
+; CHECK-LABEL: @insertelement_sequene_may_not_be_splat(
+; CHECK-NEXT:    [[T0:%.*]] = insertelement <vscale x 4 x float> undef, float [[X:%.*]], i32 0
+; CHECK-NEXT:    [[T1:%.*]] = insertelement <vscale x 4 x float> [[T0]], float [[X]], i32 1
+; CHECK-NEXT:    [[T2:%.*]] = insertelement <vscale x 4 x float> [[T1]], float [[X]], i32 2
+; CHECK-NEXT:    [[T3:%.*]] = insertelement <vscale x 4 x float> [[T2]], float [[X]], i32 3
+; CHECK-NEXT:    ret <vscale x 4 x float> [[T3]]
+;
+  %t0 = insertelement <vscale x 4 x float> undef, float %x, i32 0
+  %t1 = insertelement <vscale x 4 x float> %t0, float %x, i32 1
+  %t2 = insertelement <vscale x 4 x float> %t1, float %x, i32 2
+  %t3 = insertelement <vscale x 4 x float> %t2, float %x, i32 3
+  ret <vscale x 4 x float> %t3
+}