[IRCE] Relax profitability check

[JanJecmen]

IRCE currently has two profitability checks:

1. min number of iterations (10 by default)
2. branch is highly biased (> 15/16)

However, it may still be profitable to eliminate range checks even if the branch isn't as biased. Consider, for example, a loop with 100 iterations, where IRCE currently eliminates all 100 range checks. The same range checks performed over a loop with 200 iterations aren't eliminated because the branch is 50-50.

This patch proposes to relax the profitability checks of IRCE. Namely, instead of the two checks currenly in place, consider IRCE profitable if the branch probability scaled by the expected number of iterations (i.e., the estimated number of eliminated checks) is over a threshold. This covers the minimum number of iterations check (there are at least as many iterations as eliminated range checks), and changes the bias check from a percent of iterations to at least a constant threshold of eliminated checks.

If the number of iterations can't be estimated, the check falls back to the current 15/16 likelihood check.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Jan Ječmen (JanJecmen)

Changes

IRCE currently has two profitability checks:

1. min number of iterations (10 by default)
2. branch is highly biased (> 15/16)

However, it may still be profitable to eliminate range checks even if the branch isn't as biased. Consider, for example, a loop with 100 iterations, where IRCE currently eliminates all 100 range checks. The same range checks performed over a loop with 200 iterations aren't eliminated because the branch is 50-50.

This patch proposes to relax the profitability checks of IRCE. Namely, instead of the two checks currenly in place, consider IRCE profitable if the branch probability scaled by the expected number of iterations (i.e., the estimated number of eliminated checks) is over a threshold. This covers the minimum number of iterations check (there are at least as many iterations as eliminated range checks), and changes the bias check from a percent of iterations to at least a constant threshold of eliminated checks.

If the number of iterations can't be estimated, the check falls back to the current 15/16 likelihood check.

---

Full diff: https://github.com/llvm/llvm-project/pull/104659.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp (+48-52)
• (modified) llvm/test/Transforms/IRCE/low-iterations.ll (+2-2)

