[InstCombine] Remove scalable vector restriction in foldVectorBinop

[junparser]

Summary

As is the title.

TestPlan: check-llvm

[RKSimon]

llvm/test/Transforms/InstCombine/fold-bin-operand.ll

58

please can you run this through the update_test_checks.py script?

[junparser]

address the comment, update the testcase.

[RKSimon]

LGTM with one minor

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1607–1608

Might be worth pulling out the repeated isa/cast on Inst.getType():

auto *InstVTy = dyn_cast<FixedVectorType>(Inst.getType());

[junparser]

In D93289#2454658, @RKSimon wrote:

LGTM with one minor

Thanks for review. I'll update when I commit the patch.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (auto *BC = dyn_cast<BitCastInst>(&Op)) {
     VectorType *DestTy = dyn_cast<VectorType>(BC->getDestTy());
     VectorType *SrcTy = dyn_cast<VectorType>(BC->getSrcTy());
 
     // Verify that either both or neither are vectors.
     if ((SrcTy == nullptr) != (DestTy == nullptr))
       return nullptr;
 
     // If vectors, verify that they have the same number of elements.
-    if (SrcTy && cast<FixedVectorType>(SrcTy)->getNumElements() !=
-                     cast<FixedVectorType>(DestTy)->getNumElements())
+    if (SrcTy && SrcTy->getElementCount() != DestTy->getElementCount())
       return nullptr;
   }
 
   // Test if a CmpInst instruction is used exclusively by a select as
   // part of a minimum or maximum operation. If so, refrain from doing
   // any other folding. This helps out other analyses which understand
   // non-obfuscated minimum and maximum idioms, such as ScalarEvolution
   // and CodeGen. And in this case, at least one of the comparison
do {
 
     // Move up one level in the expression.
     assert(Ancestor->hasOneUse() && "Drilled down when more than one use!");
     Ancestor = Ancestor->user_back();
   } while (true);
 }
 
 Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) {
-  // FIXME: some of this is likely fine for scalable vectors
-  if (!isa<FixedVectorType>(Inst.getType()))
+  if (!isa<VectorType>(Inst.getType()))
     return nullptr;
 
   BinaryOperator::BinaryOps Opcode = Inst.getOpcode();
   Value *LHS = Inst.getOperand(0), *RHS = Inst.getOperand(1);
   assert(cast<VectorType>(LHS->getType())->getElementCount() ==
          cast<VectorType>(Inst.getType())->getElementCount());
   assert(cast<VectorType>(RHS->getType())->getElementCount() ==
          cast<VectorType>(Inst.getType())->getElementCount());
if (Inst.isCommutative() &&
     }
   }
 
   // If one argument is a shuffle within one vector and the other is a constant,
   // try moving the shuffle after the binary operation. This canonicalization
   // intends to move shuffles closer to other shuffles and binops closer to
   // other binops, so they can be folded. It may also enable demanded elements
   // transforms.
-  unsigned NumElts = cast<FixedVectorType>(Inst.getType())->getNumElements();
   Constant *C;
