[SLP]Fix narrow-tree gate width and extract miscount

Measure the narrowness of in-loop trees by the actual vectorization
width instead of the feeder-load width, and skip vector-typed scalars
in the instruction count check to match the cost model.

Fixes llvm#216715

Reviewers: efriedma-quic, bababuck, RKSimon, kartcq

Pull Request: llvm#216797

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -13773,6 +13773,12 @@ uint64_t BoUpSLP::getNumVectorInsts(bool HasTreeLoop) {
       continue;
     if (EU.User && EphValues.count(EU.User))
       continue;
+    // Vector-typed scalars are consumed as whole vectors, no extract
+    // instruction is emitted for them.
+    if (isVectorizedTy(EU.Scalar->getType()) &&
+        (!SLPReVec ||
+         (EU.E.hasState() && EU.E.getOpcode() == Instruction::InsertElement)))
+      continue;
     if (ExternalUsesAsOriginalScalar.contains(EU.Scalar))
       continue;
     if (!CountedExtracts.insert(EU.Scalar).second)
@@ -19289,7 +19295,17 @@ BoUpSLP::calculateTreeCostAndTrimNonProfitable(ArrayRef<Value *> VectorizedVals,
     return Cost;
   // The narrow non-profitable tree in loop? Skip, may cause regressions.
   constexpr unsigned PartLimit = 2;
-  const unsigned Sz = getVectorElementSize(getRootNodeScalars().front());
+  // Measure the narrowness by the width the tree is vectorized at.
+  const TreeEntry *RootTE = &getRootNode();
+  unsigned Sz;
+  if (auto It = MinBWs.find(RootTE); It != MinBWs.end()) {
+    Sz = It->second.first;
+  } else {
+    Sz = std::max<unsigned>(
+        getVectorElementSize(RootTE->Scalars.front()),
+        DL->getTypeSizeInBits(
+            getValueType(RootTE->Scalars.front())->getScalarType()));
+  }
   const unsigned MinVF = getMinVF(Sz);
   if (Cost >= -SLPCostThreshold &&
       getRootNodeScalars().size() * PartLimit <= MinVF &&

llvm/test/Transforms/SLPVectorizer/AArch64/abs-mul-buildvector-in-loop.ll

@@ -17,34 +17,35 @@ define void @test(ptr noalias %src, ptr noalias %descriptor, i32 %N) {
 ; CHECK-NEXT:    [[I_033:%.*]] = phi i32 [ [[INC:%.*]], %[[FOR_BODY]] ], [ 0, %[[ENTRY]] ]
 ; CHECK-NEXT:    [[ARRAYIDX]] = getelementptr inbounds nuw i8, ptr [[SRC_ADDR_035]], i64 1
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i8, ptr [[ARRAYIDX]], align 1
-; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[TMP0]] to i32
 ; CHECK-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i8, ptr [[SRC_ADDR_035]], i64 -1
 ; CHECK-NEXT:    [[TMP1:%.*]] = load i8, ptr [[ARRAYIDX1]], align 1
-; CHECK-NEXT:    [[CONV2:%.*]] = zext i8 [[TMP1]] to i32
-; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i32 [[CONV]], [[CONV2]]
 ; CHECK-NEXT:    [[ARRAYIDX3:%.*]] = getelementptr inbounds nuw i8, ptr [[SRC_ADDR_035]], i64 17
 ; CHECK-NEXT:    [[TMP2:%.*]] = load i8, ptr [[ARRAYIDX3]], align 1
-; CHECK-NEXT:    [[CONV4:%.*]] = zext i8 [[TMP2]] to i32
 ; CHECK-NEXT:    [[ARRAYIDX5:%.*]] = getelementptr inbounds i8, ptr [[SRC_ADDR_035]], i64 -17
 ; CHECK-NEXT:    [[TMP3:%.*]] = load i8, ptr [[ARRAYIDX5]], align 1
-; CHECK-NEXT:    [[CONV6:%.*]] = zext i8 [[TMP3]] to i32
-; CHECK-NEXT:    [[SUB7:%.*]] = sub nsw i32 [[CONV4]], [[CONV6]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = load i8, ptr [[P_034]], align 1
 ; CHECK-NEXT:    [[CONV9:%.*]] = zext i8 [[TMP4]] to i32
 ; CHECK-NEXT:    [[INCDEC_PTR10]] = getelementptr inbounds nuw i8, ptr [[P_034]], i64 2
 ; CHECK-NEXT:    [[BIN:%.*]] = getelementptr inbounds nuw i8, ptr [[P_034]], i64 1
 ; CHECK-NEXT:    [[TMP5:%.*]] = load i8, ptr [[BIN]], align 1
 ; CHECK-NEXT:    [[IDX_EXT:%.*]] = zext i8 [[TMP5]] to i64
 ; CHECK-NEXT:    [[ADD_PTR:%.*]] = getelementptr inbounds nuw [4 x i8], ptr [[DESCRIPTOR]], i64 [[IDX_EXT]]
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[SUB]], [[CONV9]]
-; CHECK-NEXT:    [[MUL12:%.*]] = mul nsw i32 [[SUB7]], [[CONV9]]
-; CHECK-NEXT:    [[TMP6:%.*]] = tail call i32 @llvm.abs.i32(i32 [[MUL]], i1 true)
-; CHECK-NEXT:    [[TMP7:%.*]] = tail call i32 @llvm.abs.i32(i32 [[MUL12]], i1 true)
+; CHECK-NEXT:    [[CONV4:%.*]] = zext i8 [[TMP2]] to i32
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[TMP0]] to i32
+; CHECK-NEXT:    [[CONV6:%.*]] = zext i8 [[TMP3]] to i32
+; CHECK-NEXT:    [[CONV2:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    [[SUB7:%.*]] = sub nsw i32 [[CONV4]], [[CONV6]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i32 [[CONV]], [[CONV2]]
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <2 x i32> poison, i32 [[SUB]], i64 0
+; CHECK-NEXT:    [[TMP7:%.*]] = insertelement <2 x i32> [[TMP6]], i32 [[SUB7]], i64 1
+; CHECK-NEXT:    [[TMP15:%.*]] = insertelement <2 x i32> poison, i32 [[CONV9]], i64 0
+; CHECK-NEXT:    [[TMP9:%.*]] = shufflevector <2 x i32> [[TMP15]], <2 x i32> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP10:%.*]] = mul nsw <2 x i32> [[TMP7]], [[TMP9]]
+; CHECK-NEXT:    [[TMP11:%.*]] = call <2 x i32> @llvm.abs.v2i32(<2 x i32> [[TMP10]], i1 true)
 ; CHECK-NEXT:    [[TMP8:%.*]] = load <4 x i32>, ptr [[ADD_PTR]], align 4
-; CHECK-NEXT:    [[TMP9:%.*]] = insertelement <4 x i32> poison, i32 [[MUL]], i64 0
-; CHECK-NEXT:    [[TMP10:%.*]] = insertelement <4 x i32> [[TMP9]], i32 [[MUL12]], i64 1
-; CHECK-NEXT:    [[TMP11:%.*]] = insertelement <4 x i32> [[TMP10]], i32 [[TMP6]], i64 2
-; CHECK-NEXT:    [[TMP12:%.*]] = insertelement <4 x i32> [[TMP11]], i32 [[TMP7]], i64 3
+; CHECK-NEXT:    [[TMP16:%.*]] = shufflevector <2 x i32> [[TMP10]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP14:%.*]] = shufflevector <2 x i32> [[TMP11]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP12:%.*]] = shufflevector <4 x i32> [[TMP16]], <4 x i32> [[TMP14]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
 ; CHECK-NEXT:    [[TMP13:%.*]] = add nsw <4 x i32> [[TMP8]], [[TMP12]]
 ; CHECK-NEXT:    store <4 x i32> [[TMP13]], ptr [[ADD_PTR]], align 4
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[I_033]], 1

llvm/test/Transforms/SLPVectorizer/AArch64/long-non-power-of-2.ll

@@ -7,6 +7,8 @@ define i1 @test(ptr %arg, ptr %arg1, i64 %arg2, ptr %arg3) {
 ; CHECK-NEXT:  [[BB:.*:]]
 ; CHECK-NEXT:    [[GETELEMENTPTR:%.*]] = getelementptr i8, ptr [[ARG1]], i64 [[ARG2]]
 ; CHECK-NEXT:    [[GETELEMENTPTR4:%.*]] = getelementptr i8, ptr null, i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x ptr> <ptr null, ptr poison>, ptr [[ARG3]], i64 1
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, <2 x ptr> [[TMP2]], <2 x i64> <i64 -32, i64 -432>
 ; CHECK-NEXT:    [[GETELEMENTPTR5:%.*]] = getelementptr i8, ptr null, i64 -32
 ; CHECK-NEXT:    [[TMP32:%.*]] = getelementptr i8, ptr [[ARG3]], i64 -440
 ; CHECK-NEXT:    [[GETELEMENTPTR7:%.*]] = getelementptr i8, ptr [[ARG1]], i64 0
@@ -28,8 +30,8 @@ define i1 @test(ptr %arg, ptr %arg1, i64 %arg2, ptr %arg3) {
 ; CHECK-NEXT:    [[AND25:%.*]] = and i1 false, [[ICMP24]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x ptr> <ptr poison, ptr poison, ptr null, ptr null>, ptr [[ARG]], i64 1
 ; CHECK-NEXT:    [[TMP19:%.*]] = insertelement <4 x ptr> [[TMP0]], ptr [[GETELEMENTPTR7]], i64 0
-; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x ptr> <ptr poison, ptr poison, ptr null, ptr null>, ptr [[GETELEMENTPTR5]], i64 0
-; CHECK-NEXT:    [[TMP7:%.*]] = insertelement <4 x ptr> [[TMP2]], ptr [[GETELEMENTPTR8]], i64 1
+; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <2 x ptr> [[TMP1]], <2 x ptr> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <4 x ptr> <ptr poison, ptr poison, ptr null, ptr null>, <4 x ptr> [[TMP4]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
 ; CHECK-NEXT:    [[TMP20:%.*]] = icmp ult <4 x ptr> [[TMP19]], [[TMP7]]
 ; CHECK-NEXT:    [[ICMP58:%.*]] = icmp ult ptr [[GETELEMENTPTR]], null
 ; CHECK-NEXT:    [[ICMP62:%.*]] = icmp ult ptr [[GETELEMENTPTR]], null