[IRCE] Support inverted range check's predicate

[aleksandr.popov]

Summary

IRCE expects true edge of range check's branch comes to loop.
If it meets reverse case - invert the branch.

[aleksandr.popov]

rebased

[mkazantsev]

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

378

You can modify IR at this point and then return false here:

for (auto *BBI : L->getBlocks())
  if (BranchInst *TBI = dyn_cast<BranchInst>(BBI->getTerminator()))
    InductiveRangeCheck::extractRangeChecksFromBranch(TBI, L, SE, BPI,
                                                      RangeChecks);
 
if (RangeChecks.empty())
  return false;

it would mean that you modify IR and then report that no modification was done and all analyzes can be preserved, which is wrong.

[skatkov]

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

374

Can we move this transformation AFTER profitability check? The idea is that if we will not add this check to consideration, no need to do modification.

I guess it should work:

1. Get operand idx of edge coming to loop
2. use this operand number in getEdgeProbability
3. if we pass check, do transformation if operand idx != 0.

1901

Is it Dead variable?

[skatkov]

Any words in git commit message?

[skatkov]

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

374

No operand idx but successor idx...

[aleksandr.popov]

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

374

Good point, thanks!

1901

Yeap, removed https://reviews.llvm.org/D154391

[skatkov]

In the commit message adds a couple of words about the reason - IRCE expects that true branch comes to loop.

[skatkov]

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

389

dot at the end, please.

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

 #include "llvm/IR/Value.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
 #include <iterator>
std::optional<Range> computeSafeIterationSpace(ScalarEvolution &SE,
                                                  const SCEVAddRecExpr *IndVar,
                                                  bool IsLatchSigned) const;
 
   /// Parse out a set of inductive range checks from \p BI and append them to \p
   /// Checks.
   ///
   /// NB! There may be conditions feeding into \p BI that aren't inductive range
   /// checks, and hence don't end up in \p Checks.
