[LV] Account for predicated stores in instruction costs

This patch ensures that we scale the estimated cost of predicated stores by
block probability. This is a follow-on patch for r284123.

llvm-svn: 284126

2 files changed, 45 insertions, 1 deletions

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index e41f414ef3c..ab5dc930a32 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6700,6 +6700,12 @@ unsigned LoopVectorizationCostModel::getInstructionCost(Instruction *I,
       // we might create due to scalarization.
       Cost += getScalarizationOverhead(I, VF, TTI);
 
+      // If we have a predicated store, it may not be executed for each vector
+      // lane. Scale the cost by the probability of executing the predicated
+      // block.
+      if (Legal->isScalarWithPredication(I))
+        Cost /= getReciprocalPredBlockProb();
+
       return Cost;
     }
 
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/predication_costs.ll b/llvm/test/Transforms/LoopVectorize/AArch64/predication_costs.ll
index 150073aa2ab..8307a962b13 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/predication_costs.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/predication_costs.ll
@@ -1,5 +1,5 @@
 ; REQUIRES: asserts
-; RUN: opt < %s -force-vector-width=2 -loop-vectorize -debug-only=loop-vectorize -disable-output 2>&1 | FileCheck %s
+; RUN: opt < %s -force-vector-width=2 -enable-cond-stores-vec -loop-vectorize -debug-only=loop-vectorize -disable-output 2>&1 | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-gnu"
@@ -51,3 +51,41 @@ for.end:
   %tmp7 = phi i32 [ %tmp6, %for.inc ]
   ret i32 %tmp7
 }
+
+; CHECK-LABEL: predicated_store
+;
+; This test checks that we correctly compute the cost of the predicated store
+; instruction. If we assume the block probability is 50%, we compute the cost
+; as:
+;
+; Cost for vector lane zero:
+;   (store(2) + 2 * extractelement(0)) / 2 = 1
+; Cost for vector lane one:
+;   (store(2) + 2 * extractelement(3)) / 2 = 4
+;
+; CHECK: Found an estimated cost of 5 for VF 2 For instruction: store i32 %tmp2, i32* %tmp0, align 4
+; CHECK: Scalarizing and predicating: store i32 %tmp2, i32* %tmp0, align 4
+;
+define void @predicated_store(i32* %a, i1 %c, i32 %x, i64 %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ 0, %entry ], [ %i.next, %for.inc ]
+  %tmp0 = getelementptr inbounds i32, i32* %a, i64 %i
+  %tmp1 = load i32, i32* %tmp0, align 4
+  br i1 %c, label %if.then, label %for.inc
+
+if.then:
+  %tmp2 = add nsw i32 %tmp1, %x
+  store i32 %tmp2, i32* %tmp0, align 4
+  br label %for.inc
+
+for.inc:
+  %i.next = add nuw nsw i64 %i, 1
+  %cond = icmp slt i64 %i.next, %n
+  br i1 %cond, label %for.body, label %for.end
+
+for.end:
+  ret void
+}