[LoopPredication] Add profitability check based on BPI

[anna]

Summary

LoopPredication is not profitable when the loop is known to always exit
through some block other than the latch block.
A coarse grained latch check can cause loop predication to predicate the
loop, making the guard invariant out of the loop and always
unconditionally deoptimize. However, the guard may never fail within the
loop because dynamically, it always exits earlier through some other block and never through
the latch block.
We teach LP about this using BranchProfileInfo pass.

This methodology was chosen instead of adding metadata to prevent loop
predication: https://reviews.llvm.org/D44535

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

800

I'm not against lambda, but don't you think that this big code would be better to put in the separate method of LoopPredication?

825

Don't you want to introduce some delta? Usually branch probability is something is better to compare with some delta.

[anna]

lib/Transforms/Scalar/LoopPredication.cpp

800

agreed. Will do.

825

Agreed. The delta would need to be cl::opt argument as well.

[anna]

seperated out logic into a function, instead of lambda. Added a threshold for the BPI check.

[anna]

fixed a bug in internal testing with Loop predication pass initialization - we need to register
the BPI pass dependency.

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

212

I think it is better to add some comment here, what the value of this option actually means.

726

BasicBlock has an utility function getExitEdges.
I propose to use it instead of iteration over ExitingLoops.
This way, your loop will be simplified a lot I guess.

I would also consider inserting an assert that latch is an exiting block.
By definition latch should not require to be an exiting block. But the logic below expects it as I understand...
As I understand the latch is an exiting block due to it is a loop pass and loop is in canonical form. So it is ok but I would still consider inserting an assert... Or alternatively I would find an edge from latch to exit block in an array returned from getExitEdges (this will imply an assert and gives you an input to get probability).

750

It is not a big deal and compiler should do it instead of you but you can hoist LatchExitProbability * ProbabilityThreshold to pre-header :)

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

750

Also please consider renaming of the option ProbabilityThreshold. To me threshold is actually "LatchExitProbability * ProbabilityThreshold". However I might be wrong.

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

726

I'm sorry, not a BasicBlock but Loop of course.

[anna]

lib/Transforms/Scalar/LoopPredication.cpp

726

BasicBlock has an utility function getExitEdges.

Yes, that's useful (the one in LoopInfo). Thanks! I haven't seen that function being used anywhere in the code, so it may not be well tested. However, the functionality is exactly what I want and it simplifies the logic below. I'll still extract the LatchExitProbability the way I've done right now, without going through the array.

I would also consider inserting an assert that latch is an exiting block.

Yup, Latch is an exiting block at this point because Loop Predication requires that form. I'll add an assert here.

750

Maybe something like loop-pred-latch-probability-scale?

[anna]

use getExitEdges instead of getting exiting blocks.

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

727

you do not need ExitingBlocks more.

737

Please consider re-ordering the lines in this piece of code (it is a bit difficult to read it).

I would suggest to move all code related to getting edge probability below to computation of LatchProbabilityThreshold (to me it is connected parts). Additionally you will have an earlier bail out due to more than one exit and you do not need to do anything for latch in this case.

And separate by new line the code related to LatchProbabilityThreshold and getting ExitEdges.

749

You still need to skip latch exiting basic block.

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

748

Just a side comment... Instead of "for loop" you can consider "any_of" usage in this case. However I do not know whether it looks better :) So it is up to you.

[anna]

lib/Transforms/Scalar/LoopPredication.cpp

727

stale line will remove.

737

ok.

748

since there's an extra computation of probability here, I'm not sure if any_of is better :)

749

you don't need to skip it. The check at line 750 will never be true for latch exit basic block (exitingblock probability > latchExitingBlockProbability). I intentionally ignored checking for latch exiting basic block, it saves us a check for every exiting edge...

[skatkov]

lib/Transforms/Scalar/LoopPredication.cpp

749

That is true until someone will not set LatchExitProbabilityScale to 0.5.

[anna]

lib/Transforms/Scalar/LoopPredication.cpp

749

hmm.. that changes the whole point of the profitability check. It inverts it. Thanks for pointing that out. We need to strengthen against such use cases...
if LatchExitProbabilityScale < 1, we set it to 1.

[anna]

Address Serguei's comments.

[skatkov]

LGTM with some comments.

lib/Transforms/Scalar/LoopPredication.cpp

215

