[SLP] Use getExtractWithExtendCost() to compute the scalar cost of extractelement/ext pair

[haicheng]

Summary

We use getExtractWithExtendCost to calculate the cost of extractelement and s|zext together when computing the extract cost after vectorization, but we calculate the cost of extractelement and s|zext separately when computing the scalar cost which is larger than it should be.

[ABataev]

test/Transforms/SLPVectorizer/AArch64/ext-trunc.ll

2 ↗

(On Diff #141772)

Could you regenerate test checks without your patch and without -instcombine first and commit as NFC? And after that generate these checks with your patch

[haicheng]

Rebased the test case. Thank you for taking a look.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2199 ↗

(On Diff #142022)

I think you need to check that Ext user is GEP and Ext is used as GEP index.

[haicheng]

Check the users of Ext now.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2201 ↗

(On Diff #142063)

I think you can use just GetElementPtr::classof here as a predicate. Also, it would be good to have negative test(s) for the situation where this should not work.

[haicheng]

Address Alexey's comments.

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2201 ↗

(On Diff #142063)

No, I meant that instead of lambda you can just use GetElementPtrInst::classof

llvm::all_of(Ext->users(), GetElementPtrInst::classof)

[haicheng]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2201 ↗

(On Diff #142063)

I tried, but I got this error when compiling

llvm/ADT/STLExtras.h:846:6: note: candidate template ignored: couldn't infer template argument 'UnaryPredicate'

[ABataev]

lib/Transforms/Vectorize/SLPVectorizer.cpp

2201 ↗

(On Diff #142063)

Ok then, revert back to isa

[haicheng]

Address Alexey's comments. Thank you.

[ABataev]

LG

llvm/trunk/lib/Transforms/Vectorize/SLPVectorizer.cpp

case Instruction::ExtractElement:
           DeadCost += TTI->getShuffleCost(
               TargetTransformInfo::SK_PermuteSingleSrc, VecTy);
         }
         for (unsigned i = 0, e = VL.size(); i < e; ++i) {
           Instruction *E = cast<Instruction>(VL[i]);
           // If all users are going to be vectorized, instruction can be
           // considered as dead.
           // The same, if have only one user, it will be vectorized for sure.
-          if (areAllUsersVectorized(E))
+          if (areAllUsersVectorized(E)) {
             // Take credit for instruction that will become dead.
+            if (E->hasOneUse()) {
+              Instruction *Ext = E->user_back();
+              if ((isa<SExtInst>(Ext) || isa<ZExtInst>(Ext)) &&
+                  all_of(Ext->users(),
+                         [](User *U) { return isa<GetElementPtrInst>(U); })) {
+                // Use getExtractWithExtendCost() to calculate the cost of
+                // extractelement/ext pair.
+                DeadCost -= TTI->getExtractWithExtendCost(
+                    Ext->getOpcode(), Ext->getType(), VecTy, i);
+                // Add back the cost of s|zext which is subtracted seperately.
+                DeadCost += TTI->getCastInstrCost(
+                    Ext->getOpcode(), Ext->getType(), E->getType(), Ext);
+                continue;
+              }
+            }
             DeadCost -=
                 TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, i);
           }
+        }
         return DeadCost;
       }
       return ReuseShuffleCost + getGatherCost(VL);
 
     case Instruction::ZExt:
     case Instruction::SExt:
     case Instruction::FPToUI:
     case Instruction::FPToSI:

llvm/trunk/test/Transforms/SLPVectorizer/AArch64/ext-trunc.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -slp-vectorizer -mtriple=aarch64--linux-gnu < %s | FileCheck %s
 
 target datalayout = "e-m:e-i32:64-i128:128-n32:64-S128"
 
 declare void @foo(i64, i64, i64, i64)
 
-define void @test(<4 x i16> %a, <4 x i16> %b, i64* %p) {
+define void @test1(<4 x i16> %a, <4 x i16> %b, i64* %p) {
 ; Make sure types of sub and its sources are not extended.
-; CHECK-LABEL: @test(
+; CHECK-LABEL: @test1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[Z0:%.*]] = zext <4 x i16> [[A:%.*]] to <4 x i32>
 ; CHECK-NEXT:    [[Z1:%.*]] = zext <4 x i16> [[B:%.*]] to <4 x i32>
 ; CHECK-NEXT:    [[SUB0:%.*]] = sub <4 x i32> [[Z0]], [[Z1]]
-; CHECK-NEXT:    [[TMP0:%.*]] = sext <4 x i32> [[SUB0]] to <4 x i64>
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc <4 x i64> [[TMP0]] to <4 x i32>
-; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <4 x i32> [[TMP1]], i32 0
-; CHECK-NEXT:    [[TMP3:%.*]] = sext i32 [[TMP2]] to i64
-; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds i64, i64* [[P:%.*]], i64 [[TMP3]]
+; CHECK-NEXT:    [[E0:%.*]] = extractelement <4 x i32> [[SUB0]], i32 0
+; CHECK-NEXT:    [[S0:%.*]] = sext i32 [[E0]] to i64
+; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds i64, i64* [[P:%.*]], i64 [[S0]]
 ; CHECK-NEXT:    [[LOAD0:%.*]] = load i64, i64* [[GEP0]]
-; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <4 x i32> [[TMP1]], i32 1
-; CHECK-NEXT:    [[TMP5:%.*]] = sext i32 [[TMP4]] to i64
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP5]]
+; CHECK-NEXT:    [[E1:%.*]] = extractelement <4 x i32> [[SUB0]], i32 1
+; CHECK-NEXT:    [[S1:%.*]] = sext i32 [[E1]] to i64
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[S1]]
 ; CHECK-NEXT:    [[LOAD1:%.*]] = load i64, i64* [[GEP1]]
-; CHECK-NEXT:    [[TMP6:%.*]] = extractelement <4 x i32> [[TMP1]], i32 2
-; CHECK-NEXT:    [[TMP7:%.*]] = sext i32 [[TMP6]] to i64
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP7]]
+; CHECK-NEXT:    [[E2:%.*]] = extractelement <4 x i32> [[SUB0]], i32 2
+; CHECK-NEXT:    [[S2:%.*]] = sext i32 [[E2]] to i64
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[S2]]
 ; CHECK-NEXT:    [[LOAD2:%.*]] = load i64, i64* [[GEP2]]
-; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <4 x i32> [[TMP1]], i32 3
-; CHECK-NEXT:    [[TMP9:%.*]] = sext i32 [[TMP8]] to i64
-; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP9]]
+; CHECK-NEXT:    [[E3:%.*]] = extractelement <4 x i32> [[SUB0]], i32 3
+; CHECK-NEXT:    [[S3:%.*]] = sext i32 [[E3]] to i64
+; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[S3]]
 ; CHECK-NEXT:    [[LOAD3:%.*]] = load i64, i64* [[GEP3]]
 ; CHECK-NEXT:    call void @foo(i64 [[LOAD0]], i64 [[LOAD1]], i64 [[LOAD2]], i64 [[LOAD3]])
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %z0 = zext <4 x i16> %a to <4 x i32>
   %z1 = zext <4 x i16> %b to <4 x i32>
   %sub0 = sub <4 x i32> %z0, %z1
entry:
   %load2 = load i64, i64* %gep2
   %e3 = extractelement <4 x i32> %sub0, i32 3
   %s3 = sext i32 %e3 to i64
   %gep3 = getelementptr inbounds i64, i64* %p, i64 %s3
   %load3 = load i64, i64* %gep3
   call void @foo(i64 %load0, i64 %load1, i64 %load2, i64 %load3)
   ret void
 }
+
+define void @test2(<4 x i16> %a, <4 x i16> %b, i64 %c0, i64 %c1, i64 %c2, i64 %c3, i64* %p) {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[Z0:%.*]] = zext <4 x i16> [[A:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[Z1:%.*]] = zext <4 x i16> [[B:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[SUB0:%.*]] = sub <4 x i32> [[Z0]], [[Z1]]
+; CHECK-NEXT:    [[TMP0:%.*]] = sext <4 x i32> [[SUB0]] to <4 x i64>
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i64> undef, i64 [[C0:%.*]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x i64> [[TMP1]], i64 [[C1:%.*]], i32 1
+; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <4 x i64> [[TMP2]], i64 [[C2:%.*]], i32 2
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <4 x i64> [[TMP3]], i64 [[C3:%.*]], i32 3
+; CHECK-NEXT:    [[TMP5:%.*]] = add <4 x i64> [[TMP4]], [[TMP0]]
+; CHECK-NEXT:    [[TMP6:%.*]] = extractelement <4 x i64> [[TMP5]], i32 0
+; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds i64, i64* [[P:%.*]], i64 [[TMP6]]
+; CHECK-NEXT:    [[LOAD0:%.*]] = load i64, i64* [[GEP0]]
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <4 x i64> [[TMP5]], i32 1
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP7]]
+; CHECK-NEXT:    [[LOAD1:%.*]] = load i64, i64* [[GEP1]]
+; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <4 x i64> [[TMP5]], i32 2
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP8]]
+; CHECK-NEXT:    [[LOAD2:%.*]] = load i64, i64* [[GEP2]]
+; CHECK-NEXT:    [[TMP9:%.*]] = extractelement <4 x i64> [[TMP5]], i32 3
+; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i64, i64* [[P]], i64 [[TMP9]]
+; CHECK-NEXT:    [[LOAD3:%.*]] = load i64, i64* [[GEP3]]
+; CHECK-NEXT:    call void @foo(i64 [[LOAD0]], i64 [[LOAD1]], i64 [[LOAD2]], i64 [[LOAD3]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %z0 = zext <4 x i16> %a to <4 x i32>
+  %z1 = zext <4 x i16> %b to <4 x i32>
+  %sub0 = sub <4 x i32> %z0, %z1
+  %e0 = extractelement <4 x i32> %sub0, i32 0
+  %s0 = sext i32 %e0 to i64
+  %a0 = add i64 %s0, %c0
+  %gep0 = getelementptr inbounds i64, i64* %p, i64 %a0
+  %load0 = load i64, i64* %gep0
+  %e1 = extractelement <4 x i32> %sub0, i32 1
+  %s1 = sext i32 %e1 to i64
+  %a1 = add i64 %s1, %c1
+  %gep1 = getelementptr inbounds i64, i64* %p, i64 %a1
+  %load1 = load i64, i64* %gep1
+  %e2 = extractelement <4 x i32> %sub0, i32 2
+  %s2 = sext i32 %e2 to i64
+  %a2 = add i64 %s2, %c2
+  %gep2 = getelementptr inbounds i64, i64* %p, i64 %a2
+  %load2 = load i64, i64* %gep2
+  %e3 = extractelement <4 x i32> %sub0, i32 3
+  %s3 = sext i32 %e3 to i64
+  %a3 = add i64 %s3, %c3
+  %gep3 = getelementptr inbounds i64, i64* %p, i64 %a3
+  %load3 = load i64, i64* %gep3
+  call void @foo(i64 %load0, i64 %load1, i64 %load2, i64 %load3)
+  ret void
+}