[IRCE] Recognize loops with ne/eq latch conditions

[mkazantsev]

Summary

In some particular cases eq/ne conditions can be turned into equivalent slt/sgt conditions.
This patch teaches parseLoopStructure to handle some of these cases.

[sanjoy]

Did you consider a somewhat more general trick of using the loop's BE taken count to derive an equivalent BE condition? E.g. if the loops be taken count is 100, and the post-incremented indvar (the one the BE condition uses) is {5,+,1} (everything is i32) then the BE taken condition can be ++I s< 105 (and you can derive this without looking at whether the BE condition was expressed using SLT or NE or EQ etc.).

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

810 ↗

(On Diff #105264)

I don't think this equivalence holds if the RHS is negative.

[mkazantsev]

Seems that the patch has a bug, needs rework.

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

810 ↗

(On Diff #105264)

What's the problem with negative RHS? There is no conceptual difference between

while (++i != 5) ---> while (++i < 5)

and

while (++i != -5) ---> while (++i < -5)

Maybe what you meant was "if the RHS is less than the IV's initial value"? Such loops get rejected earlier, test_02 exercises this situation. Will add an assert on that.

850 ↗

(On Diff #105264)

Bug here: if we have reduced RightSCEV by one on line 822, we should not increase RightValue by 1.

901 ↗

(On Diff #105264)

Same here.

[mkazantsev]

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

810 ↗

(On Diff #105264)

Actually not earlier, but under this condition we will bail out on line 842.

[mkazantsev]

Extended test base, fix a bug with shifting comparison borders when we shouldn't.

[mkazantsev]

In D35010#802926, @sanjoy wrote:

Did you consider a somewhat more general trick of using the loop's BE taken count to derive an equivalent BE condition? E.g. if the loops be taken count is 100, and the post-incremented indvar (the one the BE condition uses) is {5,+,1} (everything is i32) then the BE taken condition can be ++I s< 105 (and you can derive this without looking at whether the BE condition was expressed using SLT or NE or EQ etc.).

I think we still need to parse the loop structure to be able to duplicate it correctly, but the idea is interesting. I will think around it.

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

810 ↗

(On Diff #105264)

Tests 2, 4, 6, 8 check such situations and make sure that the IRCE doesn't apply to them.

[mkazantsev]

Ping

[sanjoy]

Some comments / questions inline.

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

818 ↗

(On Diff #105808)

I'm not sure why you need to handle this case differently than the above. That is,

while (true) {
  if (++i == len)
    break;
}

seems equivalent to

while (++i != len) {
}

853 ↗

(On Diff #105808)

I'm missing a link here. IIUC:

• You started with the backedge taken condition as !(++i == len)
• !(++i == len) is equivalent to !(++i >s (len - 1)), after checking len - 1 does not overflow (this is the new step you added)
• !(++i >s (len - 1)) is equivalent to (++i s<= (len - 1))
• (++i s<= (len - 1)) is equivalent to (++i s< len) after checking that (len - 1) + 1 does not overflow

Where does DecreasedRightValueByOne come in?

[mkazantsev]

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

818 ↗

(On Diff #105808)

The difference here is in LatchBrExitIdx. Here is a tricky moment. What we have here is

loop:
  %cond = icmp eq i32 %iv.next, %len
  br i1 %cond, label %exit, label %loop

and

loop:
  %cond = icmp ne i32 %iv.next, %len
  br i1 %cond, label %loop, label %exit

We could transform them both into

loop:
  %cond = icmp le i32 %iv.next, %len
  br i1 %cond, label %loop, label %exit

And set LatchBrExitIdx to 1. But LatchBrExitIdx is then stored into a structure and used later. It means that what we should not just change value of LatchBrExitIdx, but also change actual successors of branch instruction. And we should either remember that we did this interchange and swap them in the end of the method, or swap them just after this replacement (risking to have CFG changes before we checked that the gates pass).

This approach looks more complex for understanding, and also looks more error-prone. That's why I separate these cases here.

853 ↗

(On Diff #105808)

In case 2. After we go from (++i == length) to (++i >s (len - 1)), the new RighValue we are working with is actually len - 1. We don't need to create such an instruction, though. RightValue is only used to calculate RightValue + 1. Knowing that we have decreased the old RightValue by 1, we could just not add this 1 and use the old value of len instead.

[mkazantsev]

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

818 ↗

(On Diff #105808)

should be lt, not le

[sanjoy]

lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

853 ↗

(On Diff #105808)

I see -- do you mind (when you have time) refactoring this code a bit to make things like this more clear? For instance, it is confusing that we have RightValue and RightSCEV that are not in sync at certain points in the control flow.

One way to do this would be to have LoopStructure contain a "virtual" version of the latch branch (i.e. ICmpInst::Predicate Pred, const SCEV *LHS, const SCEV *RHS, BasicBlock *TrueBranch, BasicBlock *FalseBranch) and always rewrite the latch branch to be this "canonical" version during IRCE (after the gating checks have passed).

[mkazantsev]

Yes, I will try to make it more clear once I have some time for that. Thanks!

llvm/trunk/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

LoopStructure::parseLoopStructure(ScalarEvolution &SE, BranchProbabilityInfo &BPI,
 
   const SCEV *StartNext = IndVarNext->getStart();
   const SCEV *Addend = SE.getNegativeSCEV(IndVarNext->getStepRecurrence(SE));
   const SCEV *IndVarStart = SE.getAddExpr(StartNext, Addend);
 
   ConstantInt *One = ConstantInt::get(IndVarTy, 1);
   // TODO: generalize the predicates here to also match their unsigned variants.
   if (IsIncreasing) {
+    bool DecreasedRightValueByOne = false;
+    // Try to turn eq/ne predicates to those we can work with.
+    if (Pred == ICmpInst::ICMP_NE && LatchBrExitIdx == 1)
+      // while (++i != len) {         while (++i < len) {
+      //   ...                 --->     ...
+      // }                            }
+      Pred = ICmpInst::ICMP_SLT;
+    else if (Pred == ICmpInst::ICMP_EQ && LatchBrExitIdx == 0 &&
+             !CanBeSMin(SE, RightSCEV)) {
+      // while (true) {               while (true) {
+      //   if (++i == len)     --->     if (++i > len - 1)
+      //     break;                       break;
+      //   ...                          ...
+      // }                            }
+      Pred = ICmpInst::ICMP_SGT;
+      RightSCEV = SE.getMinusSCEV(RightSCEV, SE.getOne(RightSCEV->getType()));
+      DecreasedRightValueByOne = true;
+    }
+
     bool FoundExpectedPred =
         (Pred == ICmpInst::ICMP_SLT && LatchBrExitIdx == 1) ||
         (Pred == ICmpInst::ICMP_SGT && LatchBrExitIdx == 0);
 
     if (!FoundExpectedPred) {
       FailureReason = "expected icmp slt semantically, found something else";
       return None;
     }
 
     if (LatchBrExitIdx == 0) {
       if (CanBeSMax(SE, RightSCEV)) {
         // TODO: this restriction is easily removable -- we just have to
         // remember that the icmp was an slt and not an sle.
         FailureReason = "limit may overflow when coercing sle to slt";
         return None;
       }
 
       if (!SE.isLoopEntryGuardedByCond(
               &L, CmpInst::ICMP_SLT, IndVarStart,
               SE.getAddExpr(RightSCEV, SE.getOne(RightSCEV->getType())))) {
         FailureReason = "Induction variable start not bounded by upper limit";
         return None;
       }
 
+      // We need to increase the right value unless we have already decreased
+      // it virtually when we replaced EQ with SGT.
+      if (!DecreasedRightValueByOne) {
         IRBuilder<> B(Preheader->getTerminator());
         RightValue = B.CreateAdd(RightValue, One);
+      }
     } else {
       if (!SE.isLoopEntryGuardedByCond(&L, CmpInst::ICMP_SLT, IndVarStart,
                                        RightSCEV)) {
         FailureReason = "Induction variable start not bounded by upper limit";
         return None;
       }
+      assert(!DecreasedRightValueByOne &&
+             "Right value can be decreased only for LatchBrExitIdx == 0!");
     }
   } else {
+    bool IncreasedRightValueByOne = false;
+    // Try to turn eq/ne predicates to those we can work with.
+    if (Pred == ICmpInst::ICMP_NE && LatchBrExitIdx == 1)
+      // while (--i != len) {         while (--i > len) {
+      //   ...                 --->     ...
+      // }                            }
+      Pred = ICmpInst::ICMP_SGT;
+    else if (Pred == ICmpInst::ICMP_EQ && LatchBrExitIdx == 0 &&
+             !CanBeSMax(SE, RightSCEV)) {
+      // while (true) {               while (true) {
+      //   if (--i == len)     --->     if (--i < len + 1)
+      //     break;                       break;
+      //   ...                          ...
+      // }                            }
+      Pred = ICmpInst::ICMP_SLT;
+      RightSCEV = SE.getAddExpr(RightSCEV, SE.getOne(RightSCEV->getType()));
+      IncreasedRightValueByOne = true;
+    }
+
     bool FoundExpectedPred =
         (Pred == ICmpInst::ICMP_SGT && LatchBrExitIdx == 1) ||
         (Pred == ICmpInst::ICMP_SLT && LatchBrExitIdx == 0);
 
     if (!FoundExpectedPred) {
       FailureReason = "expected icmp sgt semantically, found something else";
       return None;
     }
 
     if (LatchBrExitIdx == 0) {
       if (CanBeSMin(SE, RightSCEV)) {
         // TODO: this restriction is easily removable -- we just have to
         // remember that the icmp was an sgt and not an sge.
         FailureReason = "limit may overflow when coercing sge to sgt";
         return None;
       }
 
       if (!SE.isLoopEntryGuardedByCond(
               &L, CmpInst::ICMP_SGT, IndVarStart,
               SE.getMinusSCEV(RightSCEV, SE.getOne(RightSCEV->getType())))) {
         FailureReason = "Induction variable start not bounded by lower limit";
         return None;
       }
 
+      // We need to decrease the right value unless we have already increased
+      // it virtually when we replaced EQ with SLT.
+      if (!IncreasedRightValueByOne) {
         IRBuilder<> B(Preheader->getTerminator());
         RightValue = B.CreateSub(RightValue, One);
+      }
     } else {
       if (!SE.isLoopEntryGuardedByCond(&L, CmpInst::ICMP_SGT, IndVarStart,
                                        RightSCEV)) {
         FailureReason = "Induction variable start not bounded by lower limit";
         return None;
       }
+      assert(!IncreasedRightValueByOne &&
+             "Right value can be increased only for LatchBrExitIdx == 0!");
     }
   }
 
   BasicBlock *LatchExit = LatchBr->getSuccessor(LatchBrExitIdx);
 
   assert(SE.getLoopDisposition(LatchCount, &L) ==
              ScalarEvolution::LoopInvariant &&
          "loop variant exit count doesn't make sense!");

llvm/trunk/test/Transforms/IRCE/eq_ne.ll

+; RUN: opt -verify-loop-info -irce-print-changed-loops -irce -S < %s 2>&1 | FileCheck %s
+
+; CHECK: irce: in function test_01: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK-NOT: irce: in function test_02: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK: irce: in function test_03: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK-NOT: irce: in function test_04: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK: irce: in function test_05: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK-NOT: irce: in function test_06: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK: irce: in function test_07: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+; CHECK-NOT: irce: in function test_08: constrained Loop at depth 1 containing: %loop<header><exiting>,%in.bounds<latch><exiting>
+
+; Show that IRCE can turn 'ne' condition to 'slt' in increasing IV.
+define void @test_01(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK:      test_01
+; CHECK:        main.exit.selector:
+; CHECK-NEXT:     [[PSEUDO_PHI:%[^ ]+]] = phi i32 [ %idx.next, %in.bounds ]
+; CHECK-NEXT:     [[COND:%[^ ]+]] = icmp slt i32 [[PSEUDO_PHI]], 100
+; CHECK-NEXT:     br i1 [[COND]]
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp ne i32 %idx.next, 100
+  br i1 %next, label %loop, label %exit
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that if n is not known to be greater than the starting value, IRCE
+; doesn't apply.
+define void @test_02(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_02(
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp ne i32 %idx.next, -100
+  br i1 %next, label %loop, label %exit
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that IRCE can turn 'eq' condition to 'sge' in increasing IV.
+define void @test_03(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_03(
+; CHECK:        main.exit.selector:
+; CHECK-NEXT:     [[PSEUDO_PHI:%[^ ]+]] = phi i32 [ %idx.next, %in.bounds ]
+; CHECK-NEXT:     [[COND:%[^ ]+]] = icmp slt i32 [[PSEUDO_PHI]], 100
+; CHECK-NEXT:     br i1 [[COND]]
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp eq i32 %idx.next, 100
+  br i1 %next, label %exit, label %loop
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that if n is not known to be greater than the starting value, IRCE
+; doesn't apply.
+define void @test_04(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_04(
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp eq i32 %idx.next, -100
+  br i1 %next, label %exit, label %loop
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that IRCE can turn 'ne' condition to 'sgt' in decreasing IV.
+define void @test_05(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_05(
+; CHECK:        preloop.exit.selector:
+; CHECK-NEXT:     [[PSEUDO_PHI:%[^ ]+]] = phi i32 [ %idx.next.preloop, %in.bounds.preloop ]
+; CHECK-NEXT:     [[COND:%[^ ]+]] = icmp sgt i32 [[PSEUDO_PHI]], 0
+; CHECK-NEXT:     br i1 [[COND]]
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+loop:
+  %idx = phi i32 [ 100, %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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp ne i32 %idx.next, 0
+  br i1 %next, label %loop, label %exit
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that IRCE cannot turn 'ne' condition to 'sgt' in decreasing IV if the end
+; value is not proved to be less than the start value.
+define void @test_06(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_06(
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+loop:
+  %idx = phi i32 [ 100, %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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp ne i32 %idx.next, 120
+  br i1 %next, label %loop, label %exit
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that IRCE can turn 'eq' condition to 'slt' in decreasing IV.
+define void @test_07(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_07(
+; CHECK:        preloop.exit.selector:
+; CHECK-NEXT:     [[PSEUDO_PHI:%[^ ]+]] = phi i32 [ %idx.next.preloop, %in.bounds.preloop ]
+; CHECK-NEXT:     [[COND:%[^ ]+]] = icmp sgt i32 [[PSEUDO_PHI]], 0
+; CHECK-NEXT:     br i1 [[COND]]
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+loop:
+  %idx = phi i32 [ 100, %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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp eq i32 %idx.next, 0
+  br i1 %next, label %exit, label %loop
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+; Show that IRCE cannot turn 'eq' condition to 'slt' in decreasing IV if the end
+; value is not proved to be less than the start value.
+define void @test_08(i32* %arr, i32* %a_len_ptr) #0 {
+
+; CHECK: test_08(
+
+entry:
+  %len = load i32, i32* %a_len_ptr, !range !0
+  br label %loop
+
+loop:
+  %idx = phi i32 [ 100, %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, !prof !1
+
+in.bounds:
+  %addr = getelementptr i32, i32* %arr, i32 %idx
+  store i32 0, i32* %addr
+  %next = icmp eq i32 %idx.next, 120
+  br i1 %next, label %exit, label %loop
+
+out.of.bounds:
+  ret void
+
+exit:
+  ret void
+}
+
+!0 = !{i32 0, i32 50}
+!1 = !{!"branch_weights", i32 64, i32 4}