[llvm] [LV] Determine NumPredStores before computing widening costs. (PR #199341)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/199341

Determine NumPredStores before computing widening costs, so we consistently apply large predicated store cost to all stores, matching the VPlan cost model. In practice that should not impact vectorization decisions, as the huge cost for any predicated store other than the first already effectively disables vectorization.

>From a010f5db2877484b2a594f1bf9eae441232f361c Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Fri, 22 May 2026 16:15:33 +0100
Subject: [PATCH] [LV] Determine NumPredStores before computing widening costs.

Determine NumPredStores before computing widening costs, so we
consistently apply large predicated store cost to all stores, matching
the VPlan cost model. In practice that should not impact vectorization
decisions, as the huge cost for any predicated store other than the first
already effectively disables vectorization.
---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 17 ++++++++++-------
 .../CostModel/masked-interleaved-store-i16.ll   |  8 ++++----
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index b65d19054c824..8e6de07864101 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4635,7 +4635,17 @@ LoopVectorizationCostModel::getScalarizationOverhead(Instruction *I,
 void LoopVectorizationCostModel::setCostBasedWideningDecision(ElementCount VF) {
   if (VF.isScalar())
     return;
+
+  // TODO: We should generate better code and update the cost model for
+  // predicated uniform stores. Today they are treated as any other
+  // predicated store (see added test cases in
+  // invariant-store-vectorization.ll).
   NumPredStores = 0;
+  for (BasicBlock *BB : TheLoop->blocks())
+    for (Instruction &I : *BB)
+      if (isa<StoreInst>(&I) && isScalarWithPredication(&I, VF))
+        ++NumPredStores;
+
   for (BasicBlock *BB : TheLoop->blocks()) {
     // For each instruction in the old loop.
     for (Instruction &I : *BB) {
@@ -4643,13 +4653,6 @@ void LoopVectorizationCostModel::setCostBasedWideningDecision(ElementCount VF) {
       if (!Ptr)
         continue;
 
-      // TODO: We should generate better code and update the cost model for
-      // predicated uniform stores. Today they are treated as any other
-      // predicated store (see added test cases in
-      // invariant-store-vectorization.ll).
-      if (isa<StoreInst>(&I) && isScalarWithPredication(&I, VF))
-        NumPredStores++;
-
       if (Legal->isUniformMemOp(I, VF)) {
         auto IsLegalToScalarize = [&]() {
           if (!VF.isScalable())
diff --git a/llvm/test/Transforms/LoopVectorize/X86/CostModel/masked-interleaved-store-i16.ll b/llvm/test/Transforms/LoopVectorize/X86/CostModel/masked-interleaved-store-i16.ll
index 865942265146d..7e0f17a42bc9d 100644
--- a/llvm/test/Transforms/LoopVectorize/X86/CostModel/masked-interleaved-store-i16.ll
+++ b/llvm/test/Transforms/LoopVectorize/X86/CostModel/masked-interleaved-store-i16.ll
@@ -76,13 +76,13 @@ define void @test2(ptr noalias nocapture %points, i32 %numPoints, ptr noalias no
 ; DISABLED_MASKED_STRIDED-LABEL: 'test2'
 ; DISABLED_MASKED_STRIDED:  LV: Found an estimated cost of 1 for VF 1 For instruction: store i16 %0, ptr %arrayidx2, align 2
 ; DISABLED_MASKED_STRIDED:  LV: Found an estimated cost of 1 for VF 1 For instruction: store i16 %2, ptr %arrayidx7, align 2
-; DISABLED_MASKED_STRIDED:  Cost of 8 for VF 2: REPLICATE store ir<%0>, ir<%arrayidx2>
+; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 2: REPLICATE store ir<%0>, ir<%arrayidx2>
 ; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 2: REPLICATE store ir<%2>, ir<%arrayidx7>
-; DISABLED_MASKED_STRIDED:  Cost of 17 for VF 4: REPLICATE store ir<%0>, ir<%arrayidx2>
+; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 4: REPLICATE store ir<%0>, ir<%arrayidx2>
 ; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 4: REPLICATE store ir<%2>, ir<%arrayidx7>
-; DISABLED_MASKED_STRIDED:  Cost of 35 for VF 8: REPLICATE store ir<%0>, ir<%arrayidx2>
+; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 8: REPLICATE store ir<%0>, ir<%arrayidx2>
 ; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 8: REPLICATE store ir<%2>, ir<%arrayidx7>
-; DISABLED_MASKED_STRIDED:  Cost of 71 for VF 16: REPLICATE store ir<%0>, ir<%arrayidx2>
+; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 16: REPLICATE store ir<%0>, ir<%arrayidx2>
 ; DISABLED_MASKED_STRIDED:  Cost of 3000000 for VF 16: REPLICATE store ir<%2>, ir<%arrayidx7>
 ;
 ; ENABLED_MASKED_STRIDED-LABEL: 'test2'