[llvm] [LLVM] [SeparateConstOffsetFromGEP] Fix sep-const-offset-from-gep invalid assumption (PR #183402)

Meredith Julian via llvm-commits llvm-commits at lists.llvm.org
Fri Mar 6 15:47:05 PST 2026


https://github.com/mjulian31 updated https://github.com/llvm/llvm-project/pull/183402

>From e3541d8798b1c2edc7c24b1fec5733ad78d05ee5 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Wed, 4 Feb 2026 17:09:14 -0800
Subject: [PATCH 1/8] Fix sep-const-offset-from-gep invalid assumption

---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 33 +++++++++-
 .../NVPTX/split-gep.ll                        | 63 ++++++++++++++++++-
 2 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index c298daff30108..160c08b8aedf5 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -775,14 +775,41 @@ static bool allowsPreservingNUW(const User *U) {
   return true;
 }
 
+// An inbounds GEP does not guarantee that the index is non-negative.
+// This helper checks first if the index is known non-negative. If it is not,
+// it checks whether the GEP is directly based on a global or an alloca with
+// zero offset, in which case inbounds is sufficient to prove non-negativity.
+static bool isGEPNonNegative(const GetElementPtrInst *GEP,
+                             const Value *Idx,
+                             const DataLayout &DL) {
+  if (isKnownNonNegative(Idx, DL))
+    return true;
+
+  if (!GEP->isInBounds())
+    return false;
+
+  const Value *Ptr = GEP->getPointerOperand();
+  int64_t Offset = 0;
+  const Value *Base = GetPointerBaseWithConstantOffset(
+      const_cast<Value *>(Ptr), Offset, DL, /*AllowNonInbounds=*/false);
+  if (!Base || Offset != 0)
+    return false;
+
+  if (!isa<AllocaInst>(Base) && !isa<GlobalObject>(Base))
+    return false;
+
+  return Base == getUnderlyingObject(Ptr);
+}
+
 Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
                                         User *&UserChainTail,
                                         bool &PreservesNUW) {
   ConstantOffsetExtractor Extractor(GEP->getIterator());
+  bool GEPNonNegative = isGEPNonNegative(GEP, Idx, Extractor.DL);
   // Find a non-zero constant offset first.
   APInt ConstantOffset =
       Extractor.find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-                     GEP->isInBounds());
+                     /* NonNegative */ GEPNonNegative);
   if (ConstantOffset == 0) {
     UserChainTail = nullptr;
     PreservesNUW = true;
@@ -798,10 +825,10 @@ Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
 }
 
 APInt ConstantOffsetExtractor::Find(Value *Idx, GetElementPtrInst *GEP) {
-  // If Idx is an index of an inbound GEP, Idx is guaranteed to be non-negative.
+  bool GEPNonNegative = isGEPNonNegative(GEP, Idx, GEP->getDataLayout());
   return ConstantOffsetExtractor(GEP->getIterator())
       .find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-            GEP->isInBounds());
+            /* NonNegative */ GEPNonNegative);
 }
 
 bool SeparateConstOffsetFromGEP::canonicalizeArrayIndicesToIndexSize(
diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
index 77b3434f4f159..b8b07052dcf2d 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
@@ -30,8 +30,8 @@ entry:
 }
 
 ; We should be able to trace into sext(a + b) if a + b is non-negative
