[InstCombine] Check nowrap flags when folding comparison of GEPs with the same base pointer

[dtcxzyw]

Alive2: https://alive2.llvm.org/ce/z/P5XbMx
Closes #121890

TODO: It is still safe to perform this transform without nowrap flags if the corresponding scale factor is 1 byte: https://alive2.llvm.org/ce/z/J-JCJd

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Alive2: https://alive2.llvm.org/ce/z/P5XbMx
Closes #121890

TODO: It is still safe to perform this transform without nowrap flags if the corresponding scale factor is 1 byte: https://alive2.llvm.org/ce/z/J-JCJd

---

Full diff: https://github.com/llvm/llvm-project/pull/121892.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp (+5-2)
• (modified) llvm/test/Transforms/InstCombine/opaque-ptr.ll (+35)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 8b23583c510637..9f136860724db8 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -833,8 +833,11 @@ Instruction *InstCombinerImpl::foldGEPICmp(GEPOperator *GEPLHS, Value *RHS,
       if (NumDifferences == 0)   // SAME GEP?
         return replaceInstUsesWith(I, // No comparison is needed here.
           ConstantInt::get(I.getType(), ICmpInst::isTrueWhenEqual(Cond)));
-
-      else if (NumDifferences == 1 && CanFold(NW)) {
+      // If two GEPs only differ by an index/the base pointer, compare them.
+      // Note that nowrap flags are always needed when comparing two indices.
+      else if (NumDifferences == 1 &&
+               (DiffOperand == 0 ? CanFold(NW)
+                                 : NW != GEPNoWrapFlags::none())) {
         Value *LHSV = GEPLHS->getOperand(DiffOperand);
         Value *RHSV = GEPRHS->getOperand(DiffOperand);
         return NewICmp(NW, LHSV, RHSV);
diff --git a/llvm/test/Transforms/InstCombine/opaque-ptr.ll b/llvm/test/Transforms/InstCombine/opaque-ptr.ll
index b05274658e812e..be734243d14a10 100644
--- a/llvm/test/Transforms/InstCombine/opaque-ptr.ll
+++ b/llvm/test/Transforms/InstCombine/opaque-ptr.ll
@@ -467,6 +467,41 @@ define i1 @cmp_gep_same_base_same_type(ptr %ptr, i64 %idx1, i64 %idx2) {
   ret i1 %cmp
 }
 
+define i1 @cmp_gep_same_base_same_type_maywrap(ptr %ptr, i64 %idx1, i64 %idx2) {
+; CHECK-LABEL: @cmp_gep_same_base_same_type_maywrap(
+; CHECK-NEXT:    [[CMP_UNSHIFTED:%.*]] = xor i64 [[IDX1:%.*]], [[IDX2:%.*]]
+; CHECK-NEXT:    [[CMP_MASK:%.*]] = and i64 [[CMP_UNSHIFTED]], 4611686018427387903
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[CMP_MASK]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr i32, ptr %ptr, i64 %idx1
+  %gep2 = getelementptr i32, ptr %ptr, i64 %idx2
+  %cmp = icmp eq ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
+define i1 @cmp_gep_same_base_same_type_nuw(ptr %ptr, i64 %idx1, i64 %idx2) {
+; CHECK-LABEL: @cmp_gep_same_base_same_type_nuw(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[IDX1:%.*]], [[IDX2:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nuw i32, ptr %ptr, i64 %idx1
+  %gep2 = getelementptr nuw i32, ptr %ptr, i64 %idx2
+  %cmp = icmp eq ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
+define i1 @cmp_gep_same_base_same_type_nusw(ptr %ptr, i64 %idx1, i64 %idx2) {
+; CHECK-LABEL: @cmp_gep_same_base_same_type_nusw(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[IDX1:%.*]], [[IDX2:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nusw i32, ptr %ptr, i64 %idx1
+  %gep2 = getelementptr nusw i32, ptr %ptr, i64 %idx2
+  %cmp = icmp eq ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
 define i1 @cmp_gep_same_base_different_type(ptr %ptr, i64 %idx1, i64 %idx2) {
 ; CHECK-LABEL: @cmp_gep_same_base_different_type(
 ; CHECK-NEXT:    [[GEP1_IDX:%.*]] = shl nsw i64 [[IDX1:%.*]], 2

[nikic]

I don't think DiffOperand == 0 can happen here? The above loop only goes over the index operands. The case of base pointer difference is handled by different code further up.

[nikic]

LGTM. We'll still fall back to the generic handling below, so I think the simple check is fine.

[nikic]

Reverse ping :)

[dtcxzyw]

Reverse ping :)

Is it necessary to merge this PR now? I was previously planning to merge it after fixing the regression.

[nikic]

Reverse ping :)

Is it necessary to merge this PR now? I was previously planning to merge it after fixing the regression.

As it fixes a miscompile and the regression doesn't seem particularly significant, I think it's better to merge it earlier than later.

[nikic]

/cherry-pick 9725595

[llvmbot]

Failed to cherry-pick: 9725595

https://github.com/llvm/llvm-project/actions/runs/13097312131

Please manually backport the fix and push it to your github fork. Once this is done, please create a pull request