[InstCombine] Support gep nuw in icmp folds

Unsigned icmp of gep nuw folds to unsigned icmp of offsets.
Unsigned icmp of gep nusw nuw folds to unsigned samesign icmp of
offsets.

Proofs: https://alive2.llvm.org/ce/z/VEwQY8

Author: nikic

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -690,13 +690,32 @@ Instruction *InstCombinerImpl::foldGEPICmp(GEPOperator *GEPLHS, Value *RHS,
   if (!isa<GetElementPtrInst>(RHS))
     RHS = RHS->stripPointerCasts();
 
+  auto CanFold = [Cond](GEPNoWrapFlags NW) {
+    if (ICmpInst::isEquality(Cond))
+      return true;
+
+    // Unsigned predicates can be folded if the GEPs have *any* nowrap flags.
+    assert(ICmpInst::isUnsigned(Cond));
+    return NW != GEPNoWrapFlags::none();
+  };
+
+  auto NewICmp = [Cond](GEPNoWrapFlags NW, Value *Op1, Value *Op2) {
+    if (!NW.hasNoUnsignedWrap()) {
+      // Convert signed to unsigned comparison.
+      return new ICmpInst(ICmpInst::getSignedPredicate(Cond), Op1, Op2);
+    }
+
+    auto *I = new ICmpInst(Cond, Op1, Op2);
+    I->setSameSign(NW.hasNoUnsignedSignedWrap());
+    return I;
+  };
+
   Value *PtrBase = GEPLHS->getOperand(0);
-  if (PtrBase == RHS &&
-      (GEPLHS->hasNoUnsignedSignedWrap() || ICmpInst::isEquality(Cond))) {
+  if (PtrBase == RHS && CanFold(GEPLHS->getNoWrapFlags())) {
     // ((gep Ptr, OFFSET) cmp Ptr)   ---> (OFFSET cmp 0).
     Value *Offset = EmitGEPOffset(GEPLHS);
-    return new ICmpInst(ICmpInst::getSignedPredicate(Cond), Offset,
-                        Constant::getNullValue(Offset->getType()));
+    return NewICmp(GEPLHS->getNoWrapFlags(), Offset,
+                   Constant::getNullValue(Offset->getType()));
   }
 
   if (GEPLHS->isInBounds() && ICmpInst::isEquality(Cond) &&
@@ -814,19 +833,18 @@ Instruction *InstCombinerImpl::foldGEPICmp(GEPOperator *GEPLHS, Value *RHS,
         return replaceInstUsesWith(I, // No comparison is needed here.
           ConstantInt::get(I.getType(), ICmpInst::isTrueWhenEqual(Cond)));
 
-      else if (NumDifferences == 1 && NW.hasNoUnsignedSignedWrap()) {
+      else if (NumDifferences == 1 && CanFold(NW)) {
         Value *LHSV = GEPLHS->getOperand(DiffOperand);
         Value *RHSV = GEPRHS->getOperand(DiffOperand);
-        // Make sure we do a signed comparison here.
-        return new ICmpInst(ICmpInst::getSignedPredicate(Cond), LHSV, RHSV);
+        return NewICmp(NW, LHSV, RHSV);
       }
     }
 
-    if (NW.hasNoUnsignedSignedWrap() || CmpInst::isEquality(Cond)) {
+    if (CanFold(NW)) {
       // ((gep Ptr, OFFSET1) cmp (gep Ptr, OFFSET2)  --->  (OFFSET1 cmp OFFSET2)
       Value *L = EmitGEPOffset(GEPLHS, /*RewriteGEP=*/true);
       Value *R = EmitGEPOffset(GEPRHS, /*RewriteGEP=*/true);
-      return new ICmpInst(ICmpInst::getSignedPredicate(Cond), L, R);
+      return NewICmp(NW, L, R);
     }
   }
 

llvm/test/Transforms/InstCombine/icmp-gep.ll

@@ -143,6 +143,44 @@ define i1 @ult_base_nusw(ptr %x, i64 %y) {
   ret i1 %r
 }
 
+define i1 @ugt_base_nuw(ptr %x, i64 %y) {
+; CHECK-LABEL: @ugt_base_nuw(
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i64 [[Y:%.*]], 0
+; CHECK-NEXT:    ret i1 [[R]]
+;
+  %g = getelementptr nuw i8, ptr %x, i64 %y
+  %r = icmp ugt ptr %g, %x
+  ret i1 %r
+}
+
+define i1 @ugt_base_nusw_nuw(ptr %x, i64 %y) {
+; CHECK-LABEL: @ugt_base_nusw_nuw(
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i64 [[Y:%.*]], 0
+; CHECK-NEXT:    ret i1 [[R]]
+;
+  %g = getelementptr nusw nuw i8, ptr %x, i64 %y
+  %r = icmp ugt ptr %g, %x
+  ret i1 %r
+}
+
+define i1 @uge_base_nuw(ptr %x, i64 %y) {
+; CHECK-LABEL: @uge_base_nuw(
+; CHECK-NEXT:    ret i1 true
+;
+  %g = getelementptr nuw i8, ptr %x, i64 %y
+  %r = icmp uge ptr %g, %x
+  ret i1 %r
+}
+
+define i1 @uge_base_nusw_nuw(ptr %x, i64 %y) {
+; CHECK-LABEL: @uge_base_nusw_nuw(
+; CHECK-NEXT:    ret i1 true
+;
+  %g = getelementptr nusw nuw i8, ptr %x, i64 %y
+  %r = icmp uge ptr %g, %x
+  ret i1 %r
+}
+
 define i1 @ugt_base_inbounds_commute(i64 %y) {
 ; CHECK-LABEL: @ugt_base_inbounds_commute(
 ; CHECK-NEXT:    [[X:%.*]] = call ptr @getptr()
@@ -319,6 +357,43 @@ define i1 @test60_nusw_inbounds(ptr %foo, i64 %i, i64 %j) {
   ret i1 %cmp
 }
 
+define i1 @test60_nuw(ptr %foo, i64 %i, i64 %j) {
+; CHECK-LABEL: @test60_nuw(
+; CHECK-NEXT:    [[GEP1_IDX:%.*]] = shl nuw i64 [[I:%.*]], 2
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[GEP1_IDX]], [[J:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nuw i32, ptr %foo, i64 %i
+  %gep2 = getelementptr nuw i8, ptr %foo, i64 %j
+  %cmp = icmp ult ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
+define i1 @test60_nusw_nuw(ptr %foo, i64 %i, i64 %j) {
+; CHECK-LABEL: @test60_nusw_nuw(
+; CHECK-NEXT:    [[GEP1_IDX:%.*]] = shl nuw nsw i64 [[I:%.*]], 2
+; CHECK-NEXT:    [[CMP:%.*]] = icmp samesign ult i64 [[GEP1_IDX]], [[J:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nusw nuw i32, ptr %foo, i64 %i
+  %gep2 = getelementptr nusw nuw i8, ptr %foo, i64 %j
+  %cmp = icmp ult ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
+define i1 @test60_nusw_nuw_mix(ptr %foo, i64 %i, i64 %j) {
+; CHECK-LABEL: @test60_nusw_nuw_mix(
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr nuw i32, ptr [[FOO:%.*]], i64 [[I:%.*]]
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr nusw i8, ptr [[FOO]], i64 [[J:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult ptr [[GEP1]], [[GEP2]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nuw i32, ptr %foo, i64 %i
+  %gep2 = getelementptr nusw i8, ptr %foo, i64 %j
+  %cmp = icmp ult ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
 define i1 @test_gep_ult_no_inbounds(ptr %foo, i64 %i, i64 %j) {
 ; CHECK-LABEL: @test_gep_ult_no_inbounds(
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr i32, ptr [[FOO:%.*]], i64 [[I:%.*]]