cl::opt has a way to add a description for the option: cl::desc("")
I would suggest always to provide a description which is printed in help message.
The notice about >=1 and ignoring the value <1 can be also added to desc.

726

new line?

750

DEBUG(dbgs() with message that we ignore user settings?
It is not a good idea to update the value of opt. I would suggest to use local variable for that.

lib/Transforms/Scalar/LoopPredication.cpp

 //     guardStart u< guardLimit && latchLimit s>= 1.
 //   For uge condition the widened condition is:
 //     guardStart u< guardLimit && latchLimit u> 1.
 //   For sge condition the widened condition is:
 //     guardStart u< guardLimit && latchLimit s> 1.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LoopPredication.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 
 #define DEBUG_TYPE "loop-predication"
 
 using namespace llvm;
 
 static cl::opt<bool> EnableIVTruncation("loop-predication-enable-iv-truncation",
                                         cl::Hidden, cl::init(true));
 
 static cl::opt<bool> EnableCountDownLoop("loop-predication-enable-count-down-loop",
                                         cl::Hidden, cl::init(true));
+
+static cl::opt<bool>
+    SkipProfitabilityChecks("loop-predication-skip-profitability-checks",
+                            cl::Hidden, cl::init(false));
+
+// This is the scale factor for the latch probability. We use this during
+// profitability analysis to find other exiting blocks that have a much higher
+// probability of exiting the loop instead of loop exiting via latch.
+// This value should be greater than 1 for a sane profitability check.
+static cl::opt<float>
+    LatchExitProbabilityScale("loop-predication-latch-probability-scale",
+                         cl::Hidden, cl::init(2.0));
+
 namespace {
 class LoopPredication {
   /// Represents an induction variable check:
   ///   icmp Pred, <induction variable>, <loop invariant limit>
   struct LoopICmp {
     ICmpInst::Predicate Pred;
     const SCEVAddRecExpr *IV;
     const SCEV *Limit;
     LoopICmp(ICmpInst::Predicate Pred, const SCEVAddRecExpr *IV,
              const SCEV *Limit)
         : Pred(Pred), IV(IV), Limit(Limit) {}
     LoopICmp() {}
     void dump() {
       dbgs() << "LoopICmp Pred = " << Pred << ", IV = " << *IV
              << ", Limit = " << *Limit << "\n";
     }
   };
 
   ScalarEvolution *SE;
+  BranchProbabilityInfo *BPI;
 
   Loop *L;
   const DataLayout *DL;
   BasicBlock *Preheader;
   LoopICmp LatchCheck;
 
   bool isSupportedStep(const SCEV* Step);
   Optional<LoopICmp> parseLoopICmp(ICmpInst *ICI) {
Optional<Value *> widenICmpRangeCheckIncrementingLoop(LoopICmp LatchCheck,
                                                         SCEVExpander &Expander,
                                                         IRBuilder<> &Builder);
   Optional<Value *> widenICmpRangeCheckDecrementingLoop(LoopICmp LatchCheck,
                                                         LoopICmp RangeCheck,
                                                         SCEVExpander &Expander,
                                                         IRBuilder<> &Builder);
   bool widenGuardConditions(IntrinsicInst *II, SCEVExpander &Expander);
 
+  // If the loop always exits through another block in the loop, we should not
+  // predicate based on the latch check. For example, the latch check can be a
+  // very coarse grained check and there can be more fine grained exit checks
+  // within the loop. We identify such unprofitable loops through BPI.
+  bool isLoopProfitableToPredicate();
+
   // When the IV type is wider than the range operand type, we can still do loop
   // predication, by generating SCEVs for the range and latch that are of the
   // same type. We achieve this by generating a SCEV truncate expression for the
   // latch IV. This is done iff truncation of the IV is a safe operation,
   // without loss of information.
   // Another way to achieve this is by generating a wider type SCEV for the
   // range check operand, however, this needs a more involved check that
   // operands do not overflow. This can lead to loss of information when the
   // range operand is of the form: add i32 %offset, %iv. We need to prove that
   // sext(x + y) is same as sext(x) + sext(y).
   // This function returns true if we can safely represent the IV type in
   // the RangeCheckType without loss of information.
   bool isSafeToTruncateWideIVType(Type *RangeCheckType);
   // Return the loopLatchCheck corresponding to the RangeCheckType if safe to do
   // so.
   Optional<LoopICmp> generateLoopLatchCheck(Type *RangeCheckType);
 
 public:
-  LoopPredication(ScalarEvolution *SE) : SE(SE){};
+  LoopPredication(ScalarEvolution *SE, BranchProbabilityInfo *BPI)
+      : SE(SE), BPI(BPI){};
   bool runOnLoop(Loop *L);
 };
 
 class LoopPredicationLegacyPass : public LoopPass {
 public:
   static char ID;
   LoopPredicationLegacyPass() : LoopPass(ID) {
     initializeLoopPredicationLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<BranchProbabilityInfoWrapperPass>();
     getLoopAnalysisUsage(AU);
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override {
     if (skipLoop(L))
       return false;
     auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
-    LoopPredication LP(SE);
+    BranchProbabilityInfo &BPI =
+        getAnalysis<BranchProbabilityInfoWrapperPass>().getBPI();
+    LoopPredication LP(SE, &BPI);
     return LP.runOnLoop(L);
   }
 };
 
 char LoopPredicationLegacyPass::ID = 0;
 } // end namespace llvm
 
 INITIALIZE_PASS_BEGIN(LoopPredicationLegacyPass, "loop-predication",
                       "Loop predication", false, false)
+INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
 INITIALIZE_PASS_END(LoopPredicationLegacyPass, "loop-predication",
                     "Loop predication", false, false)
 
 Pass *llvm::createLoopPredicationPass() {
   return new LoopPredicationLegacyPass();
 }
 
 PreservedAnalyses LoopPredicationPass::run(Loop &L, LoopAnalysisManager &AM,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
-  LoopPredication LP(&AR.SE);
+  const auto &FAM =
+      AM.getResult<FunctionAnalysisManagerLoopProxy>(L, AR).getManager();
+  Function *F = L.getHeader()->getParent();
+  auto *BPI = FAM.getCachedResult<BranchProbabilityAnalysis>(*F);
+  LoopPredication LP(&AR.SE, BPI);
   if (!LP.runOnLoop(&L))
     return PreservedAnalyses::all();
 
   return getLoopPassPreservedAnalyses();
 }
 
 Optional<LoopPredication::LoopICmp>
 LoopPredication::parseLoopICmp(ICmpInst::Predicate Pred, Value *LHS,
bool LoopPredication::isSafeToTruncateWideIVType(Type *RangeCheckType) {
   // The active bits should be less than the bits in the RangeCheckType. This
   // guarantees that truncating the latch check to RangeCheckType is a safe
   // operation.
   auto RangeCheckTypeBitSize = DL->getTypeSizeInBits(RangeCheckType);
   return Start->getAPInt().getActiveBits() < RangeCheckTypeBitSize &&
          Limit->getAPInt().getActiveBits() < RangeCheckTypeBitSize;
 }
 
+bool LoopPredication::isLoopProfitableToPredicate() {
+  if (SkipProfitabilityChecks || !BPI)
+    return true;
+  SmallVector<std::pair<const BasicBlock *, const BasicBlock *>, 8> ExitEdges;
+  L->getExitEdges(ExitEdges);
+  // If there is only one exiting edge in the loop, it is always profitable to
+  // predicate the loop.
+  if (ExitEdges.size() == 1)
+    return true;
+
+  // Calculate the exiting probabilities of all exiting edges from the loop,
+  // starting with the LatchExitProbability.
+  // Heuristic for profitability: If any of the exiting blocks' probability of
+  // exiting the loop is larger than exiting through the latch block, it's not
+  // profitable to predicate the loop.
+  auto *LatchBlock = L->getLoopLatch();
+  assert(LatchBlock && "Should have a single latch at this point!");
+  auto *LatchTerm = LatchBlock->getTerminator();
+  assert(LatchTerm->getNumSuccessors() == 2 &&
+         "expected to be an exiting block with 2 succs!");
+  unsigned LatchBrExitIdx =
+      LatchTerm->getSuccessor(0) == L->getHeader() ? 1 : 0;
+  BranchProbability LatchExitProbability =
+      BPI->getEdgeProbability(LatchBlock, LatchBrExitIdx);
+  // Protect against degenerate inputs provided by the user. Providing a value
+  // less than one, can invert the definition of profitable loop predication.
+  if (LatchExitProbabilityScale < 1)
+    LatchExitProbabilityScale = 1.0;
+  const auto LatchProbabilityThreshold =
+      LatchExitProbability * LatchExitProbabilityScale;
+
+  for (const auto &ExitEdge : ExitEdges) {
+    BranchProbability ExitingBlockProbability =
+        BPI->getEdgeProbability(ExitEdge.first, ExitEdge.second);
+    // Some exiting edge has higher probability than the latch exiting edge.
+    // No longer profitable to predicate.
+    if (ExitingBlockProbability > LatchProbabilityThreshold)
+      return false;
+  }
+  // Using BPI, we have concluded that the most probable way to exit from the
+  // loop is through the latch (or there's no profile information and all
+  // exits are equally likely).
+  return true;
+}
+
 bool LoopPredication::runOnLoop(Loop *Loop) {
   L = Loop;
 
   DEBUG(dbgs() << "Analyzing ");
   DEBUG(L->dump());
 
   Module *M = L->getHeader()->getModule();
 
bool LoopPredication::runOnLoop(Loop *Loop) {
   auto LatchCheckOpt = parseLoopLatchICmp();
   if (!LatchCheckOpt)
     return false;
   LatchCheck = *LatchCheckOpt;
 
   DEBUG(dbgs() << "Latch check:\n");
   DEBUG(LatchCheck.dump());
 
+  if (!isLoopProfitableToPredicate()) {
+    DEBUG(dbgs()<< "Loop not profitable to predicate!\n");
+    return false;
+  }
   // Collect all the guards into a vector and process later, so as not
   // to invalidate the instruction iterator.
   SmallVector<IntrinsicInst *, 4> Guards;
   for (const auto BB : L->blocks())
     for (auto &I : *BB)
       if (auto *II = dyn_cast<IntrinsicInst>(&I))
         if (II->getIntrinsicID() == Intrinsic::experimental_guard)
           Guards.push_back(II);
 
   if (Guards.empty())
     return false;
 
   SCEVExpander Expander(*SE, *DL, "loop-predication");
 
   bool Changed = false;
   for (auto *Guard : Guards)
     Changed |= widenGuardConditions(Guard, Expander);
 
   return Changed;
 }

test/Transforms/LoopPredication/profitability.ll

+; RUN: opt -S -loop-predication -loop-predication-skip-profitability-checks=false < %s 2>&1 | FileCheck %s
+; RUN: opt -S -loop-predication-skip-profitability-checks=false -passes='require<scalar-evolution>,require<branch-prob>,loop(loop-predication)' < %s 2>&1 | FileCheck %s
+
+; latch block exits to a speculation block. BPI already knows (without prof
+; data) that deopt is very rarely
+; taken. So we do not predicate this loop using that coarse latch check.
+; LatchExitProbability: 0x04000000 / 0x80000000 = 3.12%
+; ExitingBlockProbability: 0x7ffa572a / 0x80000000 = 99.98%
+define i64 @donot_predicate(i64* nocapture readonly %arg, i32 %length, i64* nocapture readonly %arg2, i64* nocapture readonly %n_addr, i64 %i) {
+; CHECK-LABEL: donot_predicate(
+entry:
+  %length.ext = zext i32 %length to i64
+  %n.pre = load i64, i64* %n_addr, align 4
+  br label %Header
+
+; CHECK-LABEL: Header:
+; CHECK:         %within.bounds = icmp ult i64 %j2, %length.ext
+; CHECK-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9)
+Header:                                          ; preds = %entry, %Latch
+  %result.in3 = phi i64* [ %arg2, %entry ], [ %arg, %Latch ]
+  %j2 = phi i64 [ 0, %entry ], [ %j.next, %Latch ]
+  %within.bounds = icmp ult i64 %j2, %length.ext
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+  %innercmp = icmp eq i64 %j2, %n.pre
+  %j.next = add nuw nsw i64 %j2, 1
+  br i1 %innercmp, label %Latch, label %exit, !prof !0
+
+Latch:                                           ; preds = %Header
+  %speculate_trip_count = icmp ult i64 %j.next, 1048576
+  br i1 %speculate_trip_count, label %Header, label %deopt
+
+deopt:                                            ; preds = %Latch
+  %counted_speculation_failed = call i64 (...) @llvm.experimental.deoptimize.i64(i64 30) [ "deopt"(i32 0) ]
+  ret i64 %counted_speculation_failed
+
+exit:                                             ; preds = %Header
+  %result.in3.lcssa = phi i64* [ %result.in3, %Header ]
+  %result.le = load i64, i64* %result.in3.lcssa, align 8
+  ret i64 %result.le
+}
+!0 = !{!"branch_weights", i32 18, i32 104200}
+
+; predicate loop since there's no profile information and BPI concluded all
+; exiting blocks have same probability of exiting from loop.
+define i64 @predicate(i64* nocapture readonly %arg, i32 %length, i64* nocapture readonly %arg2, i64* nocapture readonly %n_addr, i64 %i) {
+; CHECK-LABEL: predicate(
+; CHECK-LABEL: entry:
+; CHECK:           [[limit_check:[^ ]+]] = icmp ule i64 1048576, %length.ext
+; CHECK-NEXT:      [[first_iteration_check:[^ ]+]] = icmp ult i64 0, %length.ext
+; CHECK-NEXT: [[wide_cond:[^ ]+]] = and i1 [[first_iteration_check]], [[limit_check]]
+entry:
+  %length.ext = zext i32 %length to i64
+  %n.pre = load i64, i64* %n_addr, align 4
+  br label %Header
+
+; CHECK-LABEL: Header:
+; CHECK: call void (i1, ...) @llvm.experimental.guard(i1 [[wide_cond]], i32 9) [ "deopt"() ]
+Header:                                          ; preds = %entry, %Latch
+  %result.in3 = phi i64* [ %arg2, %entry ], [ %arg, %Latch ]
+  %j2 = phi i64 [ 0, %entry ], [ %j.next, %Latch ]
+  %within.bounds = icmp ult i64 %j2, %length.ext
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+  %innercmp = icmp eq i64 %j2, %n.pre
+  %j.next = add nuw nsw i64 %j2, 1
+  br i1 %innercmp, label %Latch, label %exit
+
+Latch:                                           ; preds = %Header
+  %speculate_trip_count = icmp ult i64 %j.next, 1048576
+  br i1 %speculate_trip_count, label %Header, label %exitLatch
+
+exitLatch:                                            ; preds = %Latch
+  ret i64 1
+
+exit:                                             ; preds = %Header
+  %result.in3.lcssa = phi i64* [ %result.in3, %Header ]
+  %result.le = load i64, i64* %result.in3.lcssa, align 8
+  ret i64 %result.le
+}
+
+; Same as test above but with profiling data that the most probable exit from
+; the loop is the header exiting block (not the latch block). So do not predicate.
+; LatchExitProbability: 0x000020e1 / 0x80000000 = 0.00%
+; ExitingBlockProbability: 0x7ffcbb86 / 0x80000000 = 99.99%
+define i64 @donot_predicate_prof(i64* nocapture readonly %arg, i32 %length, i64* nocapture readonly %arg2, i64* nocapture readonly %n_addr, i64 %i) {
+; CHECK-LABEL: donot_predicate_prof(
+; CHECK-LABEL: entry:
+entry:
+  %length.ext = zext i32 %length to i64
+  %n.pre = load i64, i64* %n_addr, align 4
+  br label %Header
+
+; CHECK-LABEL: Header:
+; CHECK:         %within.bounds = icmp ult i64 %j2, %length.ext
+; CHECK-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9)
+Header:                                          ; preds = %entry, %Latch
+  %result.in3 = phi i64* [ %arg2, %entry ], [ %arg, %Latch ]
+  %j2 = phi i64 [ 0, %entry ], [ %j.next, %Latch ]
+  %within.bounds = icmp ult i64 %j2, %length.ext
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+  %innercmp = icmp eq i64 %j2, %n.pre
+  %j.next = add nuw nsw i64 %j2, 1
+  br i1 %innercmp, label %Latch, label %exit, !prof !1
+
+Latch:                                           ; preds = %Header
+  %speculate_trip_count = icmp ult i64 %j.next, 1048576
+  br i1 %speculate_trip_count, label %Header, label %exitLatch, !prof !2
+
+exitLatch:                                            ; preds = %Latch
+  ret i64 1
+
+exit:                                             ; preds = %Header
+  %result.in3.lcssa = phi i64* [ %result.in3, %Header ]
+  %result.le = load i64, i64* %result.in3.lcssa, align 8
+  ret i64 %result.le
+}
+declare i64 @llvm.experimental.deoptimize.i64(...)
+declare void @llvm.experimental.guard(i1, ...)
+
+!1 = !{!"branch_weights", i32 104, i32 1042861}
+!2 = !{!"branch_weights", i32 255129, i32 1}