Rename getReciprocalPredBlockProb and move CostKind handling into it.

Author: john-brawn-arm

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -3345,7 +3345,7 @@ LoopVectorizationCostModel::getDivRemSpeculationCost(Instruction *I,
     // Scale the cost by the probability of executing the predicated blocks.
     // This assumes the predicated block for each vector lane is equally
     // likely.
-    ScalarizationCost = ScalarizationCost / getReciprocalPredBlockProb();
+    ScalarizationCost = ScalarizationCost / getPredBlockCostDivisor(CostKind);
   }
   InstructionCost SafeDivisorCost = 0;
 
@@ -5515,8 +5515,7 @@ InstructionCost LoopVectorizationCostModel::computePredInstDiscount(
       }
 
     // Scale the total scalar cost by block probability.
-    if (CostKind != TTI::TCK_CodeSize)
-      ScalarCost /= getReciprocalPredBlockProb();
+    ScalarCost /= getPredBlockCostDivisor(CostKind);
 
     // Compute the discount. A non-negative discount means the vector version
     // of the instruction costs more, and scalarizing would be beneficial.
@@ -5568,9 +5567,8 @@ InstructionCost LoopVectorizationCostModel::expectedCost(ElementCount VF) {
     // the predicated block, if it is an if-else block. Thus, scale the block's
     // cost by the probability of executing it. blockNeedsPredication from
     // Legal is used so as to not include all blocks in tail folded loops.
-    if (VF.isScalar() && Legal->blockNeedsPredication(BB) &&
-        CostKind != TTI::TCK_CodeSize)
-      BlockCost /= getReciprocalPredBlockProb();
+    if (VF.isScalar() && Legal->blockNeedsPredication(BB))
+      BlockCost /= getPredBlockCostDivisor(CostKind);
 
     Cost += BlockCost;
   }
@@ -5648,8 +5646,7 @@ LoopVectorizationCostModel::getMemInstScalarizationCost(Instruction *I,
   // conditional branches, but may not be executed for each vector lane. Scale
   // the cost by the probability of executing the predicated block.
   if (isPredicatedInst(I)) {
-    if (CostKind != TTI::TCK_CodeSize)
-      Cost /= getReciprocalPredBlockProb();
+    Cost /= getPredBlockCostDivisor(CostKind);
 
     // Add the cost of an i1 extract and a branch
     auto *VecI1Ty =

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -793,8 +793,8 @@ InstructionCost VPRegionBlock::cost(ElementCount VF, VPCostContext &Ctx) {
 
   // For the scalar case, we may not always execute the original predicated
   // block, Thus, scale the block's cost by the probability of executing it.
-  if (VF.isScalar() && Ctx.CostKind != TTI::TCK_CodeSize)
-    return ThenCost / getReciprocalPredBlockProb();
+  if (VF.isScalar())
+    return ThenCost / getPredBlockCostDivisor(Ctx.CostKind);
 
   return ThenCost;
 }

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -83,13 +83,20 @@ Value *getRuntimeVF(IRBuilderBase &B, Type *Ty, ElementCount VF);
 Value *createStepForVF(IRBuilderBase &B, Type *Ty, ElementCount VF,
                        int64_t Step);
 
-/// A helper function that returns the reciprocal of the block probability of
-/// predicated blocks. If we return X, we are assuming the predicated block
-/// will execute once for every X iterations of the loop header.
+/// A helper function that returns how much we should divide the cost of a
+/// predicated block by. Typically this is the reciprocal of the block
+/// probability, i.e. if we return X we are assuming the predicated block will
+/// execute once for every X iterations of the loop header so the block should
+/// only contribute 1/X of its cost to the total cost calculation, but when
+/// optimizing for code size it will just be 1 as code size costs don't depend
+/// on execution probabilities.
 ///
 /// TODO: We should use actual block probability here, if available. Currently,
 ///       we always assume predicated blocks have a 50% chance of executing.
-inline unsigned getReciprocalPredBlockProb() { return 2; }
+inline unsigned
+getPredBlockCostDivisor(TargetTransformInfo::TargetCostKind CostKind) {
+  return CostKind == TTI::TCK_CodeSize ? 1 : 2;
+}
 
 /// A range of powers-of-2 vectorization factors with fixed start and
 /// adjustable end. The range includes start and excludes end, e.g.,: