[SimplifyIndVar] Canonicalize comparisons to unsigned while eliminating truncs

[mkazantsev]

Summary

This is a follow-up for the patch rL335020. When we replace compares against
trunc with compares against wide IV, we can also replace signed predicates with
unsigned where it is legal.

[greened]

lib/Transforms/Utils/SimplifyIndVar.cpp

558

That name seems misleading. It's not checking if they match, it's checking if the higest bit is zero for both. Did you mean to add a check for both being negative as well? The comments below seem to indicate so.

[greened]

lib/Transforms/Utils/SimplifyIndVar.cpp

558

Why does this need to be a lambda?

563

Does this need to be a lambda? Seems unnecessarily complicated.

[mkazantsev]

lib/Transforms/Utils/SimplifyIndVar.cpp

558

Yes, actually in the initial version it also checked negative, but then I realized that for negative values zext(trunc(x)) is never equal to x, so in practice this case will never fire. I will change the naming here.

563

It is the simplest way to avoid something like if (Cond1 || (!Cond2 && Cond3 && Cond4)) which is harder to understand.

[mkazantsev]

lib/Transforms/Utils/SimplifyIndVar.cpp

558

It doesn't, I'll merge it with the next one.

[mkazantsev]

Added some comments, removed one unnecessary lambda.

[mkazantsev]

Ping.

[greened]

LGTM.

lib/Transforms/Utils/SimplifyIndVar.cpp

563

Actually, I was wondering if it should be a freestanding function, but thinking about it now it seems like a good idea to define it where it's used. It's acting like a nested function in that sense, so I think this is fine. Thanks.

[reames]

Marking as accepted since LGTMed by previous reviewer. Just getting it off my queue.

lib/Transforms/Utils/SimplifyIndVar.cpp

if (ICmpInst *ICI = dyn_cast<ICmpInst>(U)) {
         // For equality, either signed or unsigned works.
         ICmpUsers.push_back(ICI);
       } else
         return false;
     } else
       return false;
   }
 
+  auto HighestBitsMatch = [&](Value *X, Value *Y) {
+    const SCEV *SCEVX = SE->getSCEV(X);
+    const SCEV *SCEVY = SE->getSCEV(Y);
+    return SE->isKnownNonNegative(SCEVX) && SE->isKnownNonNegative(SCEVY);
+  };
+  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.
+    return HighestBitsMatch(ICI->getOperand(0), ICI->getOperand(1));
+  };
   // Replace all comparisons against trunc with comparisons against IV.
   for (auto *ICI : ICmpUsers) {
     auto *Op1 = 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.
-    if (ICI->isUnsigned() || (ICI->isEquality() && DoesZExtCollapse)) {
+    ICmpInst::Predicate Pred = ICI->getPredicate();
+    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!");
     }
     bool Changed;
     L->makeLoopInvariant(Ext, Changed);
     (void)Changed;
-    ICmpInst *NewICI = new ICmpInst(ICI, ICI->getPredicate(), IV, Ext);
+    ICmpInst *NewICI = new ICmpInst(ICI, Pred, IV, Ext);
     ICI->replaceAllUsesWith(NewICI);
     DeadInsts.emplace_back(ICI);
   }
 
   // Trunc no longer needed.
   TI->replaceAllUsesWith(UndefValue::get(TI->getType()));
   DeadInsts.emplace_back(TI);
   return true;

test/Transforms/IndVarSimplify/eliminate-trunc.ll

 
 bb6:                                              ; preds = %bb4, %bb1
   %_tmp32 = add nuw nsw i32 %e.5.0, 1
   br i1 false, label %bb1, label %bb7
 
 bb7:                                             ; preds = %bb6
   ret void
 }
+
+; Show that we can turn signed comparison to unsigned and use zext while
+; comparing non-negative values.
+define void @test_12(i32* %p) {
+; CHECK-LABEL: @test_12(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[N:%.*]] = load i32, i32* [[P:%.*]], !range !0
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i32 [[N]] 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]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i64 [[IV_NEXT]], [[ZEXT]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %n = load i32, i32* %p, !range !0
+  br label %loop
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %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
+}
+
+!0 = !{i32 0, i32 1000}