[Loop Vectorize] Block Probability for Predicated Blocks

[DIVYA]

Summary

LoopVectorize Pass uses reciprocal of block probability as 2 for all predicated blocks.Here, we have used BlockFrequencyInfo to get the block probabilities

Worked in collaboration with Aditya Kumar

[fhahn]

I'm just curious, did you run any benchmarks with this change? I think it's quite clear that using the block frequencies is a good idea, but it would be awesome to know if/by how much it improves things :)

Also, is it worth adding a dedicated test case making sure the frequencies are used by the cost model as expected?

lib/Transforms/Vectorize/LoopVectorize.cpp

343

It looks like this commit fixes this TODO, remove the TODO?

346

Any reasons for the name change?

Also, given that this function takes BFI and the headerBB as arguments I think making it a member function of LoopVectorizationCostModel would make it easier to use (both BFI and HeaderBB should be available as member variables there)

349

At least the outermost brackets can be omitted here I think

[inouehrs]

lib/Transforms/Vectorize/LoopVectorize.cpp

350

Don't you need to handle the case of zero frequency (e.g. if BFI is not available) to avoid divide-by-zero?

[DIVYA]

• Added getReciprocalPredBlockProb() as a member function of LoopVectorizationCostModel class
• Added a testcase

[fhahn]

lib/Transforms/Vectorize/LoopVectorize.cpp

2016

I think this could be const? and BB could also be const?

Is BFI guaranteed to be non-NULL?

[DIVYA]

[DIVYA]

ping

[DIVYA]

ping

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

2020–2021

I too am curious if you saw any performance improvements with this change. Can you share any data?

If the goal is to be more precise, it probably makes sense for this function to return a float for block probability and not an unsigned for the reciprocal. I think nearly all the users of this function use it as the denominator of a division. So with the added division here and the ones that follow, we're losing precision. I'm not sure this is any better than always returning 2.

test/Transforms/LoopVectorize/AArch64/aarch64-predication.ll

25

FWIW, the current patch breaks this test. For it to be testing what it's intended to test, the add should be scalarized.

[DIVYA]

This comment was removed by DIVYA.

lib/Transforms/Vectorize/LoopVectorize.cpp

static bool hasIrregularType(Type *Ty, const DataLayout &DL, unsigned VF) {
   // requires padding between elements.
   return DL.getTypeAllocSizeInBits(Ty) != DL.getTypeSizeInBits(Ty);
 }
 
 /// 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 for every X iterations of the loop header.
 ///
 /// TODO: We should use actual block probability here, if available. Currently,
 ///       we always assume predicated blocks have a 50% chance of executing.
-static unsigned getReciprocalPredBlockProb() { return 2; }
+
+static unsigned getReciprocalPredBlockProb1(BlockFrequencyInfo *BFI,
+                                            BasicBlock *BB,
+                                            BasicBlock *HeaderBB) {
+  return ((BFI->getBlockFreq(HeaderBB)).getFrequency()) /
+         ((BFI->getBlockFreq(BB)).getFrequency());
+}
 
 /// A helper function that adds a 'fast' flag to floating-point operations.
 static Value *addFastMathFlag(Value *V) {
   if (isa<FPMathOperator>(V)) {
     FastMathFlags Flags;
     Flags.setUnsafeAlgebra();
     cast<Instruction>(V)->setFastMathFlags(Flags);
   }

 class LoopVectorizationCostModel {
 public:
   LoopVectorizationCostModel(Loop *L, PredicatedScalarEvolution &PSE,
                              LoopInfo *LI, LoopVectorizationLegality *Legal,
                              const TargetTransformInfo &TTI,
                              const TargetLibraryInfo *TLI, DemandedBits *DB,
                              AssumptionCache *AC,
                              OptimizationRemarkEmitter *ORE, const Function *F,
-                             const LoopVectorizeHints *Hints)
+                             const LoopVectorizeHints *Hints,
+                             BlockFrequencyInfo *BFI )
       : TheLoop(L), PSE(PSE), LI(LI), Legal(Legal), TTI(TTI), TLI(TLI), DB(DB),
-        AC(AC), ORE(ORE), TheFunction(F), Hints(Hints) {}
+        AC(AC), ORE(ORE), TheFunction(F), Hints(Hints), BFI(BFI) {}
 
   /// \return An upper bound for the vectorization factor, or None if
   /// vectorization should be avoided up front.
   Optional<unsigned> computeMaxVF(bool OptForSize);
 
   /// Information about vectorization costs
   struct VectorizationFactor {
     unsigned Width; // Vector width with best cost
bool isOptimizableIVTruncate(Instruction *I, unsigned VF) {
     // If the truncated value is not an induction variable, return false.
     return Legal->isInductionVariable(Op);
   }
 
 private:
   /// \return An upper bound for the vectorization factor, larger than zero.
   /// One is returned if vectorization should best be avoided due to cost.
   unsigned computeFeasibleMaxVF(bool OptForSize);
 
   /// The vectorization cost is a combination of the cost itself and a boolean
   /// indicating whether any of the contributing operations will actually
   /// operate on
   /// vector values after type legalization in the backend. If this latter value
   /// is
   /// false, then all operations will be scalarized (i.e. no vectorization has
   /// actually taken place).
   typedef std::pair<unsigned, bool> VectorizationCostTy;
 
   /// Returns the expected execution cost. The unit of the cost does
   /// not matter because we use the 'cost' units to compare different
   /// vector widths. The cost that is returned is *not* normalized by
   /// the factor width.
public:
   /// Assumption cache.
   AssumptionCache *AC;
   /// Interface to emit optimization remarks.
   OptimizationRemarkEmitter *ORE;
 
   const Function *TheFunction;
   /// Loop Vectorize Hint.
   const LoopVectorizeHints *Hints;
+  /// Block Frequency Info
+  BlockFrequencyInfo *BFI;
   /// Values to ignore in the cost model.
   SmallPtrSet<const Value *, 16> ValuesToIgnore;
   /// Values to ignore in the cost model when VF > 1.
   SmallPtrSet<const Value *, 16> VecValuesToIgnore;
 };
 
 /// LoopVectorizationPlanner - drives the vectorization process after having
 /// passed Legality checks.
for (Use &U : I->operands())
                "Instruction has non-scalar type");
         if (canBeScalarized(J))
           Worklist.push_back(J);
         else if (needsExtract(J))
           ScalarCost += TTI.getScalarizationOverhead(
                               ToVectorTy(J->getType(),VF), false, true);
       }
 
+
     // Scale the total scalar cost by block probability.
-    ScalarCost /= getReciprocalPredBlockProb();
+    ScalarCost /=
+        getReciprocalPredBlockProb1(BFI,I->getParent(),TheLoop->getHeader());
 
     // Compute the discount. A non-negative discount means the vector version
     // of the instruction costs more, and scalarizing would be beneficial.
     Discount += VectorCost - ScalarCost;
     ScalarCosts[I] = ScalarCost;
   }
 
   return Discount;
for (BasicBlock *BB : TheLoop->blocks()) {
 
     // If we are vectorizing a predicated block, it will have been
     // if-converted. This means that the block's instructions (aside from
     // stores and instructions that may divide by zero) will now be
     // unconditionally executed. For the scalar case, we may not always execute
     // the predicated block. Thus, scale the block's cost by the probability of
     // executing it.
     if (VF == 1 && Legal->blockNeedsPredication(BB))
-      BlockCost.first /= getReciprocalPredBlockProb();
+      BlockCost.first /=
+          getReciprocalPredBlockProb1(BFI, BB, TheLoop->getHeader());
 
     Cost.first += BlockCost.first;
     Cost.second |= BlockCost.second;
   }
 
   return Cost;
 }
 
unsigned LoopVectorizationCostModel::getMemInstScalarizationCost(Instruction *I,
   // Get the overhead of the extractelement and insertelement instructions
   // 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();
+    Cost /=
+        getReciprocalPredBlockProb1(BFI, I->getParent(), TheLoop->getHeader());
 
   return Cost;
 }
 
 unsigned LoopVectorizationCostModel::getConsecutiveMemOpCost(Instruction *I,
                                                              unsigned VF) {
   Type *ValTy = getMemInstValueType(I);
   Type *VectorTy = ToVectorTy(ValTy, VF);
if (VF > 1 && Legal->isScalarWithPredication(I)) {
 
       // The cost of insertelement and extractelement instructions needed for
       // scalarization.
       Cost += getScalarizationOverhead(I, VF, TTI);
 
       // Scale the cost by the probability of executing the predicated blocks.
       // This assumes the predicated block for each vector lane is equally
       // likely.
-      return Cost / getReciprocalPredBlockProb();
+      return Cost / getReciprocalPredBlockProb1(BFI, I->getParent(),
+                                                TheLoop->getHeader());
     }
     LLVM_FALLTHROUGH;
   case Instruction::Add:
   case Instruction::FAdd:
   case Instruction::Sub:
   case Instruction::FSub:
   case Instruction::Mul:
   case Instruction::FMul:
ORE->emit(
         createMissedAnalysis(Hints.vectorizeAnalysisPassName(), "UnsafeFP", L)
         << "loop not vectorized due to unsafe FP support.");
     emitMissedWarning(F, L, Hints, ORE);
     return false;
   }
 
   // Use the cost model.
   LoopVectorizationCostModel CM(L, PSE, LI, &LVL, *TTI, TLI, DB, AC, ORE, F,
-                                &Hints);
+                                &Hints, BFI);
   CM.collectValuesToIgnore();
 
   // Use the planner for vectorization.
   LoopVectorizationPlanner LVP(L, LI, &LVL, CM);
 
   // Get user vectorization factor.
   unsigned UserVF = Hints.getWidth();
 

test/Transforms/LoopVectorize/AArch64/aarch64-predication.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -loop-vectorize -disable-output -debug-only=loop-vectorize 2>&1 | FileCheck %s --check-prefix=COST
 ; RUN: opt < %s -loop-vectorize -force-vector-width=2 -instcombine -simplifycfg -S | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-gnu"
 
 ; This test checks that we correctly compute the scalarized operands for a
 ; user-specified vectorization factor when interleaving is disabled. We use the
 ; "optsize" attribute to disable all interleaving calculations.  A cost of 4
 ; for %tmp4 indicates that we would scalarize it's operand (%tmp3), giving
 ; %tmp4 a lower scalarization overhead.
 ;
 ; COST-LABEL:  predicated_udiv_scalarized_operand
-; COST:        LV: Found an estimated cost of 4 for VF 2 For instruction: %tmp4 = udiv i64 %tmp2, %tmp3
+; COST:        LV: Found an estimated cost of 11 for VF 2 For instruction: %tmp4 = udiv i64 %tmp2, %tmp3
 ;
 ; CHECK-LABEL: @predicated_udiv_scalarized_operand(
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %entry ], [ [[INDEX_NEXT:%.*]], %[[PRED_UDIV_CONTINUE2:.*]] ]
 ; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <2 x i64> [ zeroinitializer, %entry ], [ [[TMP17:%.*]], %[[PRED_UDIV_CONTINUE2]] ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i64, i64* %a, i64 [[INDEX]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i64* [[TMP0]] to <2 x i64>*
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <2 x i64>, <2 x i64>* [[TMP1]], align 4
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt <2 x i64> [[WIDE_LOAD]], zeroinitializer
+; CHECK-NEXT:    [[TMP6:%.*]] = add nsw <2 x i64> [[WIDE_LOAD]], %broadcast.splat2
 ; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <2 x i1> [[TMP2]], i32 0
 ; CHECK-NEXT:    br i1 [[TMP3]], label %[[PRED_UDIV_IF:.*]], label %[[PRED_UDIV_CONTINUE:.*]]
 ; CHECK:       [[PRED_UDIV_IF]]:
 ; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <2 x i64> [[WIDE_LOAD]], i32 0
-; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x i64> [[WIDE_LOAD]], i32 0
-; CHECK-NEXT:    [[TMP6:%.*]] = add nsw i64 [[TMP5]], %x
-; CHECK-NEXT:    [[TMP7:%.*]] = udiv i64 [[TMP4]], [[TMP6]]
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x i64> [[TMP6]], i32 0
+; CHECK-NEXT:    [[TMP7:%.*]] = udiv i64 [[TMP4]], [[TMP5]]
 ; CHECK-NEXT:    [[TMP8:%.*]] = insertelement <2 x i64> undef, i64 [[TMP7]], i32 0
 ; CHECK-NEXT:    br label %[[PRED_UDIV_CONTINUE]]
 ; CHECK:       [[PRED_UDIV_CONTINUE]]:
 ; CHECK-NEXT:    [[TMP9:%.*]] = phi <2 x i64> [ undef, %vector.body ], [ [[TMP8]], %[[PRED_UDIV_IF]] ]
 ; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <2 x i1> [[TMP2]], i32 1
 ; CHECK-NEXT:    br i1 [[TMP10]], label %[[PRED_UDIV_IF1:.*]], label %[[PRED_UDIV_CONTINUE2]]
 ; CHECK:       [[PRED_UDIV_IF1]]:
 ; CHECK-NEXT:    [[TMP11:%.*]] = extractelement <2 x i64> [[WIDE_LOAD]], i32 1
-; CHECK-NEXT:    [[TMP12:%.*]] = extractelement <2 x i64> [[WIDE_LOAD]], i32 1
-; CHECK-NEXT:    [[TMP13:%.*]] = add nsw i64 [[TMP12]], %x
-; CHECK-NEXT:    [[TMP14:%.*]] = udiv i64 [[TMP11]], [[TMP13]]
+; CHECK-NEXT:    [[TMP12:%.*]] = extractelement <2 x i64> [[TMP6]], i32 1
+; CHECK-NEXT:    [[TMP14:%.*]] = udiv i64 [[TMP11]], [[TMP12]]
 ; CHECK-NEXT:    [[TMP15:%.*]] = insertelement <2 x i64> [[TMP9]], i64 [[TMP14]], i32 1
 ; CHECK-NEXT:    br label %[[PRED_UDIV_CONTINUE2]]
 ; CHECK:       [[PRED_UDIV_CONTINUE2]]:
 ; CHECK-NEXT:    [[TMP16:%.*]] = phi <2 x i64> [ [[TMP9]], %[[PRED_UDIV_CONTINUE]] ], [ [[TMP15]], %[[PRED_UDIV_IF1]] ]
 ; CHECK-NEXT:    [[PREDPHI:%.*]] = select <2 x i1> [[TMP2]], <2 x i64> [[TMP16]], <2 x i64> [[WIDE_LOAD]]
 ; CHECK-NEXT:    [[TMP17]] = add <2 x i64> [[VEC_PHI]], [[PREDPHI]]
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add i64 [[INDEX]], 2
 ; CHECK:         br i1 {{.*}}, label %middle.block, label %vector.body