From fe56f19c6705d652183aaa0fcfc7cf67ec88f1fc Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Tue, 4 Jun 2024 09:55:57 +0200
Subject: [PATCH] [Local] Use nusw and nuw flags in emitGEPOffset()

---
 llvm/lib/Analysis/Local.cpp                  | 11 +++----
 llvm/test/Transforms/InstCombine/icmp-gep.ll | 34 ++++++++++++++++++++
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Analysis/Local.cpp b/llvm/lib/Analysis/Local.cpp
index f5e080d2c78e..a668beabbb29 100644
--- a/llvm/lib/Analysis/Local.cpp
+++ b/llvm/lib/Analysis/Local.cpp
@@ -25,13 +25,13 @@ Value *llvm::emitGEPOffset(IRBuilderBase *Builder, const DataLayout &DL,
   Type *IntIdxTy = DL.getIndexType(GEP->getType());
   Value *Result = nullptr;
 
-  // If the GEP is inbounds, we know that none of the addressing operations will
-  // overflow in a signed sense.
-  bool isInBounds = GEPOp->isInBounds() && !NoAssumptions;
+  // nusw implies nsw for the offset arithmetic.
+  bool NSW = GEPOp->hasNoUnsignedSignedWrap() && !NoAssumptions;
+  bool NUW = GEPOp->hasNoUnsignedWrap() && !NoAssumptions;
   auto AddOffset = [&](Value *Offset) {
     if (Result)
       Result = Builder->CreateAdd(Result, Offset, GEP->getName() + ".offs",
-                                  false /*NUW*/, isInBounds /*NSW*/);
+                                  NUW, NSW);
     else
       Result = Offset;
   };
@@ -71,8 +71,7 @@ Value *llvm::emitGEPOffset(IRBuilderBase *Builder, const DataLayout &DL,
         Scale = Builder->CreateVectorSplat(
             cast<VectorType>(IntIdxTy)->getElementCount(), Scale);
       // We'll let instcombine(mul) convert this to a shl if possible.
-      Op = Builder->CreateMul(Op, Scale, GEP->getName() + ".idx", false /*NUW*/,
-                              isInBounds /*NSW*/);
+      Op = Builder->CreateMul(Op, Scale, GEP->getName() + ".idx", NUW, NSW);
     }
     AddOffset(Op);
   }
diff --git a/llvm/test/Transforms/InstCombine/icmp-gep.ll b/llvm/test/Transforms/InstCombine/icmp-gep.ll
index 29d0c941ac5c..ce64ab1c6305 100644
--- a/llvm/test/Transforms/InstCombine/icmp-gep.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-gep.ll
@@ -546,3 +546,37 @@ define i1 @test_scalable_ij(ptr %foo, i64 %i, i64 %j) {
   %cmp = icmp ult ptr %gep1, %gep2
   ret i1 %cmp
 }
+
+define i1 @gep_nuw(ptr %p, i64 %a, i64 %b, i64 %c, i64 %d) {
+; CHECK-LABEL: @gep_nuw(
+; CHECK-NEXT:    [[GEP1_IDX:%.*]] = shl nuw i64 [[A:%.*]], 2
+; CHECK-NEXT:    [[GEP1_IDX1:%.*]] = shl nuw i64 [[B:%.*]], 1
+; CHECK-NEXT:    [[GEP1_OFFS:%.*]] = add nuw i64 [[GEP1_IDX]], [[GEP1_IDX1]]
+; CHECK-NEXT:    [[GEP2_IDX:%.*]] = shl nuw i64 [[C:%.*]], 3
+; CHECK-NEXT:    [[GEP2_IDX2:%.*]] = shl nuw i64 [[D:%.*]], 2
+; CHECK-NEXT:    [[GEP2_OFFS:%.*]] = add nuw i64 [[GEP2_IDX]], [[GEP2_IDX2]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[GEP1_OFFS]], [[GEP2_OFFS]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nuw [2 x i16], ptr %p, i64 %a, i64 %b
+  %gep2 = getelementptr nuw [2 x i32], ptr %p, i64 %c, i64 %d
+  %cmp = icmp eq ptr %gep1, %gep2
+  ret i1 %cmp
+}
+
+define i1 @gep_nusw(ptr %p, i64 %a, i64 %b, i64 %c, i64 %d) {
+; CHECK-LABEL: @gep_nusw(
+; CHECK-NEXT:    [[GEP1_IDX:%.*]] = shl nsw i64 [[A:%.*]], 2
+; CHECK-NEXT:    [[GEP1_IDX1:%.*]] = shl nsw i64 [[B:%.*]], 1
+; CHECK-NEXT:    [[GEP1_OFFS:%.*]] = add nsw i64 [[GEP1_IDX]], [[GEP1_IDX1]]
+; CHECK-NEXT:    [[GEP2_IDX:%.*]] = shl nsw i64 [[C:%.*]], 3
+; CHECK-NEXT:    [[GEP2_IDX2:%.*]] = shl nsw i64 [[D:%.*]], 2
+; CHECK-NEXT:    [[GEP2_OFFS:%.*]] = add nsw i64 [[GEP2_IDX]], [[GEP2_IDX2]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[GEP1_OFFS]], [[GEP2_OFFS]]
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %gep1 = getelementptr nusw [2 x i16], ptr %p, i64 %a, i64 %b
+  %gep2 = getelementptr nusw [2 x i32], ptr %p, i64 %c, i64 %d
+  %cmp = icmp eq ptr %gep1, %gep2
+  ret i1 %cmp
+}