Teach IRCE to look at branch weights when recognizing range checks

[sanjoy]

Summary

Splitting a loop to make range checks redundant is profitable only if the range check "never" fails. Make this fact a part of recognizing a range check -- a branch is a range check only if it is expected to pass (via branch_weights metadata).

[hfinkel]

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

361

Please don't access the branch_weights metadata directly. Use the BranchProbabilityInfo analysis (which also includes some useful static heuristics).

[sanjoy]

Address Hal's comment and use BranchProbabilityInfo instead of manually parsing branch_weights metadata.

[hfinkel]

LGTM.

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

 //     } else {
 //       throw_out_of_bounds();
 //     }
 //   }
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/Optional.h"
 
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/InstructionSimplify.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/Analysis/ValueTracking.h"
 
#endif
   /// Computes a range for the induction variable in which the range check is
   /// redundant and can be constant-folded away.
   Optional<Range> computeSafeIterationSpace(ScalarEvolution &SE,
                                             IRBuilder<> &B) const;
 
   /// Create an inductive range check out of BI if possible, else return
   /// nullptr.
   static InductiveRangeCheck *create(AllocatorTy &Alloc, BranchInst *BI,
-                                     Loop *L, ScalarEvolution &SE);
+                                     Loop *L, ScalarEvolution &SE,
+                                     BranchProbabilityInfo &BPI);
 };
 
 class InductiveRangeCheckElimination : public LoopPass {
   InductiveRangeCheck::AllocatorTy Allocator;
 
 public:
   static char ID;
   InductiveRangeCheckElimination() : LoopPass(ID) {
     initializeInductiveRangeCheckEliminationPass(
         *PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addRequiredID(LoopSimplifyID);
     AU.addRequiredID(LCSSAID);
     AU.addRequired<ScalarEvolution>();
+    AU.addRequired<BranchProbabilityInfo>();
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override;
 };
 
 char InductiveRangeCheckElimination::ID = 0;
 }
 
DEBUG(dbgs() << " in function: " << L->getHeader()->getParent()->getName()
           dbgs() << " UpperLimit is not loop invariant: "
                  << UpperLimit->getName() << "\n";);
     return false;
   }
 
   return true;
 }
 
