[LoopPredication] Handle the case when the guard and the latch IV have different offsets

[apilipenko]

Summary

This is a follow up change for D37569.

Currently the transformation is limited to the case when:

• The loop has a single latch with the condition of the form: ++i <pred> latchLimit, where <pred> is u<, u<=, s<, or s<=.
• The step of the IV used in the latch condition is 1.
• The IV of the latch condition is the same as the post increment IV of the guard condition.
• The guard condition is of the form i u< guardLimit.

This patch enables the transform in the case when the latch is

latchStart + i <pred> latchLimit, where <pred> is u<, u<=, s<, or s<=.

And the guard is

guardStart + i u< guardLimit

This is the last change in this chain.

[anna]

@apilipenko please add patch with full context.

[apilipenko]

Upload diff with context. Minor fix in the header comment.

[anna]

Artur, before this patch, we were checking whether RangeCheckIV's postIncExpr is latchCheckIV.

Now, the header comments and proof seems to point to the fact that the guard and latch are using the same offset (lines 103 - 110), but the code does not have any check that we are using the same offset. Or, is it that the offset is the same, but the difference in the latchCheckIV versus rangeCheckIV is caught in their respective start values, and the offset for both latch and guard remains the same?

Pls see inline comment for more clarification and specific example.

lib/Transforms/Scalar/LoopPredication.cpp

61

Now that arbitrary offsets are supported for both guard and latch, the guard can even be guard(G(I.INC)), while the backedge condition can be B(I)?
In other words, do we need the same offset 'i' for both guard and latch?

Please add a testcase when the guard uses the I.INC and backedge uses I.

[apilipenko]

Clarified the description, added the requested test case.

[apilipenko]

In D39097#905174, @anna wrote:

Artur, before this patch, we were checking whether RangeCheckIV's postIncExpr is latchCheckIV.

Now, the header comments and proof seems to point to the fact that the guard and latch are using the same offset (lines 103 - 110), but the code does not have any check that we are using the same offset. Or, is it that the offset is the same, but the difference in the latchCheckIV versus rangeCheckIV is caught in their respective start values, and the offset for both latch and guard remains the same?

I changed the loop in the description so it doesn't use Start. Instead the offsets are sunk into predicates G and B.

I also explicitly specified the predicates for the supported cases.

lib/Transforms/Scalar/LoopPredication.cpp

61

No, we don't. I changed the loop so to make this clear, added the requested test case.

[anna]

LGTM.

lib/Transforms/Scalar/LoopPredication.cpp

56

I'm not sure if this is clearer :)
My concern was that this loop makes it look like we support only the case where B and G use the same value of I.

It becomes clearer when we define what B and G predicates are.

103

The key point here is that latchStart and guardStart can be different values.
So, B(X) can be X, while G(X) is X+1. Various different combinations are allowed.

lib/Transforms/Scalar/LoopPredication.cpp

 //
 // For now the transformation is limited to the following case:
 //   * The loop has a single latch with the condition of the form:
-//      ++i <pred> latchLimit, where <pred> is u<, u<=, s<, or s<=.
+//      latchStart + i <pred> latchLimit, where <pred> is u<, u<=, s<, or s<=.
 //   * The step of the IV used in the latch condition is 1.
-//   * The IV of the latch condition is the same as the post increment IV of the
-//   guard condition.
-//   * The guard condition is
-//     i u< guardLimit.
+//   * The guard condition is of the form
+//     guardStart + i u< guardLimit
 //
 // For the ult latch comparison case M is:
-//   forall X . X u< guardLimit && (X + 1) u< latchLimit =>
-//      (X + 1) u< guardLimit
+//   forall X . guardStart + X u< guardLimit && latchStart + X <u latchLimit =>
+//      guardStart + X + 1 u< guardLimit
 //
-// This is true if latchLimit u<= guardLimit since then
-//   (X + 1) u< latchLimit u<= guardLimit == (X + 1) u< guardLimit.
-//
-// So for ult condition the widened condition is:
-//   i.start u< guardLimit && latchLimit u<= guardLimit
-// Similarly for ule condition the widened condition is:
-//   i.start u< guardLimit && latchLimit u< guardLimit
-//
-// For the signed latch comparison case M is:
-//   forall X . X u< guardLimit && (X + 1) s< latchLimit =>
-//      (X + 1) u< guardLimit
-//
-// The only way the antecedent can be true and the consequent can be false is 
+// The only way the antecedent can be true and the consequent can be false is
 // if
-//   X == guardLimit - 1
+//   X == guardLimit - 1 - guardStart
 // (and guardLimit is non-zero, but we won't use this latter fact).