diff --git a/llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp b/llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp
index 104e8ceb796700..96850e11f1af24 100644
--- a/llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp
@@ -107,8 +107,8 @@ static cl::opt<bool> PrintRangeChecks("irce-print-range-checks", cl::Hidden,
 static cl::opt<bool> SkipProfitabilityChecks("irce-skip-profitability-checks",
                                              cl::Hidden, cl::init(false));
 
-static cl::opt<unsigned> MinRuntimeIterations("irce-min-runtime-iterations",
-                                              cl::Hidden, cl::init(10));
+static cl::opt<unsigned> MinEliminatedChecks("irce-min-eliminated-checks",
+                                             cl::Hidden, cl::init(10));
 
 static cl::opt<bool> AllowUnsignedLatchCondition("irce-allow-unsigned-latch",
                                                  cl::Hidden, cl::init(true));
@@ -132,15 +132,9 @@ static cl::opt<bool>
 
 namespace {
 
-/// An inductive range check is conditional branch in a loop with
-///
-///  1. a very cold successor (i.e. the branch jumps to that successor very
-///     rarely)
-///
-///  and
-///
-///  2. a condition that is provably true for some contiguous range of values
-///     taken by the containing loop's induction variable.
+/// An inductive range check is conditional branch in a loop with a condition
+/// that is provably true for some contiguous range of values taken by the
+/// containing loop's induction variable.
 ///
 class InductiveRangeCheck {
 
@@ -235,6 +229,7 @@ class InductiveRangeCheck {
   /// checks, and hence don't end up in \p Checks.
   static void extractRangeChecksFromBranch(
       BranchInst *BI, Loop *L, ScalarEvolution &SE, BranchProbabilityInfo *BPI,
+      std::optional<unsigned int> EstimatedTripCount,
       SmallVectorImpl<InductiveRangeCheck> &Checks, bool &Changed);
 };
 
@@ -248,9 +243,11 @@ class InductiveRangeCheckElimination {
       std::optional<llvm::function_ref<llvm::BlockFrequencyInfo &()>>;
   GetBFIFunc GetBFI;
 
-  // Returns true if it is profitable to do a transform basing on estimation of
-  // number of iterations.
-  bool isProfitableToTransform(const Loop &L, LoopStructure &LS);
+  // Returns the estimated number of iterations based on block frequency info if
+  // available, or on branch probability info. Nullopt is returned if the number
+  // of iterations cannot be estimated.
+  std::optional<unsigned int> estimatedTripCount(const Loop &L,
+                                                 LoopStructure &LS);
 
 public:
   InductiveRangeCheckElimination(ScalarEvolution &SE,
@@ -524,6 +521,7 @@ void InductiveRangeCheck::extractRangeChecksFromCond(
 
 void InductiveRangeCheck::extractRangeChecksFromBranch(
     BranchInst *BI, Loop *L, ScalarEvolution &SE, BranchProbabilityInfo *BPI,
+    std::optional<unsigned int> EstimatedTripCount,
     SmallVectorImpl<InductiveRangeCheck> &Checks, bool &Changed) {
   if (BI->isUnconditional() || BI->getParent() == L->getLoopLatch())
     return;
@@ -531,11 +529,19 @@ void InductiveRangeCheck::extractRangeChecksFromBranch(
   unsigned IndexLoopSucc = L->contains(BI->getSuccessor(0)) ? 0 : 1;
   assert(L->contains(BI->getSuccessor(IndexLoopSucc)) &&
          "No edges coming to loop?");
-  BranchProbability LikelyTaken(15, 16);
 
-  if (!SkipProfitabilityChecks && BPI &&
-      BPI->getEdgeProbability(BI->getParent(), IndexLoopSucc) < LikelyTaken)
-    return;
+  if (!SkipProfitabilityChecks && BPI) {
+    auto SuccessProbability =
+        BPI->getEdgeProbability(BI->getParent(), IndexLoopSucc);
+    if (EstimatedTripCount) {
+      if (SuccessProbability.scale(*EstimatedTripCount) < MinEliminatedChecks)
+        return;
+    } else {
+      BranchProbability LikelyTaken(15, 16);
+      if (SuccessProbability < LikelyTaken)
+        return;
+    }
+  }
 
   // IRCE expects branch's true edge comes to loop. Invert branch for opposite
   // case.
@@ -940,35 +946,25 @@ PreservedAnalyses IRCEPass::run(Function &F, FunctionAnalysisManager &AM) {
   return getLoopPassPreservedAnalyses();
 }
 
-bool
-InductiveRangeCheckElimination::isProfitableToTransform(const Loop &L,
-                                                        LoopStructure &LS) {
-  if (SkipProfitabilityChecks)
-    return true;
+std::optional<unsigned int>
+InductiveRangeCheckElimination::estimatedTripCount(const Loop &L,
+                                                   LoopStructure &LS) {
   if (GetBFI) {
     BlockFrequencyInfo &BFI = (*GetBFI)();
     uint64_t hFreq = BFI.getBlockFreq(LS.Header).getFrequency();
     uint64_t phFreq = BFI.getBlockFreq(L.getLoopPreheader()).getFrequency();
-    if (phFreq != 0 && hFreq != 0 && (hFreq / phFreq < MinRuntimeIterations)) {
-      LLVM_DEBUG(dbgs() << "irce: could not prove profitability: "
-                        << "the estimated number of iterations basing on "
-                           "frequency info is " << (hFreq / phFreq) << "\n";);
-      return false;
-    }
-    return true;
+    if (phFreq == 0 || hFreq == 0)
+      return std::nullopt;
+    return {hFreq / phFreq};
   }
 
   if (!BPI)
-    return true;
+    return std::nullopt;
   BranchProbability ExitProbability =
       BPI->getEdgeProbability(LS.Latch, LS.LatchBrExitIdx);
-  if (ExitProbability > BranchProbability(1, MinRuntimeIterations)) {
-    LLVM_DEBUG(dbgs() << "irce: could not prove profitability: "
-                      << "the exit probability is too big " << ExitProbability
-                      << "\n";);
-    return false;
-  }
-  return true;
+  if (ExitProbability.isUnknown() || ExitProbability.isZero())
+    return std::nullopt;
+  return {ExitProbability.getDenominator() / ExitProbability.getNumerator()};
 }
 
 bool InductiveRangeCheckElimination::run(
@@ -985,13 +981,25 @@ bool InductiveRangeCheckElimination::run(
   }
 
   LLVMContext &Context = Preheader->getContext();
+
+  const char *FailureReason = nullptr;
+  std::optional<LoopStructure> MaybeLoopStructure =
+      LoopStructure::parseLoopStructure(SE, *L, AllowUnsignedLatchCondition,
+                                        FailureReason);
+  if (!MaybeLoopStructure) {
+    LLVM_DEBUG(dbgs() << "irce: could not parse loop structure: "
+                      << FailureReason << "\n";);
+    return false;
+  }
+  LoopStructure LS = *MaybeLoopStructure;
+
   SmallVector<InductiveRangeCheck, 16> RangeChecks;
   bool Changed = false;
 
   for (auto *BBI : L->getBlocks())
     if (BranchInst *TBI = dyn_cast<BranchInst>(BBI->getTerminator()))
-      InductiveRangeCheck::extractRangeChecksFromBranch(TBI, L, SE, BPI,
-                                                        RangeChecks, Changed);
+      InductiveRangeCheck::extractRangeChecksFromBranch(
+          TBI, L, SE, BPI, estimatedTripCount(*L, LS), RangeChecks, Changed);
 
   if (RangeChecks.empty())
     return Changed;
@@ -1009,18 +1017,6 @@ bool InductiveRangeCheckElimination::run(
   if (PrintRangeChecks)
     PrintRecognizedRangeChecks(errs());
 
-  const char *FailureReason = nullptr;
-  std::optional<LoopStructure> MaybeLoopStructure =
-      LoopStructure::parseLoopStructure(SE, *L, AllowUnsignedLatchCondition,
-                                        FailureReason);
-  if (!MaybeLoopStructure) {
-    LLVM_DEBUG(dbgs() << "irce: could not parse loop structure: "
-                      << FailureReason << "\n";);
-    return Changed;
-  }
-  LoopStructure LS = *MaybeLoopStructure;
-  if (!isProfitableToTransform(*L, LS))
-    return Changed;
   const SCEVAddRecExpr *IndVar =
       cast<SCEVAddRecExpr>(SE.getMinusSCEV(SE.getSCEV(LS.IndVarBase), SE.getSCEV(LS.IndVarStep)));
 
diff --git a/llvm/test/Transforms/IRCE/low-iterations.ll b/llvm/test/Transforms/IRCE/low-iterations.ll
index 071ab4d0156852..e044c455fe6e2b 100644
--- a/llvm/test/Transforms/IRCE/low-iterations.ll
+++ b/llvm/test/Transforms/IRCE/low-iterations.ll
@@ -1,5 +1,5 @@
-; RUN: opt -verify-loop-info -irce-print-changed-loops -passes=irce -irce-min-runtime-iterations=3 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK-NO
-; RUN: opt -verify-loop-info -irce-print-changed-loops -passes=irce -irce-min-runtime-iterations=0 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK-YES
+; RUN: opt -verify-loop-info -irce-print-changed-loops -passes=irce -irce-min-eliminated-checks=3 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK-NO
+; RUN: opt -verify-loop-info -irce-print-changed-loops -passes=irce -irce-min-eliminated-checks=0 < %s 2>&1 | FileCheck %s --check-prefixes=CHECK-YES
 
 ; CHECK-YES: constrained Loop
 ; CHECK-NO-NOT: constrained Loop

[aleks-tmb]

Consider, for example, a loop with 100 iterations, where IRCE currently eliminates all 100 range checks. The same range checks performed over a loop with 200 iterations aren't eliminated because the branch is 50-50.

Could you please add a test where range checks are not eliminated by the current IRCE, but are eliminated by your patch, to demonstrate the impact of the changes?

[JanJecmen]

@aleks-tmb Hi, thanks for the review. I addressed your comments and pushed a new version. I tweaked the new test to better show the effect: IRCE previously wouldn't fire because the range check branch weights are 1:1, but now it eliminates 50 of the 100 checks. Setting the threshold to 51 then prevents it from applying. I also refactored some small details (used BranchProbability existing API for inversion, hoisted a loop invariant computation of the loop trip count estimate).

[aleks-tmb]

LGTM

[github-actions]

@JanJecmen Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!