diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index e41f414ef3c0..ab5dc930a325 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 150073aa2abe..8307a962b13a 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 +}