[llvm] [InstCombine] Refine nuw propagation in `OptimizePointerDifference` (PR #147059)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/147059

After https://github.com/llvm/llvm-project/pull/146100, the offset may be generated by a previous call to `EmitGEPOffsets`, causing the nuw flag on shl to be lost. This patch handles the `shl+ptradd` case as well. It also fixes a miscompilation in the case of `mul + ptradd`.
Alive2: https://alive2.llvm.org/ce/z/BeaNzE


>From c93e3b757047d9da0524c167a95896657e53955c Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 4 Jul 2025 21:32:36 +0800
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests. NFC.

---
 llvm/test/Transforms/InstCombine/sub-gep.ll | 52 +++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/sub-gep.ll b/llvm/test/Transforms/InstCombine/sub-gep.ll
index 993a06ad1780f..bcd86243f2206 100644
--- a/llvm/test/Transforms/InstCombine/sub-gep.ll
+++ b/llvm/test/Transforms/InstCombine/sub-gep.ll
@@ -1089,3 +1089,55 @@ define <2 x i64> @splat_geps_multiple(ptr %base, i64 %idx0, <2 x i64> %idx1, <2
   %d = sub <2 x i64> %gep2.int, %gep1.int
   ret <2 x i64> %d
 }
+
+define i64 @nuw_ptrdiff_shl_nsw(ptr %base, i64 %idx) {
+; CHECK-LABEL: @nuw_ptrdiff_shl_nsw(
+; CHECK-NEXT:    [[OFFSET:%.*]] = shl nsw i64 [[IDX:%.*]], 3
+; CHECK-NEXT:    ret i64 [[OFFSET]]
+;
+  %offset = shl nsw i64 %idx, 3
+  %gep = getelementptr inbounds i8, ptr %base, i64 %offset
+  %lhs = ptrtoint ptr %gep to i64
+  %rhs = ptrtoint ptr %base to i64
+  %diff = sub nuw i64 %lhs, %rhs
+  ret i64 %diff
+}
+
+define i64 @nuw_ptrdiff_shl_nonsw(ptr %base, i64 %idx) {
+; CHECK-LABEL: @nuw_ptrdiff_shl_nonsw(
+; CHECK-NEXT:    [[OFFSET:%.*]] = shl i64 [[IDX:%.*]], 3
+; CHECK-NEXT:    ret i64 [[OFFSET]]
+;
+  %offset = shl i64 %idx, 3
+  %gep = getelementptr inbounds i8, ptr %base, i64 %offset
+  %lhs = ptrtoint ptr %gep to i64
+  %rhs = ptrtoint ptr %base to i64
+  %diff = sub nuw i64 %lhs, %rhs
+  ret i64 %diff
+}
+
+define i64 @nuw_ptrdiff_mul_nsw_nneg_scale(ptr %base, i64 %idx) {
+; CHECK-LABEL: @nuw_ptrdiff_mul_nsw_nneg_scale(
+; CHECK-NEXT:    [[OFFSET:%.*]] = mul nuw nsw i64 [[IDX:%.*]], 3
+; CHECK-NEXT:    ret i64 [[OFFSET]]
+;
+  %offset = mul nsw i64 %idx, 3
+  %gep = getelementptr inbounds i8, ptr %base, i64 %offset
+  %lhs = ptrtoint ptr %gep to i64
+  %rhs = ptrtoint ptr %base to i64
+  %diff = sub nuw i64 %lhs, %rhs
+  ret i64 %diff
+}
+
+define i64 @nuw_ptrdiff_mul_nsw_unknown_scale(ptr %base, i64 %idx, i64 %scale) {
+; CHECK-LABEL: @nuw_ptrdiff_mul_nsw_unknown_scale(
+; CHECK-NEXT:    [[OFFSET:%.*]] = mul nuw nsw i64 [[IDX:%.*]], [[SCALE:%.*]]
+; CHECK-NEXT:    ret i64 [[OFFSET]]
+;
+  %offset = mul nsw i64 %idx, %scale
+  %gep = getelementptr inbounds i8, ptr %base, i64 %offset
+  %lhs = ptrtoint ptr %gep to i64
+  %rhs = ptrtoint ptr %base to i64
+  %diff = sub nuw i64 %lhs, %rhs
+  ret i64 %diff
+}

>From 44f1ddcba9e7a1470b394027de7d374ac7d91a93 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 4 Jul 2025 22:00:18 +0800
Subject: [PATCH 2/2] [InstCombine] Refine nuw propagation in
 `OptimizePointerDifference`

---
 llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp | 9 ++++++---
 llvm/test/Transforms/InstCombine/sub-gep.ll           | 4 ++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index f727eb0a63e05..8c16dfc82f5d5 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -2164,10 +2164,13 @@ Value *InstCombinerImpl::OptimizePointerDifference(Value *LHS, Value *RHS,
 
   // If this is a single inbounds GEP and the original sub was nuw,
   // then the final multiplication is also nuw.
-  if (auto *I = dyn_cast<Instruction>(Result))
+  if (auto *I = dyn_cast<OverflowingBinaryOperator>(Result))
     if (IsNUW && match(Offset2, m_Zero()) && Base.LHSNW.isInBounds() &&
-        I->getOpcode() == Instruction::Mul)
-      I->setHasNoUnsignedWrap();
+        I->hasNoSignedWrap() && !I->hasNoUnsignedWrap() &&
+        ((I->getOpcode() == Instruction::Mul &&
+          match(I->getOperand(1), m_NonNegative())) ||
+         I->getOpcode() == Instruction::Shl))
+      cast<Instruction>(I)->setHasNoUnsignedWrap();
 
   // If we have a 2nd GEP of the same base pointer, subtract the offsets.
   // If both GEPs are inbounds, then the subtract does not have signed overflow.
diff --git a/llvm/test/Transforms/InstCombine/sub-gep.ll b/llvm/test/Transforms/InstCombine/sub-gep.ll
index bcd86243f2206..9a6c4c788159b 100644
--- a/llvm/test/Transforms/InstCombine/sub-gep.ll
+++ b/llvm/test/Transforms/InstCombine/sub-gep.ll
@@ -1092,7 +1092,7 @@ define <2 x i64> @splat_geps_multiple(ptr %base, i64 %idx0, <2 x i64> %idx1, <2
 
 define i64 @nuw_ptrdiff_shl_nsw(ptr %base, i64 %idx) {
 ; CHECK-LABEL: @nuw_ptrdiff_shl_nsw(
-; CHECK-NEXT:    [[OFFSET:%.*]] = shl nsw i64 [[IDX:%.*]], 3
+; CHECK-NEXT:    [[OFFSET:%.*]] = shl nuw nsw i64 [[IDX:%.*]], 3
 ; CHECK-NEXT:    ret i64 [[OFFSET]]
 ;
   %offset = shl nsw i64 %idx, 3
@@ -1131,7 +1131,7 @@ define i64 @nuw_ptrdiff_mul_nsw_nneg_scale(ptr %base, i64 %idx) {
 
 define i64 @nuw_ptrdiff_mul_nsw_unknown_scale(ptr %base, i64 %idx, i64 %scale) {
 ; CHECK-LABEL: @nuw_ptrdiff_mul_nsw_unknown_scale(
-; CHECK-NEXT:    [[OFFSET:%.*]] = mul nuw nsw i64 [[IDX:%.*]], [[SCALE:%.*]]
+; CHECK-NEXT:    [[OFFSET:%.*]] = mul nsw i64 [[IDX:%.*]], [[SCALE:%.*]]
 ; CHECK-NEXT:    ret i64 [[OFFSET]]
 ;
   %offset = mul nsw i64 %idx, %scale