[llvm] [SLP]Loop-scale the extract cost in the splat subtree keep/drop check (PR #223524)

[Alexey Bataev via llvm-commits]

https://github.com/alexey-bataev created https://github.com/llvm/llvm-project/pull/223524

The unscaled extract estimate kept extract-heavy subtrees in hot loops,
rejecting profitable trees.


>From c32c3c91f69fccabb9cf5306ba0b4de1ccc2cccf Mon Sep 17 00:00:00 2001
From: Alexey Bataev <a.bataev at outlook.com>
Date: Mon, 14 Sep 2026 14:07:31 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.7
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp |  4 ++++
 .../splat-gather-subtree-loop-extracts.ll       | 17 +++++++----------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 8f34e030ead346..ce7004b33c1556 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -18969,6 +18969,10 @@ BoUpSLP::calculateTreeCostAndTrimNonProfitable(ArrayRef<Value *> VectorizedVals,
         *TTI, SLPReVec, ScalarTy,
         cast<VectorType>(getWidenedType(ScalarTy, TE->getVectorFactor())),
         ExtractElts, /*Insert=*/false, /*Extract=*/true, CostKind);
+    // Scale the extract cost to the subtree's execution frequency: the
+    // subtree and gather costs it is compared against are already loop-scaled.
+    if (KeepCost.isValid() && KeepCost != 0)
+      KeepCost *= getEntryEffectiveScale(*TE);
     // Add the cost of the subtree itself, computed before any trimming:
     // trimming of the subtree's own nodes would otherwise make it look
     // artificially cheap.
diff --git a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-loop-extracts.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-loop-extracts.ll
index edec7e47051118..c2ca451ba0fb3c 100644
--- a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-loop-extracts.ll
+++ b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-loop-extracts.ll
@@ -10,20 +10,17 @@ define i32 @splat_subtree_loop_extracts(double %div.i) {
 ; CHECK-LABEL: @splat_subtree_loop_extracts(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = tail call double @llvm.fmuladd.f64(double 0.000000e+00, double 0.000000e+00, double 0.000000e+00)
-; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <3 x double> <double poison, double 0.000000e+00, double poison>, double [[TMP0]], i64 2
-; CHECK-NEXT:    [[TMP14:%.*]] = insertelement <3 x double> <double poison, double 0.000000e+00, double poison>, double [[DIV_I:%.*]], i64 2
 ; CHECK-NEXT:    br label [[FOR_BODY46_I:%.*]]
 ; CHECK:       for.body46.i:
-; CHECK-NEXT:    [[TMP15:%.*]] = tail call double @llvm.fmuladd.f64(double 0.000000e+00, double 0.000000e+00, double 0.000000e+00)
-; CHECK-NEXT:    [[DIV_I1:%.*]] = fdiv double 0.000000e+00, 0.000000e+00
 ; CHECK-NEXT:    [[ARRAYIDX19_US63_I_3_1:%.*]] = getelementptr i8, ptr poison, i64 328
-; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <3 x double> [[TMP1]], double [[TMP15]], i64 0
-; CHECK-NEXT:    [[TMP5:%.*]] = insertelement <3 x double> [[TMP14]], double [[DIV_I1]], i64 0
-; CHECK-NEXT:    [[TMP16:%.*]] = fmul <3 x double> [[TMP4]], [[TMP5]]
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <3 x double> [[TMP16]], <3 x double> poison, <2 x i32> <i32 2, i32 2>
+; CHECK-NEXT:    [[MUL5_I471:%.*]] = fmul double [[TMP0]], [[DIV_I:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x double> poison, double [[MUL5_I471]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <2 x double> [[TMP1]], <2 x double> poison, <2 x i32> zeroinitializer
 ; CHECK-NEXT:    [[TMP3:%.*]] = call <2 x double> @llvm.fmuladd.v2f64(<2 x double> [[TMP2]], <2 x double> zeroinitializer, <2 x double> zeroinitializer)
-; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <3 x double> [[TMP16]], <3 x double> poison, <2 x i32> zeroinitializer
-; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <3 x double> [[TMP16]], <3 x double> poison, <2 x i32> <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP4:%.*]] = call <2 x double> @llvm.fmuladd.v2f64(<2 x double> zeroinitializer, <2 x double> <double 0.000000e+00, double -0.000000e+00>, <2 x double> zeroinitializer)
+; CHECK-NEXT:    [[TMP5:%.*]] = fmul <2 x double> [[TMP4]], <double +qnan, double 0.000000e+00>
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <2 x double> [[TMP5]], <2 x double> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <2 x double> [[TMP5]], <2 x double> poison, <2 x i32> <i32 1, i32 1>
 ; CHECK-NEXT:    [[TMP8:%.*]] = call <2 x double> @llvm.fmuladd.v2f64(<2 x double> [[TMP6]], <2 x double> zeroinitializer, <2 x double> [[TMP7]])
 ; CHECK-NEXT:    [[TMP9:%.*]] = call <2 x double> @llvm.fmuladd.v2f64(<2 x double> zeroinitializer, <2 x double> zeroinitializer, <2 x double> [[TMP8]])
 ; CHECK-NEXT:    [[TMP10:%.*]] = fmul <2 x double> zeroinitializer, [[TMP3]]