-// If X == guardLimit - 1 then the second half of the antecedent is
-//   guardLimit s< latchLimit
+// If X == guardLimit - 1 - guardStart then the second half of the antecedent is
+//   latchStart + guardLimit - 1 - guardStart u< latchLimit
 // and its negation is
-//   latchLimit s<= guardLimit.
+//   latchStart + guardLimit - 1 - guardStart u>= latchLimit
 //
-// In other words, if latchLimit s<= guardLimit then:
+// In other words, if
+//   latchLimit u<= latchStart + guardLimit - 1 - guardStart
+// then:
 // (the ranges below are written in ConstantRange notation, where [A, B) is the
 // set for (I = A; I != B; I++ /*maywrap*/) yield(I);)
 //
-//    forall X . X u< guardLimit && (X + 1) s< latchLimit =>  (X + 1) u< guardLimit
-// == forall X . X u< guardLimit && (X + 1) s< guardLimit =>  (X + 1) u< guardLimit
-// == forall X . X in [0, guardLimit) && (X + 1) in [INT_MIN, guardLimit) =>  (X + 1) in [0, guardLimit)
-// == forall X . X in [0, guardLimit) && X in [INT_MAX, guardLimit-1) =>  X in [-1, guardLimit-1)
-// == forall X . X in [0, guardLimit-1) => X in [-1, guardLimit-1)
+//    forall X . guardStart + X u< guardLimit &&
+//               latchStart + X u< latchLimit =>
+//      guardStart + X + 1 u< guardLimit
+// == forall X . guardStart + X u< guardLimit &&
+//               latchStart + X u< latchStart + guardLimit - 1 - guardStart =>
+//      guardStart + X + 1 u< guardLimit
+// == forall X . (guardStart + X) in [0, guardLimit) &&
+//               (latchStart + X) in [0, latchStart + guardLimit - 1 - guardStart) =>
+//      (guardStart + X + 1) in [0, guardLimit)
+// == forall X . X in [-guardStart, guardLimit - guardStart) &&
+//               X in [-latchStart, guardLimit - 1 - guardStart) =>
+//       X in [-guardStart - 1, guardLimit - guardStart - 1)
 // == true
 //
 // So the widened condition is:
-//   i.start u< guardLimit && latchLimit s<= guardLimit
-// Similarly for sle condition the widened condition is:
-//   i.start u< guardLimit && latchLimit s< guardLimit
+//   guardStart u< guardLimit &&
+//   latchStart + guardLimit - 1 - guardStart u>= latchLimit
+// Similarly for ule condition the widened condition is:
+//   guardStart u< guardLimit &&
+//   latchStart + guardLimit - 1 - guardStart u> latchLimit
+// For slt condition the widened condition is:
+//   guardStart u< guardLimit &&
+//   latchStart + guardLimit - 1 - guardStart s>= latchLimit
+// For sle condition the widened condition is:
+//   guardStart u< guardLimit &&
+//   latchStart + guardLimit - 1 - guardStart s> latchLimit
 //
 //===----------------------------------------------------------------------===//
 

     return None;
   }
   auto *RangeCheckIV = RangeCheck->IV;
-  auto *PostIncRangeCheckIV = RangeCheckIV->getPostIncExpr(*SE);
-  if (LatchCheck.IV != PostIncRangeCheckIV) {
-    DEBUG(dbgs() << "Post increment range check IV (" << *PostIncRangeCheckIV
-                 << ") is not the same as latch IV (" << *LatchCheck.IV
-                 << ")!\n");
+  auto *Ty = RangeCheckIV->getType();
+  if (Ty != LatchCheck.IV->getType()) {
+    DEBUG(dbgs() << "Type mismatch between range check and latch IVs!\n");
+    return None;
+  }
+  if (!RangeCheckIV->isAffine()) {
+    DEBUG(dbgs() << "Range check IV is not affine!\n");
+    return None;
+  }
+  auto *Step = RangeCheckIV->getStepRecurrence(*SE);
+  if (Step != LatchCheck.IV->getStepRecurrence(*SE)) {
+    DEBUG(dbgs() << "Range check and latch have IVs different steps!\n");
     return None;
   }
-  assert(RangeCheckIV->getStepRecurrence(*SE)->isOne() && "must be one");
-  const SCEV *Start = RangeCheckIV->getStart();
+  assert(Step->isOne() && "must be one");
 
   // Generate the widened condition:
-  //   i.start u< guardLimit && latchLimit <pred> guardLimit
+  //   guardStart u< guardLimit &&
+  //   latchLimit <pred> guardLimit - 1 - guardStart + latchStart
   // where <pred> depends on the latch condition predicate. See the file
   // header comment for the reasoning.
+  const SCEV *GuardStart = RangeCheckIV->getStart();
+  const SCEV *GuardLimit = RangeCheck->Limit;
+  const SCEV *LatchStart = LatchCheck.IV->getStart();
+  const SCEV *LatchLimit = LatchCheck.Limit;
+
+  // guardLimit - guardStart + latchStart - 1
+  const SCEV *RHS =
+      SE->getAddExpr(SE->getMinusSCEV(GuardLimit, GuardStart),
+                     SE->getMinusSCEV(LatchStart, SE->getOne(Ty)));
+
   ICmpInst::Predicate LimitCheckPred;
   switch (LatchCheck.Pred) {
   case ICmpInst::ICMP_ULT:

     llvm_unreachable("Unsupported loop latch!");
   }
 
+  DEBUG(dbgs() << "LHS: " << *LatchLimit << "\n");
+  DEBUG(dbgs() << "RHS: " << *RHS << "\n");
+  DEBUG(dbgs() << "Pred: " << LimitCheckPred << "\n");
+
   auto CanExpand = [this](const SCEV *S) {
     return SE->isLoopInvariant(S, L) && isSafeToExpand(S, *SE);
   };
-  if (!CanExpand(Start) || !CanExpand(LatchCheck.Limit) ||
-      !CanExpand(RangeCheck->Limit))
+  if (!CanExpand(GuardStart) || !CanExpand(GuardLimit) ||
+      !CanExpand(LatchLimit) || !CanExpand(RHS)) {
+    DEBUG(dbgs() << "Can't expand limit check!\n");
     return None;
+  }
 
   Instruction *InsertAt = Preheader->getTerminator();
-  auto *LimitCheck = expandCheck(Expander, Builder, LimitCheckPred,
-                                 LatchCheck.Limit, RangeCheck->Limit, InsertAt);
+  auto *LimitCheck =
+      expandCheck(Expander, Builder, LimitCheckPred, LatchLimit, RHS, InsertAt);
   auto *FirstIterationCheck = expandCheck(Expander, Builder, RangeCheck->Pred,
-                                          Start, RangeCheck->Limit, InsertAt);
+                                          GuardStart, GuardLimit, InsertAt);
   return Builder.CreateAnd(FirstIterationCheck, LimitCheck);
 }
 

test/Transforms/LoopPredication/basic.ll

   ret i32 %result
 }
 