-  static void
-  extractRangeChecksFromBranch(BranchInst *BI, Loop *L, ScalarEvolution &SE,
-                               BranchProbabilityInfo *BPI,
-                               SmallVectorImpl<InductiveRangeCheck> &Checks);
+  static void extractRangeChecksFromBranch(
+      BranchInst *BI, Loop *L, ScalarEvolution &SE, BranchProbabilityInfo *BPI,
+      SmallVectorImpl<InductiveRangeCheck> &Checks, bool &Changed);
 };
 
 struct LoopStructure;
 
 class InductiveRangeCheckElimination {
   ScalarEvolution &SE;
   BranchProbabilityInfo *BPI;
   DominatorTree &DT;
void InductiveRangeCheck::extractRangeChecksFromCond(
   IRC.Begin = IndexAddRec->getStart();
   IRC.Step = IndexAddRec->getStepRecurrence(SE);
   IRC.CheckUse = &ConditionUse;
   Checks.push_back(IRC);
 }
 
 void InductiveRangeCheck::extractRangeChecksFromBranch(
     BranchInst *BI, Loop *L, ScalarEvolution &SE, BranchProbabilityInfo *BPI,
-    SmallVectorImpl<InductiveRangeCheck> &Checks) {
+    SmallVectorImpl<InductiveRangeCheck> &Checks, bool &Changed) {
   if (BI->isUnconditional() || BI->getParent() == L->getLoopLatch())
     return;
 
+  // IRCE expects that true successor of branch lies in a loop and false
+  // successor - outside it. Invert branch for opposite case
+  if (!L->contains(BI->getSuccessor(0)) && L->contains(BI->getSuccessor(1))) {
+    IRBuilder<> Builder(BI);
+    InvertBranch(BI, Builder);
+    if (BPI)
+      BPI->swapSuccEdgesProbabilities(BI->getParent());
+    Changed = true;
+  }
+
   BranchProbability LikelyTaken(15, 16);
 
   if (!SkipProfitabilityChecks && BPI &&
       BPI->getEdgeProbability(BI->getParent(), (unsigned)0) < LikelyTaken)
     return;
 
   SmallPtrSet<Value *, 8> Visited;
   InductiveRangeCheck::extractRangeChecksFromCond(L, SE, BI->getOperandUse(0),
                                                   Checks, Visited);
 }
 
 // Add metadata to the loop L to disable loop optimizations. Callers need to
 // confirm that optimizing loop L is not beneficial.
 static void DisableAllLoopOptsOnLoop(Loop &L) {
   // We do not care about any existing loopID related metadata for L, since we
   // are setting all loop metadata to false.
bool InductiveRangeCheckElimination::run(
   BasicBlock *Preheader = L->getLoopPreheader();
   if (!Preheader) {
     LLVM_DEBUG(dbgs() << "irce: loop has no preheader, leaving\n");
     return false;
   }
 
   LLVMContext &Context = Preheader->getContext();
   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);
+                                                        RangeChecks, Changed);
 
   if (RangeChecks.empty())
-    return false;
+    return Changed;
 
   auto PrintRecognizedRangeChecks = [&](raw_ostream &OS) {
     OS << "irce: looking at loop "; L->print(OS);
     OS << "irce: loop has " << RangeChecks.size()
        << " inductive range checks: \n";
     for (InductiveRangeCheck &IRC : RangeChecks)
       IRC.print(OS);
   };
 
   LLVM_DEBUG(PrintRecognizedRangeChecks(dbgs()));
 
   if (PrintRangeChecks)
     PrintRecognizedRangeChecks(errs());
 
   const char *FailureReason = nullptr;
   std::optional<LoopStructure> MaybeLoopStructure =
       LoopStructure::parseLoopStructure(SE, *L, FailureReason);
   if (!MaybeLoopStructure) {
     LLVM_DEBUG(dbgs() << "irce: could not parse loop structure: "
                       << FailureReason << "\n";);
-    return false;
+    return Changed;
   }
   LoopStructure LS = *MaybeLoopStructure;
   if (!isProfitableToTransform(*L, LS))
-    return false;
+    return Changed;
   const SCEVAddRecExpr *IndVar =
       cast<SCEVAddRecExpr>(SE.getMinusSCEV(SE.getSCEV(LS.IndVarBase), SE.getSCEV(LS.IndVarStep)));
 
   std::optional<InductiveRangeCheck::Range> SafeIterRange;
   Instruction *ExprInsertPt = Preheader->getTerminator();
 
   SmallVector<InductiveRangeCheck, 4> RangeChecksToEliminate;
   // Basing on the type of latch predicate, we interpret the IV iteration range
   // as signed or unsigned range. We use different min/max functions (signed or
   // unsigned) when intersecting this range with safe iteration ranges implied
   // by range checks.
   auto IntersectRange =
       LS.IsSignedPredicate ? IntersectSignedRange : IntersectUnsignedRange;
 
   IRBuilder<> B(ExprInsertPt);
   for (InductiveRangeCheck &IRC : RangeChecks) {
     auto Result = IRC.computeSafeIterationSpace(SE, IndVar,
                                                 LS.IsSignedPredicate);
     if (Result) {
       auto MaybeSafeIterRange = IntersectRange(SE, SafeIterRange, *Result);
       if (MaybeSafeIterRange) {
         assert(!MaybeSafeIterRange->isEmpty(SE, LS.IsSignedPredicate) &&
                "We should never return empty ranges!");
         RangeChecksToEliminate.push_back(IRC);
         SafeIterRange = *MaybeSafeIterRange;
       }
     }
   }
 
   if (!SafeIterRange)
-    return false;
+    return Changed;
 
   LoopConstrainer LC(*L, LI, LPMAddNewLoop, LS, SE, DT, *SafeIterRange);
-  bool Changed = LC.run();
 
-  if (Changed) {
+  if (LC.run()) {
+    Changed = true;
+
     auto PrintConstrainedLoopInfo = [L]() {
       dbgs() << "irce: in function ";
       dbgs() << L->getHeader()->getParent()->getName() << ": ";
       dbgs() << "constrained ";
       L->print(dbgs());
     };
 
     LLVM_DEBUG(PrintConstrainedLoopInfo());

llvm/test/Transforms/IRCE/non-loop-invariant-rhs-instr.ll

 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[INDVAR:%.*]] = phi i64 [ [[INDVAR_NEXT:%.*]], [[LATCH:%.*]] ], [ 1, [[LOOP_PREHEADER2]] ]
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ult i64 [[INDVAR]], [[LEN]]
 ; CHECK-NEXT:    br i1 true, label [[GUARDED:%.*]], label [[DEOPT_LOOPEXIT3:%.*]]
 ; CHECK:       guarded:
 ; CHECK-NEXT:    [[ADDR:%.*]] = getelementptr inbounds i32, ptr [[ARRAY:%.*]], i64 [[INDVAR]]
 ; CHECK-NEXT:    [[RES:%.*]] = load i32, ptr [[ADDR]], align 4
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[RES]], 0
-; CHECK-NEXT:    br i1 [[CMP]], label [[ZERO_LOOPEXIT_LOOPEXIT4:%.*]], label [[LATCH]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[RES]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[LATCH]], label [[ZERO_LOOPEXIT_LOOPEXIT4:%.*]]
 ; CHECK:       latch:
 ; CHECK-NEXT:    [[INDVAR_NEXT]] = add nuw nsw i64 [[INDVAR]], 2
 ; CHECK-NEXT:    [[RES2:%.*]] = mul i32 [[RES]], 3
 ; CHECK-NEXT:    [[TMP5:%.*]] = zext i32 [[A]] to i64
 ; CHECK-NEXT:    [[CMP2:%.*]] = icmp ugt i64 [[INDVAR_NEXT]], [[TMP5]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = icmp ult i64 [[INDVAR_NEXT]], [[EXIT_MAINLOOP_AT]]
 ; CHECK-NEXT:    [[TMP7:%.*]] = xor i1 [[TMP6]], true
 ; CHECK-NEXT:    br i1 [[TMP7]], label [[MAIN_EXIT_SELECTOR:%.*]], label [[LOOP]]

 ; CHECK-NEXT:    br label [[LOOP_POSTLOOP:%.*]]
 ; CHECK:       loop.postloop:
 ; CHECK-NEXT:    [[INDVAR_POSTLOOP:%.*]] = phi i64 [ [[INDVAR_NEXT_POSTLOOP:%.*]], [[LATCH_POSTLOOP]] ], [ [[INDVAR_COPY]], [[POSTLOOP]] ]
 ; CHECK-NEXT:    [[TMP9:%.*]] = icmp ult i64 [[INDVAR_POSTLOOP]], [[LEN]]
 ; CHECK-NEXT:    br i1 [[TMP9]], label [[GUARDED_POSTLOOP:%.*]], label [[DEOPT_LOOPEXIT:%.*]]
 ; CHECK:       guarded.postloop:
 ; CHECK-NEXT:    [[ADDR_POSTLOOP:%.*]] = getelementptr inbounds i32, ptr [[ARRAY]], i64 [[INDVAR_POSTLOOP]]
 ; CHECK-NEXT:    [[RES_POSTLOOP:%.*]] = load i32, ptr [[ADDR_POSTLOOP]], align 4
-; CHECK-NEXT:    [[CMP_POSTLOOP:%.*]] = icmp eq i32 [[RES_POSTLOOP]], 0
-; CHECK-NEXT:    br i1 [[CMP_POSTLOOP]], label [[ZERO_LOOPEXIT_LOOPEXIT:%.*]], label [[LATCH_POSTLOOP]]
+; CHECK-NEXT:    [[CMP_POSTLOOP:%.*]] = icmp ne i32 [[RES_POSTLOOP]], 0
+; CHECK-NEXT:    br i1 [[CMP_POSTLOOP]], label [[LATCH_POSTLOOP]], label [[ZERO_LOOPEXIT_LOOPEXIT:%.*]]
 ; CHECK:       latch.postloop:
 ; CHECK-NEXT:    [[INDVAR_NEXT_POSTLOOP]] = add nuw nsw i64 [[INDVAR_POSTLOOP]], 2
 ; CHECK-NEXT:    [[RES2_POSTLOOP]] = mul i32 [[RES_POSTLOOP]], 3
 ; CHECK-NEXT:    [[TMP10:%.*]] = zext i32 [[A]] to i64
 ; CHECK-NEXT:    [[CMP2_POSTLOOP:%.*]] = icmp ugt i64 [[INDVAR_NEXT_POSTLOOP]], [[TMP10]]
 ; CHECK-NEXT:    br i1 [[CMP2_POSTLOOP]], label [[LOOPEXIT_LOOPEXIT]], label [[LOOP_POSTLOOP]], !llvm.loop [[LOOP0:![0-9]+]], !irce.loop.clone !5
 ;
 preheader:

llvm/test/Transforms/IRCE/stride_more_than_1.ll

 exit:
   ret i32 %iv
 
 out_of_bounds:
   ret i32 -1
 }
 
 ; Same as test_09 but range check comparison is inversed.
-; TODO: IRCE is allowed.
+; IRCE is allowed.
 define i32 @test_10(ptr %p, ptr %capacity_p, ptr %num_elements_p) {
 ; CHECK-LABEL: define i32 @test_10
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[CAPACITY_P:%.*]], ptr [[NUM_ELEMENTS_P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CAPACITY:%.*]] = load i32, ptr [[CAPACITY_P]], align 4, !range [[RNG16]]
 ; CHECK-NEXT:    [[NUM_ELEMENTS:%.*]] = load i32, ptr [[NUM_ELEMENTS_P]], align 4, !range [[RNG16]]
 ; CHECK-NEXT:    [[LIMIT:%.*]] = sub i32 [[CAPACITY]], 3
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[CAPACITY]], -3
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw i32 [[CAPACITY]], 2147483646
+; CHECK-NEXT:    [[SMAX:%.*]] = call i32 @llvm.smax.i32(i32 [[TMP1]], i32 0)
+; CHECK-NEXT:    [[TMP2:%.*]] = sub i32 [[TMP0]], [[SMAX]]
+; CHECK-NEXT:    [[SMIN:%.*]] = call i32 @llvm.smin.i32(i32 [[LIMIT]], i32 0)
+; CHECK-NEXT:    [[SMAX2:%.*]] = call i32 @llvm.smax.i32(i32 [[SMIN]], i32 -1)
+; CHECK-NEXT:    [[TMP3:%.*]] = add nsw i32 [[SMAX2]], 1
+; CHECK-NEXT:    [[TMP4:%.*]] = mul i32 [[TMP2]], [[TMP3]]
+; CHECK-NEXT:    [[SMIN3:%.*]] = call i32 @llvm.smin.i32(i32 [[NUM_ELEMENTS]], i32 [[TMP4]])
+; CHECK-NEXT:    [[EXIT_MAINLOOP_AT:%.*]] = call i32 @llvm.smax.i32(i32 [[SMIN3]], i32 0)
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp slt i32 0, [[EXIT_MAINLOOP_AT]]
+; CHECK-NEXT:    br i1 [[TMP5]], label [[LOOP_PREHEADER:%.*]], label [[MAIN_PSEUDO_EXIT:%.*]]
+; CHECK:       loop.preheader:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[CAPACITY_CHECK:%.*]] = icmp sge i32 [[IV]], [[LIMIT]]
-; CHECK-NEXT:    br i1 [[CAPACITY_CHECK]], label [[OUT_OF_BOUNDS:%.*]], label [[BACKEDGE]], !prof [[PROF19:![0-9]+]]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ 0, [[LOOP_PREHEADER]] ]
+; CHECK-NEXT:    [[CAPACITY_CHECK:%.*]] = icmp slt i32 [[IV]], [[LIMIT]]
+; CHECK-NEXT:    br i1 true, label [[BACKEDGE]], label [[OUT_OF_BOUNDS_LOOPEXIT5:%.*]], !prof [[PROF17]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[IV_WIDE:%.*]] = zext i32 [[IV]] to i64
 ; CHECK-NEXT:    [[EL_PTR:%.*]] = getelementptr i32, ptr [[P]], i64 [[IV_WIDE]]
 ; CHECK-NEXT:    store i32 1, ptr [[EL_PTR]], align 4
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp slt i32 [[IV_NEXT]], [[NUM_ELEMENTS]]
-; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK-NEXT:    [[TMP6:%.*]] = icmp slt i32 [[IV_NEXT]], [[EXIT_MAINLOOP_AT]]
+; CHECK-NEXT:    br i1 [[TMP6]], label [[LOOP]], label [[MAIN_EXIT_SELECTOR:%.*]]
+; CHECK:       main.exit.selector:
+; CHECK-NEXT:    [[IV_NEXT_LCSSA:%.*]] = phi i32 [ [[IV_NEXT]], [[BACKEDGE]] ]
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
+; CHECK-NEXT:    [[TMP7:%.*]] = icmp slt i32 [[IV_NEXT_LCSSA]], [[NUM_ELEMENTS]]
+; CHECK-NEXT:    br i1 [[TMP7]], label [[MAIN_PSEUDO_EXIT]], label [[EXIT:%.*]]
+; CHECK:       main.pseudo.exit:
+; CHECK-NEXT:    [[IV_COPY:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT_LCSSA]], [[MAIN_EXIT_SELECTOR]] ]
+; CHECK-NEXT:    [[INDVAR_END:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV_NEXT_LCSSA]], [[MAIN_EXIT_SELECTOR]] ]
+; CHECK-NEXT:    br label [[POSTLOOP:%.*]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    [[IV_LCSSA1_PH:%.*]] = phi i32 [ [[IV_POSTLOOP:%.*]], [[BACKEDGE_POSTLOOP:%.*]] ]
+; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
+; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV_LCSSA]], [[MAIN_EXIT_SELECTOR]] ], [ [[IV_LCSSA1_PH]], [[EXIT_LOOPEXIT:%.*]] ]
 ; CHECK-NEXT:    ret i32 [[IV_LCSSA1]]
