Generalize icmp matching in IndVars' eliminateTrunc

[reames]

Summary

We were only matching RHS being a loop invariant value, not the inverse. Since there's nothing which appears to canonicalize loop invariant values to RHS, this means we missed cases.

Note: There's partial overlap between this transform and LFTR. That may worthy some deeper thinking, but for the moment, let's leave the design as is, and just make it work properly.

lib/Transforms/Utils/SimplifyIndVar.cpp

bool SimplifyIndvar::eliminateTrunc(TruncInst *TI) {
   // Collect users of the trunc that look like comparisons against invariants.
   // Bail if we find something different.
   SmallVector<ICmpInst *, 4> ICmpUsers;
   for (auto *U : TI->users()) {
     // We don't care about users in unreachable blocks.
     if (isa<Instruction>(U) &&
         !DT->isReachableFromEntry(cast<Instruction>(U)->getParent()))
       continue;
-    if (ICmpInst *ICI = dyn_cast<ICmpInst>(U)) {
-      if (ICI->getOperand(0) == TI && L->isLoopInvariant(ICI->getOperand(1))) {
-        assert(L->contains(ICI->getParent()) && "LCSSA form broken?");
-        // If we cannot get rid of trunc, bail.
-        if (ICI->isSigned() && !DoesSExtCollapse)
-          return false;
-        if (ICI->isUnsigned() && !DoesZExtCollapse)
-          return false;
-        // For equality, either signed or unsigned works.
-        ICmpUsers.push_back(ICI);
-      } else
-        return false;
-    } else
-      return false;
+    ICmpInst *ICI = dyn_cast<ICmpInst>(U);
+    if (!ICI) return false;
+    if (!(ICI->getOperand(0) == TI && L->isLoopInvariant(ICI->getOperand(1))) &&
+        !(ICI->getOperand(1) == TI && L->isLoopInvariant(ICI->getOperand(0))))
+      return false;
   }
 
   auto CanUseZExt = [&](ICmpInst *ICI) {
     // Unsigned comparison can be widened as unsigned.
     if (ICI->isUnsigned())
       return true;
     // Is it profitable to do zext?
     if (!DoesZExtCollapse)
       return false;
     // For equality, we can safely zext both parts.
     if (ICI->isEquality())
       return true;
     // Otherwise we can only use zext when comparing two non-negative or two
     // negative values. But in practice, we will never pass DoesZExtCollapse
     // check for a negative value, because zext(trunc(x)) is non-negative. So
     // it only make sense to check for non-negativity here.
     const SCEV *SCEVOP1 = SE->getSCEV(ICI->getOperand(0));
     const SCEV *SCEVOP2 = SE->getSCEV(ICI->getOperand(1));
     return SE->isKnownNonNegative(SCEVOP1) && SE->isKnownNonNegative(SCEVOP2);
   };
   // Replace all comparisons against trunc with comparisons against IV.
   for (auto *ICI : ICmpUsers) {
-    auto *Op1 = ICI->getOperand(1);
+    bool IsSwapped = L->isLoopInvariant(ICI->getOperand(0));
+    auto *Op1 = IsSwapped ? ICI->getOperand(0) : ICI->getOperand(1);
     Instruction *Ext = nullptr;
     // For signed/unsigned predicate, replace the old comparison with comparison
     // of immediate IV against sext/zext of the invariant argument. If we can
     // use either sext or zext (i.e. we are dealing with equality predicate),
     // then prefer zext as a more canonical form.
     // TODO: If we see a signed comparison which can be turned into unsigned,
     // we can do it here for canonicalization purposes.
     ICmpInst::Predicate Pred = ICI->getPredicate();
+    if (IsSwapped) Pred = ICmpInst::getSwappedPredicate(Pred);
     if (CanUseZExt(ICI)) {
       assert(DoesZExtCollapse && "Unprofitable zext?");
       Ext = new ZExtInst(Op1, IVTy, "zext", ICI);
       Pred = ICmpInst::getUnsignedPredicate(Pred);
     } else {
       assert(DoesSExtCollapse && "Unprofitable sext?");
       Ext = new SExtInst(Op1, IVTy, "sext", ICI);
       assert(Pred == ICmpInst::getSignedPredicate(Pred) && "Must be signed!");

test/Transforms/IndVarSimplify/eliminate-trunc.ll

loop:
   %iv.next = add i64 %iv, 1
   %narrow.iv = trunc i64 %iv.next to i32
   %cmp = icmp slt i32 %narrow.iv, %n
   br i1 %cmp, label %loop, label %exit
 exit:
   ret void
 }
 
+define void @test_13a(i32 %n) {
+;
+; CHECK-LABEL: @test_13a(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 1024 to i64
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ne i64 [[IV]], [[ZEXT]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %iv.next = add nsw nuw i64 %iv, 2
+  %narrow.iv = trunc i64 %iv to i32
+  %cmp = icmp ne i32 1024, %narrow.iv
+  br i1 %cmp, label %loop, label %exit
+exit:
+  ret void
+}
+
+define void @test_13b(i32 %n) {
+;
+; CHECK-LABEL: @test_13b(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 1024 to i64
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i64 [[IV]], [[ZEXT]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %iv.next = add nsw nuw i64 %iv, 2
+  %narrow.iv = trunc i64 %iv to i32
+  %cmp = icmp ugt i32 1024, %narrow.iv
+  br i1 %cmp, label %loop, label %exit
+exit:
+  ret void
+}
+
+define void @test_13c(i32 %n) {
+;
+; CHECK-LABEL: @test_13c(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 1024 to i64
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i64 [[IV]], [[ZEXT]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %iv.next = add nsw nuw i64 %iv, 2
+  %narrow.iv = trunc i64 %iv to i32
+  %cmp = icmp sgt i32 1024, %narrow.iv
+  br i1 %cmp, label %loop, label %exit
+exit:
+  ret void
+}
+
+define void @test_13d(i32 %n) {
+;
+; CHECK-LABEL: @test_13d(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i32 1024 to i64
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ -20, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nsw i64 [[IV]], 2
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp slt i64 [[IV]], [[SEXT]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %iv = phi i64 [ -20, %entry ], [ %iv.next, %loop ]
+  %iv.next = add nsw i64 %iv, 2
+  %narrow.iv = trunc i64 %iv to i32
+  %cmp = icmp sgt i32 1024, %narrow.iv
+  br i1 %cmp, label %loop, label %exit
+exit:
+  ret void
+}
+
 !0 = !{i32 0, i32 1000}