+define i32 @signed_loop_0_to_n_preincrement_latch_check(i32* %array, i32 %length, i32 %n) {
+; CHECK-LABEL: @signed_loop_0_to_n_preincrement_latch_check
+entry:
+  %tmp5 = icmp sle i32 %n, 0
+  br i1 %tmp5, label %exit, label %loop.preheader
+
+loop.preheader:
+; CHECK: loop.preheader:
+; CHECK: [[length_minus_1:[^ ]+]] = add i32 %length, -1
+; CHECK-NEXT: [[limit_check:[^ ]+]] = icmp sle i32 %n, [[length_minus_1]]
+; CHECK-NEXT: [[first_iteration_check:[^ ]+]] = icmp ult i32 0, %length
+; CHECK-NEXT: [[wide_cond:[^ ]+]] = and i1 [[first_iteration_check]], [[limit_check]]
+; CHECK-NEXT: br label %loop
+  br label %loop
+
+loop:
+; CHECK: loop:
+; CHECK: call void (i1, ...) @llvm.experimental.guard(i1 [[wide_cond]], i32 9) [ "deopt"() ]
+  %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
+  %i = phi i32 [ %i.next, %loop ], [ 0, %loop.preheader ]
+  %within.bounds = icmp ult i32 %i, %length
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+
+  %i.i64 = zext i32 %i to i64
+  %array.i.ptr = getelementptr inbounds i32, i32* %array, i64 %i.i64
+  %array.i = load i32, i32* %array.i.ptr, align 4
+  %loop.acc.next = add i32 %loop.acc, %array.i
+
+  %i.next = add i32 %i, 1
+  %continue = icmp slt i32 %i, %n
+  br i1 %continue, label %loop, label %exit
+
+exit:
+  %result = phi i32 [ 0, %entry ], [ %loop.acc.next, %loop ]
+  ret i32 %result
+}
+
+define i32 @signed_loop_0_to_n_sle_latch_offset_ult_check(i32* %array, i32 %length, i32 %n) {
+; CHECK-LABEL: @signed_loop_0_to_n_sle_latch_offset_ult_check
+entry:
+  %tmp5 = icmp sle i32 %n, 0
+  br i1 %tmp5, label %exit, label %loop.preheader
+
+loop.preheader:
+; CHECK: loop.preheader:
+; CHECK: [[length_minus_1:[^ ]+]] = add i32 %length, -1
+; CHECK-NEXT: [[limit_check:[^ ]+]] = icmp slt i32 %n, [[length_minus_1]]
+; CHECK-NEXT: [[first_iteration_check:[^ ]+]] = icmp ult i32 1, %length
+; CHECK-NEXT: [[wide_cond:[^ ]+]] = and i1 [[first_iteration_check]], [[limit_check]]
+; CHECK-NEXT: br label %loop
+  br label %loop
+
+loop:
+; CHECK: loop:
+; CHECK: call void (i1, ...) @llvm.experimental.guard(i1 [[wide_cond]], i32 9) [ "deopt"() ]
+  %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
+  %i = phi i32 [ %i.next, %loop ], [ 0, %loop.preheader ]
+  %i.offset = add i32 %i, 1
+  %within.bounds = icmp ult i32 %i.offset, %length
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+
+  %i.i64 = zext i32 %i to i64
+  %array.i.ptr = getelementptr inbounds i32, i32* %array, i64 %i.i64
+  %array.i = load i32, i32* %array.i.ptr, align 4
+  %loop.acc.next = add i32 %loop.acc, %array.i
+
+  %i.next = add i32 %i, 1
+  %continue = icmp sle i32 %i.next, %n
+  br i1 %continue, label %loop, label %exit
+
+exit:
+  %result = phi i32 [ 0, %entry ], [ %loop.acc.next, %loop ]
+  ret i32 %result
+}
+
+define i32 @signed_loop_0_to_n_offset_sle_latch_offset_ult_check(i32* %array, i32 %length, i32 %n) {
+; CHECK-LABEL: @signed_loop_0_to_n_offset_sle_latch_offset_ult_check
+entry:
+  %tmp5 = icmp sle i32 %n, 0
+  br i1 %tmp5, label %exit, label %loop.preheader
+
+loop.preheader:
+; CHECK: loop.preheader:
+; CHECK: [[limit_check:[^ ]+]] = icmp slt i32 %n, %length
+; CHECK-NEXT: [[first_iteration_check:[^ ]+]] = icmp ult i32 1, %length
+; CHECK-NEXT: [[wide_cond:[^ ]+]] = and i1 [[first_iteration_check]], [[limit_check]]
+; CHECK-NEXT: br label %loop
+  br label %loop
+
+loop:
+; CHECK: loop:
+; CHECK: call void (i1, ...) @llvm.experimental.guard(i1 [[wide_cond]], i32 9) [ "deopt"() ]
+  %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
+  %i = phi i32 [ %i.next, %loop ], [ 0, %loop.preheader ]
+  %i.offset = add i32 %i, 1
+  %within.bounds = icmp ult i32 %i.offset, %length
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+
+  %i.i64 = zext i32 %i to i64
+  %array.i.ptr = getelementptr inbounds i32, i32* %array, i64 %i.i64
+  %array.i = load i32, i32* %array.i.ptr, align 4
+  %loop.acc.next = add i32 %loop.acc, %array.i
+
+  %i.next = add i32 %i, 1
+  %i.next.offset = add i32 %i.next, 1
+  %continue = icmp sle i32 %i.next.offset, %n
+  br i1 %continue, label %loop, label %exit
+
+exit:
+  %result = phi i32 [ 0, %entry ], [ %loop.acc.next, %loop ]
+  ret i32 %result
+}
+
 define i32 @unsupported_latch_pred_loop_0_to_n(i32* %array, i32 %length, i32 %n) {
 ; CHECK-LABEL: @unsupported_latch_pred_loop_0_to_n
 entry:

   ret i32 %result
 }
 
