[indvars] Allow rotation (narrowing) of exit test when discovering trip count

This relaxes the one-use requirement on the rotation transform specifically for the case where we know we're zexting an IV of the loop.  This allows us to discover trip count information in SCEV, which seems worth a single extra loop invariant truncate.  Honestly, I'd prefer if SCEV could just compute the trip count directly (e.g. D109457), but this unblocks practical benefit.

Author: preames

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

@@ -1493,17 +1493,22 @@ bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
     }
     assert(!L->isLoopInvariant(LHS) && L->isLoopInvariant(RHS));
 
-    if (!LHS->hasOneUse())
-      // Can't rotate without increasing instruction count
-      continue;
-
     // Match (icmp unsigned-cond zext, RHS)
     // TODO: Extend to handle corresponding sext/signed-cmp case
     // TODO: Extend to other invertible functions
     Value *LHSOp = nullptr;
     if (!match(LHS, m_ZExt(m_Value(LHSOp))))
       continue;
 
+    // In general, we only rotate if we can do so without increasing the number
+    // of instructions.  The exception is when we have an zext(add-rec).  The
+    // reason for allowing this exception is that we know we need to get rid
+    // of the zext for SCEV to be able to compute a trip count for said loops;
+    // we consider the new trip count valuable enough to increase instruction
+    // count by one.
+    if (!LHS->hasOneUse() && !isa<SCEVAddRecExpr>(SE->getSCEV(LHSOp)))
+      continue;
+
     // Given a icmp unsigned-cond zext(Op) where zext(trunc(RHS)) == RHS
     // replace with an icmp of the form icmp unsigned-cond Op, trunc(RHS)
     // when zext is loop varying and RHS is loop invariant.  This converts

llvm/test/Transforms/IndVarSimplify/finite-exit-comparisons.ll

@@ -928,3 +928,64 @@ for.body:                                         ; preds = %entry, %for.body
 for.end:                                          ; preds = %for.body, %entry
   ret void
 }
+
+define i16 @ult_multiuse_profit(i16 %n.raw, i8 %start) mustprogress {
+; CHECK-LABEL: @ult_multiuse_profit(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i16 254 to i8
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[IV_NEXT]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
+; CHECK:       for.end:
+; CHECK-NEXT:    [[ZEXT_LCSSA:%.*]] = phi i16 [ [[ZEXT]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i16 [[ZEXT_LCSSA]]
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %iv = phi i8 [ %iv.next, %for.body ], [ %start, %entry ]
+  %iv.next = add i8 %iv, 1
+  %zext = zext i8 %iv.next to i16
+  %cmp = icmp ult i16 %zext, 254
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:                                          ; preds = %for.body, %entry
+  ret i16 %zext
+}
+
+define i16 @ult_multiuse_profit2(i16 %n.raw, i8 %start) mustprogress {
+; CHECK-LABEL: @ult_multiuse_profit2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i16 254 to i8
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[IV2:%.*]] = phi i16 [ [[ZEXT:%.*]], [[FOR_BODY]] ], [ 0, [[ENTRY]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
+; CHECK-NEXT:    [[ZEXT]] = zext i8 [[IV_NEXT]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[IV_NEXT]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
+; CHECK:       for.end:
+; CHECK-NEXT:    [[IV2_LCSSA:%.*]] = phi i16 [ [[IV2]], [[FOR_BODY]] ]
+; CHECK-NEXT:    ret i16 [[IV2_LCSSA]]
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %iv = phi i8 [ %iv.next, %for.body ], [ %start, %entry ]
+  %iv2 = phi i16 [%zext, %for.body], [0, %entry]
+  %iv.next = add i8 %iv, 1
+  %zext = zext i8 %iv.next to i16
+  %cmp = icmp ult i16 %zext, 254
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:                                          ; preds = %for.body, %entry
+  ret i16 %iv2
+}
+