+; CHECK:       out_of_bounds.loopexit:
+; CHECK-NEXT:    br label [[OUT_OF_BOUNDS:%.*]]
+; CHECK:       out_of_bounds.loopexit5:
+; CHECK-NEXT:    br label [[OUT_OF_BOUNDS]]
 ; CHECK:       out_of_bounds:
 ; CHECK-NEXT:    ret i32 -1
+; CHECK:       postloop:
+; CHECK-NEXT:    br label [[LOOP_POSTLOOP:%.*]]
+; CHECK:       loop.postloop:
+; CHECK-NEXT:    [[IV_POSTLOOP]] = phi i32 [ [[IV_COPY]], [[POSTLOOP]] ], [ [[IV_NEXT_POSTLOOP:%.*]], [[BACKEDGE_POSTLOOP]] ]
+; CHECK-NEXT:    [[CAPACITY_CHECK_POSTLOOP:%.*]] = icmp slt i32 [[IV_POSTLOOP]], [[LIMIT]]
+; CHECK-NEXT:    br i1 [[CAPACITY_CHECK_POSTLOOP]], label [[BACKEDGE_POSTLOOP]], label [[OUT_OF_BOUNDS_LOOPEXIT:%.*]], !prof [[PROF17]]
+; CHECK:       backedge.postloop:
+; CHECK-NEXT:    [[IV_WIDE_POSTLOOP:%.*]] = zext i32 [[IV_POSTLOOP]] to i64
+; CHECK-NEXT:    [[EL_PTR_POSTLOOP:%.*]] = getelementptr i32, ptr [[P]], i64 [[IV_WIDE_POSTLOOP]]
+; CHECK-NEXT:    store i32 1, ptr [[EL_PTR_POSTLOOP]], align 4
+; CHECK-NEXT:    [[IV_NEXT_POSTLOOP]] = add nuw nsw i32 [[IV_POSTLOOP]], 4
+; CHECK-NEXT:    [[LOOP_COND_POSTLOOP:%.*]] = icmp slt i32 [[IV_NEXT_POSTLOOP]], [[NUM_ELEMENTS]]
+; CHECK-NEXT:    br i1 [[LOOP_COND_POSTLOOP]], label [[LOOP_POSTLOOP]], label [[EXIT_LOOPEXIT]], !llvm.loop [[LOOP19:![0-9]+]], !irce.loop.clone !6
 ;
 entry:
   %capacity = load i32, ptr %capacity_p, !range !4
   %num_elements = load i32, ptr %num_elements_p, !range !4
   %limit = sub i32 %capacity, 3
   br label %loop
 
 loop: