[InstCombine] Relax the same-underlying-object constraint for the GEP canonicalization

[dtcxzyw]

7d7001b canonicalizes (gep i8, X, (ptrtoint Y) - (ptrtoint X)) into bitcast Y iff X and Y have the same underlying object.

I find that the result of this pattern is usually used as an operand of an icmp in some real-world applications. I think we can do the canonicalization if the result is only used by icmps/ptrtoints.

Alive2: https://alive2.llvm.org/ce/z/j4-HJZ

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

7d7001b canonicalizes (gep i8, X, (ptrtoint Y) - (ptrtoint X)) into bitcast Y iff X and Y have the same underlying object.

I find that the result of this pattern is usually used as an operand of an icmp in some real-world applications. I think we can do the canonicalization if the result is only used by icmps/ptrtoints.

Alive2: https://alive2.llvm.org/ce/z/bfrEjC (Alive2 cannot prove it without noundef on the base pointer X)

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+9-3)
• (modified) llvm/test/Transforms/InstCombine/getelementptr.ll (+95)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index df393d72a85bf2..52adf582f56abc 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2471,13 +2471,19 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
 
       if (TyAllocSize == 1) {
         // Canonicalize (gep i8* X, (ptrtoint Y)-(ptrtoint X)) to (bitcast Y),
-        // but only if both point to the same underlying object (otherwise
-        // provenance is not necessarily retained).
+        // but only if the result pointer is only used as if it were an integer,
+        // or both point to the same underlying object (otherwise provenance is
+        // not necessarily retained).
         Value *X = GEP.getPointerOperand();
         Value *Y;
         if (match(GEP.getOperand(1),
                   m_Sub(m_PtrToInt(m_Value(Y)), m_PtrToInt(m_Specific(X)))) &&
-            getUnderlyingObject(X) == getUnderlyingObject(Y))
+            ((all_of(GEP.users(),
+                     [](User *U) {
+                       return isa<ICmpInst>(U) || isa<PtrToIntInst>(U);
+                     }) &&
+              isGuaranteedNotToBeUndef(X, &AC, &GEP, &DT)) ||
+             getUnderlyingObject(X) == getUnderlyingObject(Y)))
           return CastInst::CreatePointerBitCastOrAddrSpaceCast(Y, GEPType);
       } else {
         // Canonicalize (gep T* X, V / sizeof(T)) to (gep i8* X, V)
diff --git a/llvm/test/Transforms/InstCombine/getelementptr.ll b/llvm/test/Transforms/InstCombine/getelementptr.ll
index 7d67f2583aa24d..68c28897aabeeb 100644
--- a/llvm/test/Transforms/InstCombine/getelementptr.ll
+++ b/llvm/test/Transforms/InstCombine/getelementptr.ll
@@ -1537,4 +1537,99 @@ define ptr @gep_ashr_without_exact(ptr %p, i64 %off) {
   ret ptr %ptr
 }
 
+define i1 @test_only_used_by_icmp(ptr noundef %a, ptr %b, ptr %c) {
+; CHECK-LABEL: @test_only_used_by_icmp(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[B:%.*]], [[C:%.*]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %pa = ptrtoint ptr %a to i64
+  %pb = ptrtoint ptr %b to i64
+  %sub = sub i64 %pb, %pa
+  %gep = getelementptr i8, ptr %a, i64 %sub
+  %cmp = icmp eq ptr %gep, %c
+  ret i1 %cmp
+}
+
+define i64 @test_only_used_by_ptrtoint(ptr noundef %a, ptr %b) {
+; CHECK-LABEL: @test_only_used_by_ptrtoint(
+; CHECK-NEXT:    [[VAL:%.*]] = ptrtoint ptr [[B:%.*]] to i64
+; CHECK-NEXT:    ret i64 [[VAL]]
+;
+  %pa = ptrtoint ptr %a to i64
+  %pb = ptrtoint ptr %b to i64
+  %sub = sub i64 %pb, %pa
+  %gep = getelementptr i8, ptr %a, i64 %sub
+  %val = ptrtoint ptr %gep to i64
+  ret i64 %val
+}
+
+define i64 @test_used_by_both(ptr noundef %a, ptr %b, ptr %c) {
+; CHECK-LABEL: @test_used_by_both(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[B:%.*]], [[C:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[VAL:%.*]] = ptrtoint ptr [[B]] to i64
+; CHECK-NEXT:    ret i64 [[VAL]]
+; CHECK:       if.else:
+; CHECK-NEXT:    ret i64 0
+;
+  %pa = ptrtoint ptr %a to i64
+  %pb = ptrtoint ptr %b to i64
+  %sub = sub i64 %pb, %pa
+  %gep = getelementptr i8, ptr %a, i64 %sub
+  %cmp = icmp eq ptr %gep, %c
+  br i1 %cmp, label %if.then, label %if.else
+if.then:
+  %val = ptrtoint ptr %gep to i64
+  ret i64 %val
+if.else:
+  ret i64 0
+}
+
+; Negative tests
+
+define i64 @test_used_by_both_invalid(ptr noundef %a, ptr %b, ptr %c) {
+; CHECK-LABEL: @test_used_by_both_invalid(
+; CHECK-NEXT:    [[PA:%.*]] = ptrtoint ptr [[A:%.*]] to i64
+; CHECK-NEXT:    [[PB:%.*]] = ptrtoint ptr [[B:%.*]] to i64
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[PB]], [[PA]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 [[SUB]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[GEP]], [[C:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[VAL:%.*]] = load i64, ptr [[GEP]], align 8
+; CHECK-NEXT:    ret i64 [[VAL]]
+; CHECK:       if.else:
+; CHECK-NEXT:    ret i64 0
+;
+  %pa = ptrtoint ptr %a to i64
+  %pb = ptrtoint ptr %b to i64
+  %sub = sub i64 %pb, %pa
+  %gep = getelementptr i8, ptr %a, i64 %sub
+  %cmp = icmp eq ptr %gep, %c
+  br i1 %cmp, label %if.then, label %if.else
+if.then:
+  %val = load i64, ptr %gep, align 8
+  ret i64 %val
+if.else:
+  ret i64 0
+}
+
+define i64 @test_only_used_by_ptrtoint_without_noundef(ptr %a, ptr %b) {
+; CHECK-LABEL: @test_only_used_by_ptrtoint_without_noundef(
+; CHECK-NEXT:    [[PA:%.*]] = ptrtoint ptr [[A:%.*]] to i64
+; CHECK-NEXT:    [[PB:%.*]] = ptrtoint ptr [[B:%.*]] to i64
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[PB]], [[PA]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 [[SUB]]
+; CHECK-NEXT:    [[VAL:%.*]] = ptrtoint ptr [[GEP]] to i64
+; CHECK-NEXT:    ret i64 [[VAL]]
+;
+  %pa = ptrtoint ptr %a to i64
+  %pb = ptrtoint ptr %b to i64
+  %sub = sub i64 %pb, %pa
+  %gep = getelementptr i8, ptr %a, i64 %sub
+  %val = ptrtoint ptr %gep to i64
+  ret i64 %val
+}
+
 !0 = !{!"branch_weights", i32 2, i32 10}

[nikic]

You can use a small bit width for the alive proof, requiring noundef is unnecessary: https://alive2.llvm.org/ce/z/5EjmcW

Do you actually need the ptrtoint case? The icmp one is definitely correct, the ptrtoint one may or may not be depending on what ptrtoint semantics we decide on (e.g. it would be incorrect if it has "expose" semantics).

[nikic]

Do you actually need the ptrtoint case? The icmp one is definitely correct, the ptrtoint one may or may not be depending on what ptrtoint semantics we decide on (e.g. it would be incorrect if it has "expose" semantics).

Nevermind, this should be fine for ptrtoint as well, regardless of the specific semantics, because we already have ptrtoints for both x and y in the first place, so the amount of exposed provenance remains the same.

[dtcxzyw]

You can use a small bit width for the alive proof, requiring noundef is unnecessary: https://alive2.llvm.org/ce/z/5EjmcW

Do you actually need the ptrtoint case? The icmp one is definitely correct, the ptrtoint one may or may not be depending on what ptrtoint semantics we decide on (e.g. it would be incorrect if it has "expose" semantics).

Done.

[nikic]

LGTM