+
 InductiveRangeCheck *
 InductiveRangeCheck::create(InductiveRangeCheck::AllocatorTy &A, BranchInst *BI,
-                            Loop *L, ScalarEvolution &SE) {
+                            Loop *L, ScalarEvolution &SE,
+                            BranchProbabilityInfo &BPI) {
 
   if (BI->isUnconditional() || BI->getParent() == L->getLoopLatch())
     return nullptr;
 
+  BranchProbability LikelyTaken(15, 16);
+
+  if (BPI.getEdgeProbability(BI->getParent(), (unsigned) 0) < LikelyTaken)
+    return nullptr;
+
   Value *Length = nullptr;
   const SCEV *IndexSCEV = nullptr;
 
   if (!SplitRangeCheckCondition(L, SE, BI->getCondition(), IndexSCEV, Length))
     return nullptr;
 
   assert(IndexSCEV && Length && "contract with SplitRangeCheckCondition!");
 
if (!Preheader) {
     DEBUG(dbgs() << "irce: loop has no preheader, leaving\n");
     return false;
   }
 
   LLVMContext &Context = Preheader->getContext();
   InductiveRangeCheck::AllocatorTy IRCAlloc;
   SmallVector<InductiveRangeCheck *, 16> RangeChecks;
   ScalarEvolution &SE = getAnalysis<ScalarEvolution>();
+  BranchProbabilityInfo &BPI = getAnalysis<BranchProbabilityInfo>();
 
   for (auto BBI : L->getBlocks())
     if (BranchInst *TBI = dyn_cast<BranchInst>(BBI->getTerminator()))
       if (InductiveRangeCheck *IRC =
-              InductiveRangeCheck::create(IRCAlloc, TBI, L, SE))
+          InductiveRangeCheck::create(IRCAlloc, TBI, L, SE, BPI))
         RangeChecks.push_back(IRC);
 
   if (RangeChecks.empty())
     return false;
 
   DEBUG(dbgs() << "irce: looking at loop "; L->print(dbgs());
         dbgs() << "irce: loop has " << RangeChecks.size()
                << " inductive range checks: \n";

test/Transforms/IRCE/multiple-access-no-preloop.ll

 ; RUN: opt -irce -S < %s | FileCheck %s
 
 define void @multiple_access_no_preloop(
     i32* %arr_a, i32* %a_len_ptr, i32* %arr_b, i32* %b_len_ptr, i32 %n) {
 
  entry:
   %len.a = load i32* %a_len_ptr, !range !0
   %len.b = load i32* %b_len_ptr, !range !0
   %first.itr.check = icmp sgt i32 %n, 0
   br i1 %first.itr.check, label %loop, label %exit
 
  loop:
   %idx = phi i32 [ 0, %entry ] , [ %idx.next, %in.bounds.b ]
   %idx.next = add i32 %idx, 1
   %abc.a = icmp slt i32 %idx, %len.a
-  br i1 %abc.a, label %in.bounds.a, label %out.of.bounds
+  br i1 %abc.a, label %in.bounds.a, label %out.of.bounds, !prof !1
 
  in.bounds.a:
   %addr.a = getelementptr i32* %arr_a, i32 %idx
   store i32 0, i32* %addr.a
   %abc.b = icmp slt i32 %idx, %len.b
-  br i1 %abc.b, label %in.bounds.b, label %out.of.bounds
+  br i1 %abc.b, label %in.bounds.b, label %out.of.bounds, !prof !1
 
  in.bounds.b:
   %addr.b = getelementptr i32* %arr_b, i32 %idx
   store i32 -1, i32* %addr.b
   %next = icmp slt i32 %idx.next, %n
   br i1 %next, label %loop, label %exit
 
  out.of.bounds:

 ; CHECK: [[main_loop_cond:[^ ]+]] = icmp slt i32 %idx.next, [[upper_bound]]
 ; CHECK: br i1 [[main_loop_cond]], label %loop, label %main.exit.selector
 
 ; CHECK-LABEL: in.bounds.b.postloop:
 ; CHECK: %next.postloop = icmp slt i32 %idx.next.postloop, %n
 ; CHECK: br i1 %next.postloop, label %loop.postloop, label %exit.loopexit
 
 !0 = !{i32 0, i32 2147483647}
+!1 = !{!"branch_weights", i32 64, i32 4}

test/Transforms/IRCE/not-likely-taken.ll

+; RUN: opt -verify-loop-info -irce-print-changed-loops -irce < %s 2>&1 | FileCheck %s
+
+; CHECK-NOT: constrained Loop
+
+define void @multiple_access_no_preloop(
+    i32* %arr_a, i32* %a_len_ptr, i32* %arr_b, i32* %b_len_ptr, i32 %n) {
+
+ entry:
+  %len.a = load i32* %a_len_ptr, !range !0
+  %len.b = load i32* %b_len_ptr, !range !0
+  %first.itr.check = icmp sgt i32 %n, 0
+  br i1 %first.itr.check, label %loop, label %exit
+
+ loop:
+  %idx = phi i32 [ 0, %entry ] , [ %idx.next, %in.bounds.b ]
+  %idx.next = add i32 %idx, 1
+  %abc.a = icmp slt i32 %idx, %len.a
+  br i1 %abc.a, label %in.bounds.a, label %out.of.bounds, !prof !1
+
+ in.bounds.a:
+  %addr.a = getelementptr i32* %arr_a, i32 %idx
+  store i32 0, i32* %addr.a
+  %abc.b = icmp slt i32 %idx, %len.b
+  br i1 %abc.b, label %in.bounds.b, label %out.of.bounds, !prof !1
+
+ in.bounds.b:
+  %addr.b = getelementptr i32* %arr_b, i32 %idx
+  store i32 -1, i32* %addr.b
+  %next = icmp slt i32 %idx.next, %n
+  br i1 %next, label %loop, label %exit
+
+ out.of.bounds:
+  ret void
+
+ exit:
+  ret void
+}
+
+!0 = !{i32 0, i32 2147483647}
+!1 = !{!"branch_weights", i32 1, i32 1}
 No newline at end of file

test/Transforms/IRCE/single-access-no-preloop.ll

 ; RUN: opt -irce -S < %s | FileCheck %s
 
 define void @single_access_no_preloop_no_offset(i32 *%arr, i32 *%a_len_ptr, i32 %n) {
  entry:
   %len = load i32* %a_len_ptr, !range !0
   %first.itr.check = icmp sgt i32 %n, 0
   br i1 %first.itr.check, label %loop, label %exit
 
  loop:
   %idx = phi i32 [ 0, %entry ] , [ %idx.next, %in.bounds ]
   %idx.next = add i32 %idx, 1
   %abc = icmp slt i32 %idx, %len
-  br i1 %abc, label %in.bounds, label %out.of.bounds
+  br i1 %abc, label %in.bounds, label %out.of.bounds, !prof !1
 
  in.bounds:
   %addr = getelementptr i32* %arr, i32 %idx
   store i32 0, i32* %addr
   %next = icmp slt i32 %idx.next, %n
   br i1 %next, label %loop, label %exit
 
  out.of.bounds:
entry:
   %first.itr.check = icmp sgt i32 %n, 0
   br i1 %first.itr.check, label %loop, label %exit
 
  loop:
   %idx = phi i32 [ 0, %entry ] , [ %idx.next, %in.bounds ]
   %idx.next = add i32 %idx, 1
   %idx.for.abc = add i32 %idx, 4
   %abc = icmp slt i32 %idx.for.abc, %len
-  br i1 %abc, label %in.bounds, label %out.of.bounds
+  br i1 %abc, label %in.bounds, label %out.of.bounds, !prof !1
 
  in.bounds:
   %addr = getelementptr i32* %arr, i32 %idx.for.abc
   store i32 0, i32* %addr
   %next = icmp slt i32 %idx.next, %n
   br i1 %next, label %loop, label %exit
 
  out.of.bounds:

 ; CHECK-LABEL: loop.postloop:
 ; CHECK: %idx.postloop = phi i32 [ %idx.next.postloop, %in.bounds.postloop ], [ %idx.copy, %postloop ]
 
 ; CHECK-LABEL: in.bounds.postloop:
 ; CHECK: %next.postloop = icmp slt i32 %idx.next.postloop, %n
 ; CHECK: br i1 %next.postloop, label %loop.postloop, label %exit.loopexit
 
 !0 = !{i32 0, i32 2147483647}
+!1 = !{!"branch_weights", i32 64, i32 4}

test/Transforms/IRCE/single-access-with-preloop.ll

 ; RUN: opt -irce -S < %s | FileCheck %s
 
 define void @single_access_with_preloop(i32 *%arr, i32 *%a_len_ptr, i32 %n, i32 %offset) {
  entry:
   %len = load i32* %a_len_ptr, !range !0
   %first.itr.check = icmp sgt i32 %n, 0
   br i1 %first.itr.check, label %loop, label %exit
 
  loop:
   %idx = phi i32 [ 0, %entry ] , [ %idx.next, %in.bounds ]
   %idx.next = add i32 %idx, 1
   %array.idx = add i32 %idx, %offset
   %abc.high = icmp slt i32 %array.idx, %len
   %abc.low = icmp sge i32 %array.idx, 0
   %abc = and i1 %abc.low, %abc.high
-  br i1 %abc, label %in.bounds, label %out.of.bounds
+  br i1 %abc, label %in.bounds, label %out.of.bounds, !prof !1
 
  in.bounds:
   %addr = getelementptr i32* %arr, i32 %array.idx
   store i32 0, i32* %addr
   %next = icmp slt i32 %idx.next, %n
   br i1 %next, label %loop, label %exit
 
  out.of.bounds:

 ; CHECK: [[preloop_its_left:[^ ]+]] = icmp slt i32 %idx.next.preloop, %n
 ; CHECK: br i1 [[preloop_its_left]], label %preloop.pseudo.exit, label %exit.loopexit
 
 ; CHECK-LABEL: in.bounds.postloop:
 ; CHECK: %next.postloop = icmp slt i32 %idx.next.postloop, %n
 ; CHECK: br i1 %next.postloop, label %loop.postloop, label %exit.loopexit
 
 !0 = !{i32 0, i32 2147483647}
+!1 = !{!"branch_weights", i32 64, i32 4}

test/Transforms/IRCE/with-parent-loops.ll

entry:
   %len = load i32* %a_len_ptr, !range !0
   %first.itr.check = icmp sgt i32 %n, 0
   br i1 %first.itr.check, label %loop, label %exit
 
 loop:                                             ; preds = %in.bounds, %entry
   %idx = phi i32 [ 0, %entry ], [ %idx.next, %in.bounds ]
   %idx.next = add i32 %idx, 1
   %abc = icmp slt i32 %idx, %len
-  br i1 %abc, label %in.bounds, label %out.of.bounds
+  br i1 %abc, label %in.bounds, label %out.of.bounds, !prof !1
 
 in.bounds:                                        ; preds = %loop
   %addr = getelementptr i32* %arr, i32 %idx
   store i32 0, i32* %addr
   %next = icmp slt i32 %idx.next, %n
   br i1 %next, label %loop, label %exit
 
 out.of.bounds:                                    ; preds = %loop
loop:                                             ; preds = %inner_loop.exit, %entry
   %len.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i, label %loop.i, label %exit.i
 
 loop.i:                                           ; preds = %in.bounds.i, %loop
   %idx.i = phi i32 [ 0, %loop ], [ %idx.next.i, %in.bounds.i ]
   %idx.next.i = add i32 %idx.i, 1
   %abc.i = icmp slt i32 %idx.i, %len.i
-  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i
+  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i, !prof !1
 
 in.bounds.i:                                      ; preds = %loop.i
   %addr.i = getelementptr i32* %arr, i32 %idx.i
   store i32 0, i32* %addr.i
   %next.i = icmp slt i32 %idx.next.i, %n
   br i1 %next.i, label %loop.i, label %exit.i
 
 out.of.bounds.i:                                  ; preds = %loop.i
loop.i:                                           ; preds = %inner_loop.exit.i, %loop
   %len.i.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i.i, label %loop.i.i, label %exit.i.i
 
 loop.i.i:                                         ; preds = %in.bounds.i.i, %loop.i
   %idx.i.i = phi i32 [ 0, %loop.i ], [ %idx.next.i.i, %in.bounds.i.i ]
   %idx.next.i.i = add i32 %idx.i.i, 1
   %abc.i.i = icmp slt i32 %idx.i.i, %len.i.i
-  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i
+  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i, !prof !1
 
 in.bounds.i.i:                                    ; preds = %loop.i.i
   %addr.i.i = getelementptr i32* %arr, i32 %idx.i.i
   store i32 0, i32* %addr.i.i
   %next.i.i = icmp slt i32 %idx.next.i.i, %n
   br i1 %next.i.i, label %loop.i.i, label %exit.i.i
 
 out.of.bounds.i.i:                                ; preds = %loop.i.i
loop:                                             ; preds = %inner_loop.exit12, %entry
   %len.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i, label %loop.i, label %exit.i
 
 loop.i:                                           ; preds = %in.bounds.i, %loop
   %idx.i = phi i32 [ 0, %loop ], [ %idx.next.i, %in.bounds.i ]
   %idx.next.i = add i32 %idx.i, 1
   %abc.i = icmp slt i32 %idx.i, %len.i
-  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i
+  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i, !prof !1
 
 in.bounds.i:                                      ; preds = %loop.i
   %addr.i = getelementptr i32* %arr, i32 %idx.i
   store i32 0, i32* %addr.i
   %next.i = icmp slt i32 %idx.next.i, %n
   br i1 %next.i, label %loop.i, label %exit.i
 
 out.of.bounds.i:                                  ; preds = %loop.i
   br label %inner_loop.exit
 
 exit.i:                                           ; preds = %in.bounds.i, %loop
   br label %inner_loop.exit
 
 inner_loop.exit:                                  ; preds = %exit.i, %out.of.bounds.i
   %len.i1 = load i32* %a_len_ptr, !range !0
   %first.itr.check.i2 = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i2, label %loop.i6, label %exit.i11
 
 loop.i6:                                          ; preds = %in.bounds.i9, %inner_loop.exit
   %idx.i3 = phi i32 [ 0, %inner_loop.exit ], [ %idx.next.i4, %in.bounds.i9 ]
   %idx.next.i4 = add i32 %idx.i3, 1
   %abc.i5 = icmp slt i32 %idx.i3, %len.i1
-  br i1 %abc.i5, label %in.bounds.i9, label %out.of.bounds.i10
+  br i1 %abc.i5, label %in.bounds.i9, label %out.of.bounds.i10, !prof !1
 
 in.bounds.i9:                                     ; preds = %loop.i6
   %addr.i7 = getelementptr i32* %arr, i32 %idx.i3
   store i32 0, i32* %addr.i7
   %next.i8 = icmp slt i32 %idx.next.i4, %n
   br i1 %next.i8, label %loop.i6, label %exit.i11
 
 out.of.bounds.i10:                                ; preds = %loop.i6
loop.i:                                           ; preds = %inner_loop.exit.i, %loop
   %len.i.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i.i, label %loop.i.i, label %exit.i.i
 
 loop.i.i:                                         ; preds = %in.bounds.i.i, %loop.i
   %idx.i.i = phi i32 [ 0, %loop.i ], [ %idx.next.i.i, %in.bounds.i.i ]
   %idx.next.i.i = add i32 %idx.i.i, 1
   %abc.i.i = icmp slt i32 %idx.i.i, %len.i.i
-  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i
+  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i, !prof !1
 
 in.bounds.i.i:                                    ; preds = %loop.i.i
   %addr.i.i = getelementptr i32* %arr, i32 %idx.i.i
   store i32 0, i32* %addr.i.i
   %next.i.i = icmp slt i32 %idx.next.i.i, %n
   br i1 %next.i.i, label %loop.i.i, label %exit.i.i
 
 out.of.bounds.i.i:                                ; preds = %loop.i.i
loop.i6:                                          ; preds = %inner_loop.exit.i16, %with_parent.exit
   %len.i.i4 = load i32* %a_len_ptr, !range !0
   %first.itr.check.i.i5 = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i.i5, label %loop.i.i10, label %exit.i.i15
 
 loop.i.i10:                                       ; preds = %in.bounds.i.i13, %loop.i6
   %idx.i.i7 = phi i32 [ 0, %loop.i6 ], [ %idx.next.i.i8, %in.bounds.i.i13 ]
   %idx.next.i.i8 = add i32 %idx.i.i7, 1
   %abc.i.i9 = icmp slt i32 %idx.i.i7, %len.i.i4
-  br i1 %abc.i.i9, label %in.bounds.i.i13, label %out.of.bounds.i.i14
+  br i1 %abc.i.i9, label %in.bounds.i.i13, label %out.of.bounds.i.i14, !prof !1
 
 in.bounds.i.i13:                                  ; preds = %loop.i.i10
   %addr.i.i11 = getelementptr i32* %arr, i32 %idx.i.i7
   store i32 0, i32* %addr.i.i11
   %next.i.i12 = icmp slt i32 %idx.next.i.i8, %n
   br i1 %next.i.i12, label %loop.i.i10, label %exit.i.i15
 
 out.of.bounds.i.i14:                              ; preds = %loop.i.i10
loop:                                             ; preds = %with_parent.exit, %entry
   %len.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i, label %loop.i, label %exit.i
 
 loop.i:                                           ; preds = %in.bounds.i, %loop
   %idx.i = phi i32 [ 0, %loop ], [ %idx.next.i, %in.bounds.i ]
   %idx.next.i = add i32 %idx.i, 1
   %abc.i = icmp slt i32 %idx.i, %len.i
-  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i
+  br i1 %abc.i, label %in.bounds.i, label %out.of.bounds.i, !prof !1
 
 in.bounds.i:                                      ; preds = %loop.i
   %addr.i = getelementptr i32* %arr, i32 %idx.i
   store i32 0, i32* %addr.i
   %next.i = icmp slt i32 %idx.next.i, %n
   br i1 %next.i, label %loop.i, label %exit.i
 
 out.of.bounds.i:                                  ; preds = %loop.i
loop.i4:                                          ; preds = %inner_loop.exit.i, %inner_loop.exit
   %len.i.i = load i32* %a_len_ptr, !range !0
   %first.itr.check.i.i = icmp sgt i32 %n, 0
   br i1 %first.itr.check.i.i, label %loop.i.i, label %exit.i.i
 
 loop.i.i:                                         ; preds = %in.bounds.i.i, %loop.i4
   %idx.i.i = phi i32 [ 0, %loop.i4 ], [ %idx.next.i.i, %in.bounds.i.i ]
   %idx.next.i.i = add i32 %idx.i.i, 1
   %abc.i.i = icmp slt i32 %idx.i.i, %len.i.i
-  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i
+  br i1 %abc.i.i, label %in.bounds.i.i, label %out.of.bounds.i.i, !prof !1
 
 in.bounds.i.i:                                    ; preds = %loop.i.i
   %addr.i.i = getelementptr i32* %arr, i32 %idx.i.i
   store i32 0, i32* %addr.i.i
   %next.i.i = icmp slt i32 %idx.next.i.i, %n
   br i1 %next.i.i, label %loop.i.i, label %exit.i.i
 
 out.of.bounds.i.i:                                ; preds = %loop.i.i

 
 exit:                                             ; preds = %with_parent.exit
   ret void
 }
 
 attributes #0 = { alwaysinline }
 
 !0 = !{i32 0, i32 2147483647}
+!1 = !{!"branch_weights", i32 64, i32 4}