[indvars] Extend canonicalizeExitConditions to inverted operands

As discussed in the original reviews, but done in a follow on.
llvm · Nov 4, 2021 · 453fdeb · 453fdeb
1 parent eb0fa8b
commit 453fdeb
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git a/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp b/llvm/lib/Transforms/Scalar/IndVarSimplify.cpp
@@ -1439,8 +1439,12 @@ bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
     // For the range reasoning, avoid computing SCEVs in the loop to avoid
     // poisoning cache with sub-optimal results.  For the must-execute case,
     // this is a neccessary precondition for correctness.
-    if (!L->isLoopInvariant(RHS))
-      continue;
+    if (!L->isLoopInvariant(RHS)) {
+      if (!L->isLoopInvariant(LHS))
+        continue;
+      // Same logic applies for the inverse case
+      std::swap(LHS, RHS);
+    }
 
     // Match (icmp signed-cond zext, RHS)
     Value *LHSOp = nullptr;
@@ -1475,11 +1479,19 @@ bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
     if (!ICmp || !ICmp->hasOneUse() || !ICmp->isUnsigned())
       continue;
 
+    bool Swapped = false;
     auto *LHS = ICmp->getOperand(0);
     auto *RHS = ICmp->getOperand(1);
-    if (L->isLoopInvariant(LHS) || !L->isLoopInvariant(RHS))
+    if (L->isLoopInvariant(LHS) == L->isLoopInvariant(RHS))
       // Nothing to rotate
       continue;
+    if (L->isLoopInvariant(LHS)) {
+      // Same logic applies for the inverse case until we actually pick
+      // which operand of the compare to update.
+      Swapped = true;
+      std::swap(LHS, RHS);
+    }
+    assert(!L->isLoopInvariant(LHS) && L->isLoopInvariant(RHS));
 
     if (!LHS->hasOneUse())
       // Can't rotate without increasing instruction count
@@ -1501,8 +1513,8 @@ bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
       auto *NewRHS =
         CastInst::Create(Instruction::Trunc, RHS, LHSOp->getType(), "",
                          L->getLoopPreheader()->getTerminator());
-      ICmp->setOperand(0, LHSOp);
-      ICmp->setOperand(1, NewRHS);
+      ICmp->setOperand(Swapped ? 1 : 0, LHSOp);
+      ICmp->setOperand(Swapped ? 0 : 1, NewRHS);
       if (LHS->use_empty())
         DeadInsts.push_back(LHS);
     };

diff --git a/llvm/test/Transforms/IndVarSimplify/finite-exit-comparisons.ll b/llvm/test/Transforms/IndVarSimplify/finite-exit-comparisons.ll
@@ -824,12 +824,12 @@ for.end:                                          ; preds = %for.body, %entry
 define void @slt_constant_lhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @slt_constant_lhs(
 ; 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 slt i16 254, [[ZEXT]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[TMP0]], [[IV_NEXT]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
@@ -878,12 +878,12 @@ for.end:                                          ; preds = %for.body, %entry
 define void @ult_constant_lhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @ult_constant_lhs(
 ; 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 i16 254, [[ZEXT]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[TMP0]], [[IV_NEXT]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void