-  if (match(&Inst,
+  if (isa<FixedVectorType>(Inst.getType()) &&
+      match(&Inst,
             m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Undef(), m_Mask(Mask))),
-                      m_Constant(C))) && !isa<ConstantExpr>(C) &&
-      cast<FixedVectorType>(V1->getType())->getNumElements() <= NumElts) {
+                      m_Constant(C))) &&
+      !isa<ConstantExpr>(C) &&
+      cast<FixedVectorType>(V1->getType())->getNumElements() <=
+          cast<FixedVectorType>(Inst.getType())->getNumElements()) {
     assert(Inst.getType()->getScalarType() == V1->getType()->getScalarType() &&
            "Shuffle should not change scalar type");
 
     // Find constant NewC that has property:
     //   shuffle(NewC, ShMask) = C
     // If such constant does not exist (example: ShMask=<0,0> and C=<1,2>)
     // reorder is not possible. A 1-to-1 mapping is not required. Example:
     // ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <undef,5,6,undef>
     bool ConstOp1 = isa<Constant>(RHS);
     ArrayRef<int> ShMask = Mask;
     unsigned SrcVecNumElts =
         cast<FixedVectorType>(V1->getType())->getNumElements();
     UndefValue *UndefScalar = UndefValue::get(C->getType()->getScalarType());
     SmallVector<Constant *, 16> NewVecC(SrcVecNumElts, UndefScalar);
     bool MayChange = true;
+    unsigned NumElts = cast<FixedVectorType>(Inst.getType())->getNumElements();
     for (unsigned I = 0; I < NumElts; ++I) {
       Constant *CElt = C->getAggregateElement(I);
       if (ShMask[I] >= 0) {
         assert(ShMask[I] < (int)NumElts && "Not expecting narrowing shuffle");
         Constant *NewCElt = NewVecC[ShMask[I]];
         // Bail out if:
         // 1. The constant vector contains a constant expression.
         // 2. The shuffle needs an element of the constant vector that can't
if (auto *BCI = dyn_cast<BitCastInst>(ASCStrippedPtrOp)) {
     auto areMatchingArrayAndVecTypes = [](Type *ArrTy, Type *VecTy,
                                           const DataLayout &DL) {
       auto *VecVTy = cast<FixedVectorType>(VecTy);
       return ArrTy->getArrayElementType() == VecVTy->getElementType() &&
              ArrTy->getArrayNumElements() == VecVTy->getNumElements() &&
              DL.getTypeAllocSize(ArrTy) == DL.getTypeAllocSize(VecTy);
     };
     if (GEP.getNumOperands() == 3 &&
-        ((GEPEltType->isArrayTy() && SrcEltType->isVectorTy() &&
+        ((GEPEltType->isArrayTy() && isa<FixedVectorType>(SrcEltType) &&
           areMatchingArrayAndVecTypes(GEPEltType, SrcEltType, DL)) ||
-         (GEPEltType->isVectorTy() && SrcEltType->isArrayTy() &&
+         (isa<FixedVectorType>(GEPEltType) && SrcEltType->isArrayTy() &&
           areMatchingArrayAndVecTypes(SrcEltType, GEPEltType, DL)))) {
 
       // Create a new GEP here, as using `setOperand()` followed by
       // `setSourceElementType()` won't actually update the type of the
       // existing GEP Value. Causing issues if this Value is accessed when
       // constructing an AddrSpaceCastInst
       Value *NGEP =
           GEP.isInBounds()

llvm/test/Transforms/InstCombine/fold-bin-operand.ll

 
 define i32 @g(i32 %x) {
 ; CHECK-LABEL: @g(
 ; CHECK: ret i32 %x
   %b = add i32 %x, zext (i1 icmp eq (i8* inttoptr (i32 1000000 to i8*), i8* inttoptr (i32 2000000 to i8*)) to i32)
   ret i32 %b
 }
 
+define i32 @h(i1 %A, i32 %B) {
+EntryBlock:
+  %cf = select i1 %A, i32 1, i32 %B
+  %op = add i32 2, %cf
+  ret i32 %op
+; CHECK-LABEL: @h(
+; CHECK: [[OP:%.*]] = add i32 %B, 2
+; CHECK: select i1 %A, i32 3, i32 [[OP]]
+}
+
+define <4 x float> @h1(i1 %A, <4 x i32> %B) {
+EntryBlock:
+  %cf = select i1 %A, <4 x i32> <i32 1, i32 1, i32 1, i32 1>, <4 x i32> %B
+  %bc = bitcast <4 x i32> %cf to <4 x float>
+  ret <4 x float> %bc
+; CHECK-LABEL: @h1(
+; CHECK: [[OP:%.*]] = bitcast <4 x i32> %B to <4 x float>
+; CHECK: select i1 %A, <4 x float> <float 0x36A0000000000000, float 0x36A0000000000000, float 0x36A0000000000000, float 0x36A0000000000000>, <4 x float> [[OP]]
+}
+
+define <vscale x 4 x float> @h2(i1 %A, <vscale x 4 x i32> %B) {
+EntryBlock:
+  %cf = select i1 %A, <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> %B
+  %bc = bitcast <vscale x 4 x i32> %cf to <vscale x 4 x float>
+  ret <vscale x 4 x float> %bc
+; CHECK-LABEL: @h2(
+; CHECK: [[OP:%.*]] = bitcast <vscale x 4 x i32> %B to <vscale x 4 x float>
+; CHECK: select i1 %A, <vscale x 4 x float> zeroinitializer, <vscale x 4 x float> [[OP]]
+}
+
+define <vscale x 2 x i64> @h3(i1 %A, <vscale x 4 x i32> %B) {
+EntryBlock:
+  %cf = select i1 %A, <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> %B
+  %bc = bitcast <vscale x 4 x i32> %cf to <vscale x 2 x i64>
+  ret <vscale x 2 x i64> %bc
+
+; CHECK-LABEL: @h3(
+; CHECK: [[CF:%.*]] = select i1 %A, <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> %B
+; CHECK: [[BC:%.*]] = bitcast <vscale x 4 x i32> [[CF]] to <vscale x 2 x i64>
+; CHECK: ret <vscale x 2 x i64> [[BC]]
+}
+

llvm/test/Transforms/InstCombine/vec-binop-select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; Non-canonical mask
 
 define <4 x i32> @and(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @and(
 ; CHECK-NEXT:    [[R:%.*]] = and <4 x i32> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <4 x i32> [[R]]
 ;
   %sel1 = shufflevector <4 x i32> %x, <4 x i32> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   %sel2 = shufflevector <4 x i32> %x, <4 x i32> %y, <4 x i32> <i32 4, i32 1, i32 2, i32 7>
   %r = and <4 x i32> %sel1, %sel2
   ret <4 x i32> %r
 }
 
+define <vscale x 4 x i32> @vscaleand(<vscale x 4 x i32> %x, <vscale x 4 x i32> %y) {
+; CHECK-LABEL: @vscaleand(
+; CHECK-NEXT:    [[R:%.*]] = and <vscale x 4 x i32> [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[S:%.*]] = shufflevector <vscale x 4 x i32> [[R]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[S]]
+;
+  %sel1 = shufflevector <vscale x 4 x i32> %x, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %sel2 = shufflevector <vscale x 4 x i32> %y, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = and <vscale x 4 x i32> %sel1, %sel2
+  ret <vscale x 4 x i32> %r
+}
+
 define <4 x i32> @or(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @or(
 ; CHECK-NEXT:    [[R:%.*]] = or <4 x i32> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <4 x i32> [[R]]
 ;
   %sel1 = shufflevector <4 x i32> %x, <4 x i32> %y, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   %sel2 = shufflevector <4 x i32> %y, <4 x i32> %x, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   %r = or <4 x i32> %sel1, %sel2

llvm/test/Transforms/InstCombine/vec_shuffle.ll

 ; CHECK-NEXT:    ret <4 x i32> [[R]]
 ;
   %splatx = shufflevector <4 x i32> %x, <4 x i32> undef, <4 x i32> zeroinitializer
   %a = add <4 x i32> %y, <i32 317426, i32 317426, i32 317426, i32 317426>
   %r = add <4 x i32> %splatx, %a
   ret <4 x i32> %r
 }
 
+define <vscale x 4 x i32> @vsplat_assoc_add(<vscale x 4 x i32> %x, <vscale x 4 x i32> %y) {
+; CHECK-LABEL: @vsplat_assoc_add(
+; CHECK-NEXT:    [[TMP1:%.*]] = add <vscale x 4 x i32> [[X:%.*]], shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 317426, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <vscale x 4 x i32> [[TMP1]], <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = add <vscale x 4 x i32> [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret <vscale x 4 x i32> [[R]]
+;
+
+  %splatx = shufflevector <vscale x 4 x i32> %x, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %a = add <vscale x 4 x i32> %y, shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 317426, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+  %r = add <vscale x 4 x i32> %splatx, %a
+  ret <vscale x 4 x i32> %r
+}
+
 ; Undefs in splat mask are replaced with defined splat index
 
 define <4 x i32> @splat_assoc_add_undef_mask_elts(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @splat_assoc_add_undef_mask_elts(
 ; CHECK-NEXT:    [[TMP1:%.*]] = add <4 x i32> [[X:%.*]], <i32 42, i32 undef, i32 undef, i32 undef>
 ; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x i32> [[TMP1]], <4 x i32> undef, <4 x i32> zeroinitializer
 ; CHECK-NEXT:    [[R:%.*]] = add <4 x i32> [[TMP2]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <4 x i32> [[R]]