-; (e.g., used as an index of an inbounds GEP) and one of a and b is
-; non-negative.
+; (e.g., used as an index of an inbounds GEP on a global base ptr) and one of a
+; or b is non-negative.
 define ptr @sext_add(i32 %i, i32 %j) {
 ; CHECK-LABEL: define ptr @sext_add(
 ; CHECK-SAME: i32 [[I:%.*]], i32 [[J:%.*]]) {
@@ -53,6 +53,65 @@ entry:
   ret ptr %p
 }
 
+; We should trace into sext(a + b) if a + b is an inbounds GEP on a known
+; base ptr (alloca) if one of a or b is non-negative.
+define ptr @sext_add_alloca(i32 %i) {
+; CHECK-LABEL: define ptr @sext_add_alloca(
+; CHECK-SAME: i32 [[I:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARR:%.*]] = alloca [32 x [32 x float]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [32 x [32 x float]], ptr [[ARR]], i64 0, i64 [[TMP0]], i64 0
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr i8, ptr [[TMP1]], i64 128
+; CHECK-NEXT:    ret ptr [[P1]]
+;
+entry:
+  %arr = alloca [32 x [32 x float]], align 4
+  %0 = add i32 %i, 1
+  %1 = sext i32 %0 to i64
+  ; inbound sext(i + 1) = sext(i) + 1 because inbounds on base ptr -> non-negative
+  %p = getelementptr inbounds [32 x [32 x float]], ptr %arr, i64 0, i64 %1, i64 0
+  ret ptr %p
+}
+
+; We cannot trace into sext(a + b) if a + b is an inbounds GEP but not on a
+; known base ptr even if one of a or b is non-negative.
+define ptr @sext_add_nonbase(i32 %i, ptr %unknown_arr) {
+; CHECK-LABEL: define ptr @sext_add_nonbase(
+; CHECK-SAME: i32 [[I:%.*]], ptr [[ARR:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[TMP0]] to i64
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr inbounds [32 x [32 x float]], ptr [[ARR]], i64 0, i64 [[TMP1]], i64 0
+; CHECK-NEXT:    ret ptr [[P1]]
+;
+entry:
+  %0 = add i32 %i, 1
+  %1 = sext i32 %0 to i64
+  ; inbound sext(i + 1) != sext(i) + 1 because a wrapped result can still be inbounds if not at start of arr
+  %p = getelementptr inbounds [32 x [32 x float]], ptr %unknown_arr, i64 0, i64 %1, i64 0
+  ret ptr %p
+}
+
+; We can trace into sext(a + b) if a + b is non-negative (nsw flag) and one of
+; a or b is non-negative, even if the gep is not inbounds
+define ptr @sext_add_nsw(i32 %i, ptr %unknown_arr) {
+; CHECK-LABEL: define ptr @sext_add_nsw(
+; CHECK-SAME: i32 [[I:%.*]], ptr [[ARR:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [32 x [32 x float]], ptr [[ARR]], i64 0, i64 [[TMP0]], i64 0
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr i8, ptr [[TMP1]], i64 128
+; CHECK-NEXT:    ret ptr [[P1]]
+;
+entry:
+  %0 = add nsw i32 %i, 1
+  %1 = sext i32 %0 to i64
+  ; sext(nsw i + 1) = sext(i) + 1
+  %p = getelementptr [32 x [32 x float]], ptr %unknown_arr, i64 0, i64 %1, i64 0
+  ret ptr %p
+}
+
 ; We should be able to trace into sext/zext if it can be distributed to both
 ; operands, e.g., sext (add nsw a, b) == add nsw (sext a), (sext b)
 ;

>From 2619a9f01cd8753aaa6ee3259a0cd194559f21cd Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Thu, 5 Feb 2026 15:10:56 -0800
Subject: [PATCH 2/8] add non-zero offset test

---
 .../NVPTX/split-gep.ll                        | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
index b8b07052dcf2d..39da11c8339ca 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
@@ -10,6 +10,7 @@
 
 @struct_array = global [1024 x %struct.S] zeroinitializer, align 16
 @float_2d_array = global [32 x [32 x float]] zeroinitializer, align 4
+ at float_array = global [128 x float] zeroinitializer, align 4
 
 ; We should not extract any struct field indices, because fields in a struct
 ; may have different types.
@@ -93,6 +94,26 @@ entry:
   ret ptr %p
 }
 
+; We cannot trace into sext(a + b) if a + b is an inbounds GEP but not a zero
+; offset from a known base ptr even if one of a or b is non-negative.
+define ptr @sext_add_nonzerooffset(i4 %i) {
+; CHECK-LABEL: define ptr @sext_add_nonzerooffset(
+; CHECK-SAME: i4 [[I:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i4 [[I]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i4 [[TMP0]] to i64
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr float, ptr @float_array, i64 [[TMP1]]
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr [128 x float], ptr [[GEP1]], i64 0, i64 64
+; CHECK-NEXT:    ret ptr [[P1]]
+;
+entry:
+  %offsetarr = getelementptr inbounds [128 x float], ptr @float_array, i64 0, i64 64
+  %add = add i4 %i, 1
+  %sext = sext i4 %add to i64
+  %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
+  ret ptr %p
+}
+
 ; We can trace into sext(a + b) if a + b is non-negative (nsw flag) and one of
 ; a or b is non-negative, even if the gep is not inbounds
 define ptr @sext_add_nsw(i32 %i, ptr %unknown_arr) {

>From d3f59a38ba21daa74710facc2810c515f83bb9d3 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Wed, 25 Feb 2026 13:08:47 -0800
Subject: [PATCH 3/8] simplify check

---
 llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 160c08b8aedf5..3daba800719af 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -792,13 +792,11 @@ static bool isGEPNonNegative(const GetElementPtrInst *GEP,
   int64_t Offset = 0;
   const Value *Base = GetPointerBaseWithConstantOffset(
       const_cast<Value *>(Ptr), Offset, DL, /*AllowNonInbounds=*/false);
-  if (!Base || Offset != 0)
-    return false;
 
-  if (!isa<AllocaInst>(Base) && !isa<GlobalObject>(Base))
+  if (!Base || Offset != 0)
     return false;
 
-  return Base == getUnderlyingObject(Ptr);
+  return isa<AllocaInst>(Base) || isa<GlobalObject>(Base);
 }
 
 Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,

>From 51fed049345fa773a9c246b03df43898bb6ef771 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Wed, 25 Feb 2026 14:21:40 -0800
Subject: [PATCH 4/8] fix clang-format

---
 llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 3daba800719af..bac7ad2ed9b51 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -779,8 +779,7 @@ static bool allowsPreservingNUW(const User *U) {
 // This helper checks first if the index is known non-negative. If it is not,
 // it checks whether the GEP is directly based on a global or an alloca with
 // zero offset, in which case inbounds is sufficient to prove non-negativity.
-static bool isGEPNonNegative(const GetElementPtrInst *GEP,
-                             const Value *Idx,
+static bool isGEPNonNegative(const GetElementPtrInst *GEP, const Value *Idx,
                              const DataLayout &DL) {
   if (isKnownNonNegative(Idx, DL))
     return true;

>From 3cfefa68b5f9f33a06553617117a4e0df0d7a634 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Wed, 25 Feb 2026 18:44:36 -0800
Subject: [PATCH 5/8] add case for sext(nuw add) into inbounds nuw gep

---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 66 ++++++++++++-------
 1 file changed, 42 insertions(+), 24 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index bac7ad2ed9b51..39074440323d0 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -233,10 +233,12 @@ class ConstantOffsetExtractor {
   /// \p ZeroExtended Whether V will be zero-extended in the computation of the
   ///                 GEP index
   /// \p NonNegative  Whether V is guaranteed to be non-negative. For example,
-  ///                 an index of an inbounds GEP is guaranteed to be
-  ///                 non-negative. Levaraging this, we can better split
-  ///                 inbounds GEPs.
-  APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative);
+  ///                 an index of an inbounds GEP of a base address is
+  ///                 guaranteed to be non-negative. Leveraging this, we can
+  ///                 better split inbounds GEPs.
+  /// \p GEPInboundsNUW  Whether the GEP is both inbounds and nuw.
+  APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative,
+             bool GEPInboundsNUW);
 
   /// A helper function to look into both operands of a binary operator.
   APInt findInEitherOperand(BinaryOperator *BO, bool SignExtended,
@@ -292,8 +294,9 @@ class ConstantOffsetExtractor {
   /// \p ZeroExtended Whether BO is surrounded by zext
   /// \p NonNegative Whether BO is known to be non-negative, e.g., an in-bound
   ///                array index.
+  /// \p GEPInboundsNUW 
   bool CanTraceInto(bool SignExtended, bool ZeroExtended, BinaryOperator *BO,
-                    bool NonNegative);
+                    bool NonNegative, bool GEPInboundsNUW);
 
   /// The path from the constant offset to the old GEP index. e.g., if the GEP
   /// index is "a * b + (c + 5)". After running function find, UserChain[0] will
@@ -474,10 +477,9 @@ FunctionPass *llvm::createSeparateConstOffsetFromGEPPass(bool LowerGEP) {
   return new SeparateConstOffsetFromGEPLegacyPass(LowerGEP);
 }
 
-bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended,
-                                            bool ZeroExtended,
-                                            BinaryOperator *BO,
-                                            bool NonNegative) {
+bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended, bool ZeroExtended,
+                                           BinaryOperator *BO, bool NonNegative,
+                                           bool GEPInboundsNUW) {
   // We only consider ADD, SUB and OR, because a non-zero constant found in
   // expressions composed of these operations can be easily hoisted as a
   // constant offset by reassociation.
@@ -530,6 +532,12 @@ bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended,
     }
   }
 
+  // For a sext(add nuw), allow tracing through when the enclosing GEP is both
+  // inbounds and nuw.
+  if (BO->getOpcode() == Instruction::Add && SignExtended && !ZeroExtended &&
+      GEPInboundsNUW && BO->hasNoUnsignedWrap())
+    return true;
+
   // sext (add/sub nsw A, B) == add/sub nsw (sext A), (sext B)
   // zext (add/sub nuw A, B) == add/sub nuw (zext A), (zext B)
   if (BO->getOpcode() == Instruction::Add ||
@@ -550,9 +558,10 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   size_t ChainLength = UserChain.size();
 
   // BO being non-negative does not shed light on whether its operands are
-  // non-negative. Clear the NonNegative flag here.
-  APInt ConstantOffset = find(BO->getOperand(0), SignExtended, ZeroExtended,
-                              /* NonNegative */ false);
+  // non-negative. Clear the NonNegative and GEPInboundsNUW flags here.
+  APInt ConstantOffset =
+      find(BO->getOperand(0), SignExtended, ZeroExtended,
+           /* NonNegative */ false, /* GEPInboundsNUW */ false);
   // If we found a constant offset in the left operand, stop and return that.
   // This shortcut might cause us to miss opportunities of combining the
   // constant offsets in both operands, e.g., (a + 4) + (b + 5) => (a + b) + 9.
@@ -565,7 +574,7 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   UserChain.resize(ChainLength);
 
   ConstantOffset = find(BO->getOperand(1), SignExtended, ZeroExtended,
-                        /* NonNegative */ false);
+                        /* NonNegative */ false, /* GEPInboundsNUW */ false);
   // If U is a sub operator, negate the constant offset found in the right
   // operand.
   if (BO->getOpcode() == Instruction::Sub)
@@ -579,7 +588,8 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
 }
 
 APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
-                                    bool ZeroExtended, bool NonNegative) {
+                                    bool ZeroExtended, bool NonNegative,
+                                    bool GEPInboundsNUW) {
   // TODO(jingyue): We could trace into integer/pointer casts, such as
   // inttoptr, ptrtoint, bitcast, and addrspacecast. We choose to handle only
   // integers because it gives good enough results for our benchmarks.
@@ -595,23 +605,27 @@ APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
     ConstantOffset = CI->getValue();
   } else if (BinaryOperator *BO = dyn_cast<BinaryOperator>(V)) {
     // Trace into subexpressions for more hoisting opportunities.
-    if (CanTraceInto(SignExtended, ZeroExtended, BO, NonNegative))
+    if (CanTraceInto(SignExtended, ZeroExtended, BO, NonNegative,
+                     GEPInboundsNUW))
       ConstantOffset = findInEitherOperand(BO, SignExtended, ZeroExtended);
   } else if (isa<TruncInst>(V)) {
-    ConstantOffset =
-        find(U->getOperand(0), SignExtended, ZeroExtended, NonNegative)
-            .trunc(BitWidth);
+    ConstantOffset = find(U->getOperand(0), SignExtended, ZeroExtended,
+                          NonNegative, GEPInboundsNUW).trunc(BitWidth);
   } else if (isa<SExtInst>(V)) {
     ConstantOffset = find(U->getOperand(0), /* SignExtended */ true,
-                          ZeroExtended, NonNegative).sext(BitWidth);
+                          ZeroExtended, NonNegative, GEPInboundsNUW)
+                         .sext(BitWidth);
   } else if (isa<ZExtInst>(V)) {
     // As an optimization, we can clear the SignExtended flag because
     // sext(zext(a)) = zext(a). Verified in @sext_zext in split-gep.ll.
     //
+    // We can also clear the GEPInboundsNUW flag because it is not required to
+    // trace into zext.
+    //
     // Clear the NonNegative flag, because zext(a) >= 0 does not imply a >= 0.
-    ConstantOffset =
-        find(U->getOperand(0), /* SignExtended */ false,
-             /* ZeroExtended */ true, /* NonNegative */ false).zext(BitWidth);
+    ConstantOffset = find(U->getOperand(0), /* SignExtended */ false,
+                          /* ZeroExtended */ true, /* NonNegative */ false,
+                          /* GEPInboundsNUW */ false).zext(BitWidth);
   }
 
   // If we found a non-zero constant offset, add it to the path for
@@ -803,10 +817,12 @@ Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
                                         bool &PreservesNUW) {
   ConstantOffsetExtractor Extractor(GEP->getIterator());
   bool GEPNonNegative = isGEPNonNegative(GEP, Idx, Extractor.DL);
+  bool GEPInboundsNUW = GEP->isInBounds() && GEP->hasNoUnsignedWrap();
   // Find a non-zero constant offset first.
   APInt ConstantOffset =
       Extractor.find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-                     /* NonNegative */ GEPNonNegative);
+                     /* NonNegative */ GEPNonNegative,
+                     /* GEPInboundsNUW */ GEPInboundsNUW);
   if (ConstantOffset == 0) {
     UserChainTail = nullptr;
     PreservesNUW = true;
@@ -823,9 +839,11 @@ Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
 
 APInt ConstantOffsetExtractor::Find(Value *Idx, GetElementPtrInst *GEP) {
   bool GEPNonNegative = isGEPNonNegative(GEP, Idx, GEP->getDataLayout());
+  bool GEPInboundsNUW = GEP->isInBounds() && GEP->hasNoUnsignedWrap();
   return ConstantOffsetExtractor(GEP->getIterator())
       .find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-            /* NonNegative */ GEPNonNegative);
+            /* NonNegative */ GEPNonNegative,
+            /* GEPInboundsNUW */ GEPInboundsNUW);
 }
 
 bool SeparateConstOffsetFromGEP::canonicalizeArrayIndicesToIndexSize(

>From 1a465de1c8cfa125587714f1cd5bc6b5db4a1fe5 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Wed, 25 Feb 2026 18:54:52 -0800
Subject: [PATCH 6/8] clang-format fix

---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 39074440323d0..1f92c4eeb7888 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -227,16 +227,16 @@ class ConstantOffsetExtractor {
   /// successful, returns C and update UserChain as a def-use chain from C to V;
   /// otherwise, UserChain is empty.
   ///
-  /// \p V            The given expression
-  /// \p SignExtended Whether V will be sign-extended in the computation of the
-  ///                 GEP index
-  /// \p ZeroExtended Whether V will be zero-extended in the computation of the
-  ///                 GEP index
-  /// \p NonNegative  Whether V is guaranteed to be non-negative. For example,
-  ///                 an index of an inbounds GEP of a base address is
-  ///                 guaranteed to be non-negative. Leveraging this, we can
-  ///                 better split inbounds GEPs.
-  /// \p GEPInboundsNUW  Whether the GEP is both inbounds and nuw.
+  /// \p V              The given expression
+  /// \p SignExtended   Whether V will be sign-extended in the computation of
+  ///                   the GEP index
+  /// \p ZeroExtended   Whether V will be zero-extended in the computation of
+  ///                   the GEP index
+  /// \p NonNegative    Whether V is guaranteed to be non-negative. For example,
+  ///                   an index of an inbounds GEP of a base address is
+  ///                   guaranteed to be non-negative. Leveraging this, we can
+  ///                   better split inbounds GEPs.
+  /// \p GEPInboundsNUW Whether the GEP is both inbounds and nuw.
   APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative,
              bool GEPInboundsNUW);
 
@@ -294,7 +294,7 @@ class ConstantOffsetExtractor {
   /// \p ZeroExtended Whether BO is surrounded by zext
   /// \p NonNegative Whether BO is known to be non-negative, e.g., an in-bound
   ///                array index.
-  /// \p GEPInboundsNUW 
+  /// \p GEPInboundsNUW Whether the GEP is both inbounds and nuw.
   bool CanTraceInto(bool SignExtended, bool ZeroExtended, BinaryOperator *BO,
                     bool NonNegative, bool GEPInboundsNUW);
 
@@ -610,7 +610,8 @@ APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
       ConstantOffset = findInEitherOperand(BO, SignExtended, ZeroExtended);
   } else if (isa<TruncInst>(V)) {
     ConstantOffset = find(U->getOperand(0), SignExtended, ZeroExtended,
-                          NonNegative, GEPInboundsNUW).trunc(BitWidth);
+                          NonNegative, GEPInboundsNUW)
+                         .trunc(BitWidth);
   } else if (isa<SExtInst>(V)) {
     ConstantOffset = find(U->getOperand(0), /* SignExtended */ true,
                           ZeroExtended, NonNegative, GEPInboundsNUW)
@@ -625,7 +626,8 @@ APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
     // Clear the NonNegative flag, because zext(a) >= 0 does not imply a >= 0.
     ConstantOffset = find(U->getOperand(0), /* SignExtended */ false,
                           /* ZeroExtended */ true, /* NonNegative */ false,
-                          /* GEPInboundsNUW */ false).zext(BitWidth);
+                          /* GEPInboundsNUW */ false)
+                         .zext(BitWidth);
   }
 
   // If we found a non-zero constant offset, add it to the path for

>From df87662480fc8ab2e724d5facf72100910069fa4 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Fri, 6 Mar 2026 15:38:57 -0800
Subject: [PATCH 7/8] add offset/object size threshold logic and tests

---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 221 +++++++++++-------
 .../NVPTX/split-gep.ll                        | 123 ++++++++--
 2 files changed, 245 insertions(+), 99 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 1f92c4eeb7888..a050cdbb99301 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -228,17 +228,16 @@ class ConstantOffsetExtractor {
   /// otherwise, UserChain is empty.
   ///
   /// \p V              The given expression
+  /// \p GEP            The base GEP instruction, used for determining relevant
+  ///                   types, flags, and non-negativity needed for safe
+  ///                   reassociation
+  /// \p Idx            The original index of the GEP
   /// \p SignExtended   Whether V will be sign-extended in the computation of
   ///                   the GEP index
   /// \p ZeroExtended   Whether V will be zero-extended in the computation of
   ///                   the GEP index
-  /// \p NonNegative    Whether V is guaranteed to be non-negative. For example,
-  ///                   an index of an inbounds GEP of a base address is
-  ///                   guaranteed to be non-negative. Leveraging this, we can
-  ///                   better split inbounds GEPs.
-  /// \p GEPInboundsNUW Whether the GEP is both inbounds and nuw.
-  APInt find(Value *V, bool SignExtended, bool ZeroExtended, bool NonNegative,
-             bool GEPInboundsNUW);
+  APInt find(Value *V, GetElementPtrInst *GEP, Value *Idx, bool SignExtended,
+             bool ZeroExtended);
 
   /// A helper function to look into both operands of a binary operator.
   APInt findInEitherOperand(BinaryOperator *BO, bool SignExtended,
@@ -292,11 +291,11 @@ class ConstantOffsetExtractor {
   ///
   /// \p SignExtended Whether BO is surrounded by sext
   /// \p ZeroExtended Whether BO is surrounded by zext
-  /// \p NonNegative Whether BO is known to be non-negative, e.g., an in-bound
-  ///                array index.
-  /// \p GEPInboundsNUW Whether the GEP is both inbounds and nuw.
-  bool CanTraceInto(bool SignExtended, bool ZeroExtended, BinaryOperator *BO,
-                    bool NonNegative, bool GEPInboundsNUW);
+  /// \p GEP          The base GEP instruction, used for determining relevant
+  ///                 types and flags needed for safe reassociation.
+  /// \p Idx          The original index of the GEP
+  bool canTraceInto(bool SignExtended, bool ZeroExtended, BinaryOperator *BO,
+                    GetElementPtrInst *GEP, Value *Idx);
 
   /// The path from the constant offset to the old GEP index. e.g., if the GEP
   /// index is "a * b + (c + 5)". After running function find, UserChain[0] will
@@ -477,9 +476,112 @@ FunctionPass *llvm::createSeparateConstOffsetFromGEPPass(bool LowerGEP) {
   return new SeparateConstOffsetFromGEPLegacyPass(LowerGEP);
 }
 
-bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended, bool ZeroExtended,
-                                           BinaryOperator *BO, bool NonNegative,
-                                           bool GEPInboundsNUW) {
+// Checks if it is safe to reorder an add/sext result used in a GEP.
+//
+// An inbounds GEP does not guarantee that the index is non-negative.
+// This helper checks first if the index is known non-negative. If the index is
+// non-negative, the transform is always safe.
+// Second, it checks whether the GEP is inbounds and directly based on a global
+// or an alloca, which are required to prove futher transform validity.
+// If the GEP:
+// - Has a zero offset from the base, the index is non-negative (any negative
+//   value would produce poison/UB)
+// - Has ObjectSize < (2^(N-1) - C + 1) * stride, where C is a constant from the
+//   add, stride is the element size of Idx, and N is bitwidth of Idx.
+//   This is because with this pattern:
+//     %add = add iN %val, C
+//     %sext = sext iN %add to i64
+//     %gep = getelementptr inbounds TYPE, %sext
+//   The worst-case is when %val sign-flips to produce the smallest magnitude
+//   negative value, at 2^(N-1)-1. In this case, the add/sext is -(2^(N-1)-C+1),
+//   and the sext/add is 2^(N-1)+C-1 (2^N difference). The original add/sext
+//   only produces a defined GEP when -(2^(N-1)-C+1) is inbounds. So, if
+//   ObjectSize < (2^(N-1) - C + 1) * stride, it is impossible for the
+//   worst-case sign-flip to be defined.
+//   Note that in this case the GEP is not neccesarily non-negative, but any
+//   negative results will still produce the same behavior in the reordered
+//   version with a defined GEP.
+//   This can also work for negative C, but the threshold is instead
+//   (2^(N-1)+C)*stride, since the sign-flip is done in reverse and is instead
+//   producing a large positive value that still needs to be inbounds to the
+//   object size. If C is negative, we cannot make any useful assumptions based
+//   on the offset, since it would need to be extremely large.
+static bool canReorderAddSextToGEP(const GetElementPtrInst *GEP,
+                                   const Value *Idx,
+                                   const BinaryOperator *Add,
+                                   const DataLayout &DL) {
+  if (isKnownNonNegative(Idx, DL))
+    return true;
+
+  if (!GEP->isInBounds())
+    return false;
+
+  const Value *Ptr = GEP->getPointerOperand();
+  int64_t Offset = 0;
+  const Value *Base = GetPointerBaseWithConstantOffset(
+      const_cast<Value *>(Ptr), Offset, DL);
+
+  // We need one of the operands to be a constant to be able to trace into the
+  // operator.
+  const ConstantInt *CI = dyn_cast<ConstantInt>(Add->getOperand(0));
+  if (!CI)
+    CI = dyn_cast<ConstantInt>(Add->getOperand(1));
+  if (!CI)
+    return false;
+  // Calculate the threshold
+  APInt Threshold;
+  unsigned N = Add->getType()->getIntegerBitWidth();
+  uint64_t Stride =
+      DL.getTypeAllocSize(GEP->getSourceElementType()).getFixedValue();
+  if (!CI->isNegative()) {
+    // (2^(N-1) - C + 1) * stride
+    Threshold = (APInt::getSignedMinValue(N).zext(128) -
+                 CI->getValue().zextOrTrunc(128) + 1) *
+                APInt(128, Stride);
+  } else {
+    // (2^(N-1) + C) * stride
+    Threshold = (APInt::getSignedMinValue(N).zext(128) +
+                 CI->getValue().zextOrTrunc(128)) *
+                APInt(128, Stride);
+  }
+
+  if (Base && (isa<AllocaInst>(Base) || isa<GlobalObject>(Base)) &&
+      !CI->isNegative()) {
+    // If the offset is zero from an alloca or global, inbounds is sufficient to
+    // prove non-negativity if one add operand is non-negative
+    if (Offset == 0)
+      return true;
+
+    // Check if the Offset < Threshold (positive CI only) otherwise
+    if (APInt(128, (uint64_t)Offset).ult(Threshold))
+      return true;
+  } else {
+    // If we can't determine the offset from the base object, we can still use
+    // the underlying object and type size constraints
+    Base = getUnderlyingObject(Ptr);
+    // Can only prove non-negativity if the base object is known
+    if (!(isa<AllocaInst>(Base) || isa<GlobalObject>(Base)))
+      return false;
+  }
+
+  // Check if the ObjectSize < Threshold (for both positive or negative C)
+  uint64_t ObjSize = 0;
+  if (const auto *AI = dyn_cast<AllocaInst>(Base)) {
+    if (auto AllocSize = AI->getAllocationSize(DL))
+      if (!AllocSize->isScalable())
+        ObjSize = AllocSize->getFixedValue();
+  } else if (const auto *GV = dyn_cast<GlobalVariable>(Base)) {
+    ObjSize = DL.getTypeAllocSize(GV->getValueType()).getFixedValue();
+  }
+  if (ObjSize > 0 && APInt(128, ObjSize).ult(Threshold))
+    return true;
+
+  return false;
+}
+
+bool ConstantOffsetExtractor::canTraceInto(bool SignExtended, bool ZeroExtended,
+                                           BinaryOperator *BO,
+                                           GetElementPtrInst *GEP, Value *Idx) {
   // We only consider ADD, SUB and OR, because a non-zero constant found in
   // expressions composed of these operations can be easily hoisted as a
   // constant offset by reassociation.
@@ -513,27 +615,23 @@ bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended, bool ZeroExtended,
   //       1       |      0       | sext(BO) == sext(A) op sext(B)
   //       1       |      1       | zext(sext(BO)) ==
   //               |              |     zext(sext(A)) op zext(sext(B))
-  if (BO->getOpcode() == Instruction::Add && !ZeroExtended && NonNegative) {
+  if (BO->getOpcode() == Instruction::Add && !ZeroExtended && GEP) {
     // If a + b >= 0 and (a >= 0 or b >= 0), then
     //   sext(a + b) = sext(a) + sext(b)
     // even if the addition is not marked nsw.
     //
     // Leveraging this invariant, we can trace into an sext'ed inbound GEP
-    // index if the constant offset is non-negative.
+    // index under certain conditions (see canReorderAddSextToGEP).
     //
     // Verified in @sext_add in split-gep.ll.
-    if (ConstantInt *ConstLHS = dyn_cast<ConstantInt>(LHS)) {
-      if (!ConstLHS->isNegative())
-        return true;
-    }
-    if (ConstantInt *ConstRHS = dyn_cast<ConstantInt>(RHS)) {
-      if (!ConstRHS->isNegative())
-        return true;
-    }
+    if (canReorderAddSextToGEP(GEP, Idx, BO, DL))
+      return true;
   }
 
   // For a sext(add nuw), allow tracing through when the enclosing GEP is both
   // inbounds and nuw.
+  bool GEPInboundsNUW =
+      GEP ? (GEP->isInBounds() && GEP->hasNoUnsignedWrap()) : false;
   if (BO->getOpcode() == Instruction::Add && SignExtended && !ZeroExtended &&
       GEPInboundsNUW && BO->hasNoUnsignedWrap())
     return true;
@@ -557,11 +655,9 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   // Save off the current height of the chain, in case we need to restore it.
   size_t ChainLength = UserChain.size();
 
-  // BO being non-negative does not shed light on whether its operands are
-  // non-negative. Clear the NonNegative and GEPInboundsNUW flags here.
+  // BO cannot use information from the base GEP at this point, so clear it.
   APInt ConstantOffset =
-      find(BO->getOperand(0), SignExtended, ZeroExtended,
-           /* NonNegative */ false, /* GEPInboundsNUW */ false);
+      find(BO->getOperand(0), nullptr, nullptr, SignExtended, ZeroExtended);
   // If we found a constant offset in the left operand, stop and return that.
   // This shortcut might cause us to miss opportunities of combining the
   // constant offsets in both operands, e.g., (a + 4) + (b + 5) => (a + b) + 9.
@@ -573,8 +669,8 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   // since visiting the LHS didn't pan out.
   UserChain.resize(ChainLength);
 
-  ConstantOffset = find(BO->getOperand(1), SignExtended, ZeroExtended,
-                        /* NonNegative */ false, /* GEPInboundsNUW */ false);
+  ConstantOffset = find(BO->getOperand(1), nullptr, nullptr, SignExtended,
+                        ZeroExtended);
   // If U is a sub operator, negate the constant offset found in the right
   // operand.
   if (BO->getOpcode() == Instruction::Sub)
@@ -587,9 +683,9 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   return ConstantOffset;
 }
 
-APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
-                                    bool ZeroExtended, bool NonNegative,
-                                    bool GEPInboundsNUW) {
+APInt ConstantOffsetExtractor::find(Value *V, GetElementPtrInst *GEP,
+                                    Value *Idx, bool SignExtended,
+                                    bool ZeroExtended) {
   // TODO(jingyue): We could trace into integer/pointer casts, such as
   // inttoptr, ptrtoint, bitcast, and addrspacecast. We choose to handle only
   // integers because it gives good enough results for our benchmarks.
@@ -605,28 +701,21 @@ APInt ConstantOffsetExtractor::find(Value *V, bool SignExtended,
     ConstantOffset = CI->getValue();
   } else if (BinaryOperator *BO = dyn_cast<BinaryOperator>(V)) {
     // Trace into subexpressions for more hoisting opportunities.
-    if (CanTraceInto(SignExtended, ZeroExtended, BO, NonNegative,
-                     GEPInboundsNUW))
+    if (canTraceInto(SignExtended, ZeroExtended, BO, GEP, Idx))
       ConstantOffset = findInEitherOperand(BO, SignExtended, ZeroExtended);
   } else if (isa<TruncInst>(V)) {
-    ConstantOffset = find(U->getOperand(0), SignExtended, ZeroExtended,
-                          NonNegative, GEPInboundsNUW)
+    ConstantOffset = find(U->getOperand(0), GEP, Idx, SignExtended,
+                          ZeroExtended)
                          .trunc(BitWidth);
   } else if (isa<SExtInst>(V)) {
-    ConstantOffset = find(U->getOperand(0), /* SignExtended */ true,
-                          ZeroExtended, NonNegative, GEPInboundsNUW)
+    ConstantOffset = find(U->getOperand(0), GEP, Idx, /* SignExtended */ true,
+                          ZeroExtended)
                          .sext(BitWidth);
   } else if (isa<ZExtInst>(V)) {
     // As an optimization, we can clear the SignExtended flag because
     // sext(zext(a)) = zext(a). Verified in @sext_zext in split-gep.ll.
-    //
-    // We can also clear the GEPInboundsNUW flag because it is not required to
-    // trace into zext.
-    //
-    // Clear the NonNegative flag, because zext(a) >= 0 does not imply a >= 0.
-    ConstantOffset = find(U->getOperand(0), /* SignExtended */ false,
-                          /* ZeroExtended */ true, /* NonNegative */ false,
-                          /* GEPInboundsNUW */ false)
+    ConstantOffset = find(U->getOperand(0), GEP, Idx, /* SignExtended */ false,
+                          /* ZeroExtended */ true)
                          .zext(BitWidth);
   }
 
@@ -791,40 +880,14 @@ static bool allowsPreservingNUW(const User *U) {
   return true;
 }
 
-// An inbounds GEP does not guarantee that the index is non-negative.
-// This helper checks first if the index is known non-negative. If it is not,
-// it checks whether the GEP is directly based on a global or an alloca with
-// zero offset, in which case inbounds is sufficient to prove non-negativity.
-static bool isGEPNonNegative(const GetElementPtrInst *GEP, const Value *Idx,
-                             const DataLayout &DL) {
-  if (isKnownNonNegative(Idx, DL))
-    return true;
-
-  if (!GEP->isInBounds())
-    return false;
-
-  const Value *Ptr = GEP->getPointerOperand();
-  int64_t Offset = 0;
-  const Value *Base = GetPointerBaseWithConstantOffset(
-      const_cast<Value *>(Ptr), Offset, DL, /*AllowNonInbounds=*/false);
-
-  if (!Base || Offset != 0)
-    return false;
-
-  return isa<AllocaInst>(Base) || isa<GlobalObject>(Base);
-}
-
 Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
                                         User *&UserChainTail,
                                         bool &PreservesNUW) {
   ConstantOffsetExtractor Extractor(GEP->getIterator());
-  bool GEPNonNegative = isGEPNonNegative(GEP, Idx, Extractor.DL);
-  bool GEPInboundsNUW = GEP->isInBounds() && GEP->hasNoUnsignedWrap();
   // Find a non-zero constant offset first.
   APInt ConstantOffset =
-      Extractor.find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-                     /* NonNegative */ GEPNonNegative,
-                     /* GEPInboundsNUW */ GEPInboundsNUW);
+      Extractor.find(Idx, GEP, Idx, /* SignExtended */ false,
+                     /* ZeroExtended */ false);
   if (ConstantOffset == 0) {
     UserChainTail = nullptr;
     PreservesNUW = true;
@@ -840,12 +903,8 @@ Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
 }
 
 APInt ConstantOffsetExtractor::Find(Value *Idx, GetElementPtrInst *GEP) {
-  bool GEPNonNegative = isGEPNonNegative(GEP, Idx, GEP->getDataLayout());
-  bool GEPInboundsNUW = GEP->isInBounds() && GEP->hasNoUnsignedWrap();
   return ConstantOffsetExtractor(GEP->getIterator())
-      .find(Idx, /* SignExtended */ false, /* ZeroExtended */ false,
-            /* NonNegative */ GEPNonNegative,
-            /* GEPInboundsNUW */ GEPInboundsNUW);
+      .find(Idx, GEP, Idx, /* SignExtended */ false, /* ZeroExtended */ false);
 }
 
 bool SeparateConstOffsetFromGEP::canonicalizeArrayIndicesToIndexSize(
diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
index 39da11c8339ca..3e76472f62407 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/NVPTX/split-gep.ll
@@ -37,11 +37,10 @@ define ptr @sext_add(i32 %i, i32 %j) {
 ; CHECK-LABEL: define ptr @sext_add(
 ; CHECK-SAME: i32 [[I:%.*]], i32 [[J:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[J]], -2
-; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[TMP0]] to i64
-; CHECK-NEXT:    [[TMP2:%.*]] = sext i32 [[I]] to i64
-; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr [32 x [32 x float]], ptr @float_2d_array, i64 0, i64 [[TMP2]], i64 [[TMP1]]
-; CHECK-NEXT:    [[P1:%.*]] = getelementptr i8, ptr [[TMP3]], i64 128
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[J]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr [32 x [32 x float]], ptr @float_2d_array, i64 0, i64 [[TMP0]], i64 [[TMP1]]
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr i8, ptr [[TMP2]], i64 120
 ; CHECK-NEXT:    ret ptr [[P1]]
 ;
 entry:
@@ -49,6 +48,8 @@ entry:
   %1 = sext i32 %0 to i64  ; inbound sext(i + 1) = sext(i) + 1
   %2 = add i32 %j, -2
   ; However, inbound sext(j + -2) != sext(j) + -2, e.g., j = INT_MIN
+  ; But j = INT_MIN would result in a very large positive result which would be
+  ; OOB (and produce poison), so there is no counter example in this case
   %3 = sext i32 %2 to i64
   %p = getelementptr inbounds [32 x [32 x float]], ptr @float_2d_array, i64 0, i64 %1, i64 %3
   ret ptr %p
@@ -94,22 +95,108 @@ entry:
   ret ptr %p
 }
 
-; We cannot trace into sext(a + b) if a + b is an inbounds GEP but not a zero
-; offset from a known base ptr even if one of a or b is non-negative.
-define ptr @sext_add_nonzerooffset(i4 %i) {
-; CHECK-LABEL: define ptr @sext_add_nonzerooffset(
-; CHECK-SAME: i4 [[I:%.*]]) {
+; We can trace into sext(a + b) if a + b is an inbounds GEP and the known
+; offset from a known base ptr is within a certain threshold relative to the
+; bitwidth of the index (offset < (2^(n-1) - C + 1) * bitwidth).
+define ptr @sext_add_nonzerooffset_inrange(i8 %i, i64 %size) {
+; CHECK-LABEL: define ptr @sext_add_nonzerooffset_inrange(
+; CHECK-SAME: i8 [[I:%.*]], i64 [[SIZE:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = add i4 [[I]], 1
-; CHECK-NEXT:    [[TMP1:%.*]] = sext i4 [[TMP0]] to i64
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr float, ptr @float_array, i64 [[TMP1]]
-; CHECK-NEXT:    [[P1:%.*]] = getelementptr [128 x float], ptr [[GEP1]], i64 0, i64 64
-; CHECK-NEXT:    ret ptr [[P1]]
+; CHECK-NEXT:    [[ARR:%.*]] = alloca float, i64 %size, align 4
+; CHECK-NEXT:    [[OFFSETARR:%.*]] = getelementptr float, ptr [[ARR]], i64 127
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i8 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr float, ptr [[OFFSETARR]], i64 [[TMP0]]
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr i8, ptr [[TMP1]], i64 4
+; CHECK-NEXT:    ret ptr [[P2]]
+;
+entry:
+  %arr = alloca float, i64 %size, align 4
+  %offsetarr = getelementptr float, ptr %arr, i64 127
+  %add = add i8 %i, 1
+  %sext = sext i8 %add to i64
+  %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
+  ret ptr %p
+}
+
+define ptr @sext_add_nonzerooffset_outofrange(i8 %i, i64 %size) {
+; CHECK-LABEL: define ptr @sext_add_nonzerooffset_outofrange(
+; CHECK-SAME: i8 [[I:%.*]], i64 [[SIZE:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARR:%.*]] = alloca float, i64 [[SIZE]], align 4
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[I]], 1
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i8 [[ADD]] to i64
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr float, ptr [[ARR]], i64 [[SEXT]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr float, ptr [[TMP0]], i64 128
+; CHECK-NEXT:    ret ptr [[TMP1]]
+;
+entry:
+  %arr = alloca float, i64 %size, align 4
+  %offsetarr = getelementptr float, ptr %arr, i64 128
+  %add = add i8 %i, 1
+  %sext = sext i8 %add to i64
+  %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
+  ret ptr %p
+}
+
+; We can trace into sext(a + b) if a + b is an inbounds GEP and the size of the
+; known base ptr is within a certain threshold relative to the bitwidth of the
+; index (offset < (2^(n-1) - C + 1) * bitwidth).
+define ptr @sext_add_unknownoffset_inrange(i8 %i, i64 %off) {
+; CHECK-LABEL: define ptr @sext_add_unknownoffset_inrange(
+; CHECK-SAME: i8 [[I:%.*]], i64 [[OFF:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARR:%.*]] = alloca float, i64 126, align 4
+; CHECK-NEXT:    [[OFFSETARR:%.*]] = getelementptr float, ptr [[ARR]], i64 [[OFF]]
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i8 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr float, ptr [[OFFSETARR]], i64 [[TMP0]]
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr i8, ptr [[TMP1]], i64 8
+; CHECK-NEXT:    ret ptr [[P2]]
+;
+entry:
+  %arr = alloca float, i64 126, align 4
+  %offsetarr = getelementptr float, ptr %arr, i64 %off
+  %add = add i8 %i, 2
+  %sext = sext i8 %add to i64
+  %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
+  ret ptr %p
+}
+
+define ptr @sext_add_unknownoffset_inrange_neg(i8 %i, i64 %off) {
+; CHECK-LABEL: define ptr @sext_add_unknownoffset_inrange_neg(
+; CHECK-SAME: i8 [[I:%.*]], i64 [[OFF:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARR:%.*]] = alloca float, i64 125, align 4
+; CHECK-NEXT:    [[OFFSETARR:%.*]] = getelementptr float, ptr [[ARR]], i64 [[OFF]]
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i8 [[I]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr float, ptr [[OFFSETARR]], i64 [[TMP0]]
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr i8, ptr [[TMP1]], i64 -8
+; CHECK-NEXT:    ret ptr [[P2]]
+;
+entry:
+  %arr = alloca float, i64 125, align 4
+  %offsetarr = getelementptr float, ptr %arr, i64 %off
+  %add = add i8 %i, -2
+  %sext = sext i8 %add to i64
+  %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
+  ret ptr %p
+}
+
+define ptr @sext_add_unknownoffset_outofrange(i8 %i, i64 %off) {
+; CHECK-LABEL: define ptr @sext_add_unknownoffset_outofrange(
+; CHECK-SAME: i8 [[I:%.*]], i64 [[OFF:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARR:%.*]] = alloca float, i64 127, align 4
+; CHECK-NEXT:    [[OFFSETARR:%.*]] = getelementptr float, ptr [[ARR]], i64 [[OFF]]
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[I]], 2
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i8 [[ADD]] to i64
+; CHECK-NEXT:    [[P:%.*]] = getelementptr inbounds float, ptr [[OFFSETARR]], i64 [[SEXT]]
+; CHECK-NEXT:    ret ptr [[P]]
 ;
 entry:
-  %offsetarr = getelementptr inbounds [128 x float], ptr @float_array, i64 0, i64 64
-  %add = add i4 %i, 1
-  %sext = sext i4 %add to i64
+  %arr = alloca float, i64 127, align 4
+  %offsetarr = getelementptr float, ptr %arr, i64 %off
+  %add = add i8 %i, 2
+  %sext = sext i8 %add to i64
   %p = getelementptr inbounds float, ptr %offsetarr, i64 %sext
   ret ptr %p
 }

>From ba8401a5886a3e9d80914c6635f00870f6595c72 Mon Sep 17 00:00:00 2001
From: Meredith Julian <mjulian at nvidia.com>
Date: Fri, 6 Mar 2026 15:46:43 -0800
Subject: [PATCH 8/8] fix clang format

---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 28 +++++++++----------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index a050cdbb99301..08494b6b3ae4b 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -507,8 +507,7 @@ FunctionPass *llvm::createSeparateConstOffsetFromGEPPass(bool LowerGEP) {
 //   object size. If C is negative, we cannot make any useful assumptions based
 //   on the offset, since it would need to be extremely large.
 static bool canReorderAddSextToGEP(const GetElementPtrInst *GEP,
-                                   const Value *Idx,
-                                   const BinaryOperator *Add,
+                                   const Value *Idx, const BinaryOperator *Add,
                                    const DataLayout &DL) {
   if (isKnownNonNegative(Idx, DL))
     return true;
@@ -518,8 +517,8 @@ static bool canReorderAddSextToGEP(const GetElementPtrInst *GEP,
 
   const Value *Ptr = GEP->getPointerOperand();
   int64_t Offset = 0;
-  const Value *Base = GetPointerBaseWithConstantOffset(
-      const_cast<Value *>(Ptr), Offset, DL);
+  const Value *Base =
+      GetPointerBaseWithConstantOffset(const_cast<Value *>(Ptr), Offset, DL);
 
   // We need one of the operands to be a constant to be able to trace into the
   // operator.
@@ -669,8 +668,8 @@ APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
   // since visiting the LHS didn't pan out.
   UserChain.resize(ChainLength);
 
-  ConstantOffset = find(BO->getOperand(1), nullptr, nullptr, SignExtended,
-                        ZeroExtended);
+  ConstantOffset =
+      find(BO->getOperand(1), nullptr, nullptr, SignExtended, ZeroExtended);
   // If U is a sub operator, negate the constant offset found in the right
   // operand.
   if (BO->getOpcode() == Instruction::Sub)
@@ -704,13 +703,13 @@ APInt ConstantOffsetExtractor::find(Value *V, GetElementPtrInst *GEP,
     if (canTraceInto(SignExtended, ZeroExtended, BO, GEP, Idx))
       ConstantOffset = findInEitherOperand(BO, SignExtended, ZeroExtended);
   } else if (isa<TruncInst>(V)) {
-    ConstantOffset = find(U->getOperand(0), GEP, Idx, SignExtended,
-                          ZeroExtended)
-                         .trunc(BitWidth);
+    ConstantOffset =
+        find(U->getOperand(0), GEP, Idx, SignExtended, ZeroExtended)
+            .trunc(BitWidth);
   } else if (isa<SExtInst>(V)) {
-    ConstantOffset = find(U->getOperand(0), GEP, Idx, /* SignExtended */ true,
-                          ZeroExtended)
-                         .sext(BitWidth);
+    ConstantOffset =
+        find(U->getOperand(0), GEP, Idx, /* SignExtended */ true, ZeroExtended)
+            .sext(BitWidth);
   } else if (isa<ZExtInst>(V)) {
     // As an optimization, we can clear the SignExtended flag because
     // sext(zext(a)) = zext(a). Verified in @sext_zext in split-gep.ll.
@@ -885,9 +884,8 @@ Value *ConstantOffsetExtractor::Extract(Value *Idx, GetElementPtrInst *GEP,
                                         bool &PreservesNUW) {
   ConstantOffsetExtractor Extractor(GEP->getIterator());
   // Find a non-zero constant offset first.
-  APInt ConstantOffset =
-      Extractor.find(Idx, GEP, Idx, /* SignExtended */ false,
-                     /* ZeroExtended */ false);
+  APInt ConstantOffset = Extractor.find(Idx, GEP, Idx, /* SignExtended */ false,
+                                        /* ZeroExtended */ false);
   if (ConstantOffset == 0) {
     UserChainTail = nullptr;
     PreservesNUW = true;



More information about the llvm-commits mailing list