[llvm-branch-commits] [llvm] release/20.x: [InstCombine] Check nowrap flags when folding comparison of GEPs with the same base pointer (#121892) (PR #125858)

[Nikita Popov via llvm-branch-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/125858

Backport of 9725595f3acc0c1aaa354e15ac4ee2b1f8ff4cc9.

>From cc618a314c9432636446c6629e08a2293a9cf24c Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 1 Feb 2025 20:41:15 +0800
Subject: [PATCH] [InstCombine] Check nowrap flags when folding comparison of
 GEPs with the same base pointer (#121892)

Alive2: https://alive2.llvm.org/ce/z/P5XbMx
Closes https://github.com/llvm/llvm-project/issues/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

(cherry picked from commit 9725595f3acc0c1aaa354e15ac4ee2b1f8ff4cc9)
---
 .../InstCombine/InstCombineCompares.cpp       |  2 +-
 .../test/Transforms/InstCombine/opaque-ptr.ll | 35 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 5a4791870ac77bf..b64ac20ab0533c4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -834,7 +834,7 @@ 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 && CanFold(NW)) {
+      else if (NumDifferences == 1 && 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 b05274658e812e7..be734243d14a107 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