-define i32 @signed_loop_0_to_n_unrelated_iv_range_check(i32* %array, i32 %start, i32 %length, i32 %n) {
-; CHECK-LABEL: @signed_loop_0_to_n_unrelated_iv_range_check
+define i32 @signed_loop_start_to_n_offset_iv_range_check(i32* %array, i32 %start.i,
+                                                         i32 %start.j, i32 %length,
+                                                         i32 %n) {
+; CHECK-LABEL: @signed_loop_start_to_n_offset_iv_range_check
+entry:
+  %tmp5 = icmp sle i32 %n, 0
+  br i1 %tmp5, label %exit, label %loop.preheader
+
+loop.preheader:
+; CHECK: loop.preheader:
+; CHECK: [[length_plus_start_i:[^ ]+]] = add i32 %length, %start.i
+; CHECK-NEXT: [[limit:[^ ]+]] = sub i32 [[length_plus_start_i]], %start.j
+; CHECK-NEXT: [[limit_check:[^ ]+]] = icmp sle i32 %n, [[limit]]
+; CHECK-NEXT: [[first_iteration_check:[^ ]+]] = icmp ult i32 %start.j, %length
+; CHECK-NEXT: [[wide_cond:[^ ]+]] = and i1 [[first_iteration_check]], [[limit_check]]
+; CHECK-NEXT: br label %loop
+  br label %loop
+
+loop:
+; CHECK: loop:
+; CHECK: call void (i1, ...) @llvm.experimental.guard(i1 [[wide_cond]], i32 9) [ "deopt"() ]
+  %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
+  %i = phi i32 [ %i.next, %loop ], [ %start.i, %loop.preheader ]
+  %j = phi i32 [ %j.next, %loop ], [ %start.j, %loop.preheader ]
+
+  %within.bounds = icmp ult i32 %j, %length
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+
+  %i.i64 = zext i32 %i to i64
+  %array.i.ptr = getelementptr inbounds i32, i32* %array, i64 %i.i64
+  %array.i = load i32, i32* %array.i.ptr, align 4
+  %loop.acc.next = add i32 %loop.acc, %array.i
+
+  %j.next = add i32 %j, 1
+  %i.next = add i32 %i, 1
+  %continue = icmp slt i32 %i.next, %n
+  br i1 %continue, label %loop, label %exit
+
+exit:
+  %result = phi i32 [ 0, %entry ], [ %loop.acc.next, %loop ]
+  ret i32 %result
+}
+
+define i32 @signed_loop_0_to_n_different_iv_types(i32* %array, i16 %length, i32 %n) {
+; CHECK-LABEL: @signed_loop_0_to_n_different_iv_types
+entry:
+  %tmp5 = icmp sle i32 %n, 0
+  br i1 %tmp5, label %exit, label %loop.preheader
+
+loop.preheader:
+; CHECK: loop.preheader:
+; CHECK-NEXT: br label %loop
+  br label %loop
+
+loop:
+; CHECK: loop:
+; CHECK: %within.bounds = icmp ult i16 %j, %length
+; CHECK-NEXT: call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+  %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
+  %i = phi i32 [ %i.next, %loop ], [ 0, %loop.preheader ]
+  %j = phi i16 [ %j.next, %loop ], [ 0, %loop.preheader ]
+
+  %within.bounds = icmp ult i16 %j, %length
+  call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
+
+  %i.i64 = zext i32 %i to i64
+  %array.i.ptr = getelementptr inbounds i32, i32* %array, i64 %i.i64
+  %array.i = load i32, i32* %array.i.ptr, align 4
+  %loop.acc.next = add i32 %loop.acc, %array.i
+
+  %j.next = add i16 %j, 1
+  %i.next = add i32 %i, 1
+  %continue = icmp slt i32 %i.next, %n
+  br i1 %continue, label %loop, label %exit
+
+exit:
+  %result = phi i32 [ 0, %entry ], [ %loop.acc.next, %loop ]
+  ret i32 %result
+}
+
+define i32 @signed_loop_0_to_n_different_iv_strides(i32* %array, i32 %length, i32 %n) {
+; CHECK-LABEL: @signed_loop_0_to_n_different_iv_strides
 entry:
   %tmp5 = icmp sle i32 %n, 0
   br i1 %tmp5, label %exit, label %loop.preheader

 ; CHECK-NEXT: call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]
   %loop.acc = phi i32 [ %loop.acc.next, %loop ], [ 0, %loop.preheader ]
   %i = phi i32 [ %i.next, %loop ], [ 0, %loop.preheader ]
-  %j = phi i32 [ %j.next, %loop ], [ %start, %loop.preheader ]
+  %j = phi i32 [ %j.next, %loop ], [ 0, %loop.preheader ]
 
   %within.bounds = icmp ult i32 %j, %length
   call void (i1, ...) @llvm.experimental.guard(i1 %within.bounds, i32 9) [ "deopt"() ]

   %array.i = load i32, i32* %array.i.ptr, align 4
   %loop.acc.next = add i32 %loop.acc, %array.i
 
-  %j.next = add nsw i32 %j, 1
+  %j.next = add nsw i32 %j, 2
   %i.next = add nsw i32 %i, 1
   %continue = icmp slt i32 %i.next, %n
   br i1 %continue, label %loop, label %exit