[llvm] [SeparateConstOffsetFromGEP] Stop distributing sext/zext over lossy trunc (PR #221381)

Fujun Han via llvm-commits llvm-commits at lists.llvm.org
Wed Sep 9 04:08:35 PDT 2026


https://github.com/Peter9606 updated https://github.com/llvm/llvm-project/pull/221381

>From 58b43cb04cc28280951c0b4bc092b1a1b7c2c751 Mon Sep 17 00:00:00 2001
From: Fujun Han <fujun.han at iluvatar.com>
Date: Sat, 5 Sep 2026 10:06:57 +0800
Subject: [PATCH 1/2] [SeparateConstOffsetFromGEP] Precommit tests for
 sext/zext of trunc (NFC)

The nuw/nsw flags on an add hold at the width of the add, which says
nothing about wrapping at the width of a truncation applied to its
result. These tests capture the current behavior of distributing a
sext/zext above a truncation into the operands of the add; the lossy
cases are miscompiled today.

@sext_of_lossy_operand is the review counter-example: the sum survives
the truncation, but the remaining operand does not, so hoisting the
constant is still wrong.

Assisted-by: Cursor (Claude Fable 5)
Signed-off-by: Fujun Han <fujun.han at iluvatar.com>
Co-authored-by: Cursor <cursoragent at cursor.com>
---
 .../ext-of-trunc-add-wrap.ll                  | 142 ++++++++++++++++++
 1 file changed, 142 insertions(+)
 create mode 100644 llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll

diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll
new file mode 100644
index 0000000000000..94f5b6c1208c6
--- /dev/null
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll
@@ -0,0 +1,142 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes='separate-const-offset-from-gep<lower-gep>' < %s | FileCheck %s
+
+; The nuw/nsw flags on an add hold at the width of the add, which says nothing
+; about wrapping at the width of a truncation applied to its result.  A
+; sext/zext above such a truncation must therefore not be distributed into the
+; operands of the add.
+
+; zext i64 (trunc i8 (add nuw i32 A, B)) wraps modulo 256, so the constant 1
+; must stay inside the truncation.  For A = 251 and B = 5 the index is 0, while
+; distributing the casts would give 251 + 5 = 256.
+define ptr @zext_of_lossy_trunc(ptr %p, i8 %a, i64 %iv) {
+; CHECK-LABEL: define ptr @zext_of_lossy_trunc(
+; CHECK-SAME: ptr [[P:%.*]], i8 [[A:%.*]], i64 [[IV:%.*]]) {
+; CHECK-NEXT:    [[AZ:%.*]] = zext i8 [[A]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[AZ]] to i8
+; CHECK-NEXT:    [[TMP2:%.*]] = zext nneg i8 [[TMP1]] to i64
+; CHECK-NEXT:    [[T:%.*]] = trunc i64 [[IV]] to i8
+; CHECK-NEXT:    [[TZ:%.*]] = zext i8 [[T]] to i32
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[TZ]] to i8
+; CHECK-NEXT:    [[IDX:%.*]] = zext nneg i8 [[SUM8]] to i64
+; CHECK-NEXT:    [[SUM2:%.*]] = add i64 [[TMP2]], [[IDX]]
+; CHECK-NEXT:    [[TMP7:%.*]] = shl i64 [[SUM2]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP7]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 4
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %iv.next = add nuw nsw i64 %iv, 1
+  %t = trunc i64 %iv.next to i8
+  %az = zext i8 %a to i32
+  %tz = zext i8 %t to i32
+  %sum = add nuw nsw i32 %az, %tz
+  %sum8 = trunc i32 %sum to i8
+  %idx = zext nneg i8 %sum8 to i64
+  %q = getelementptr inbounds i32, ptr %p, i64 %idx
+  ret ptr %q
+}
+
+; Same shape, but masking keeps both the sum and the remaining operand below
+; 256, so the truncation loses nothing and the constant can still be hoisted.
+define ptr @zext_of_lossless_trunc(ptr %p, i32 %x) {
+; CHECK-LABEL: define ptr @zext_of_lossless_trunc(
+; CHECK-SAME: ptr [[P:%.*]], i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = and i32 [[X]], 15
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[IDX:%.*]] = zext nneg i8 [[SUM8]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 12
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %a = and i32 %x, 15
+  %sum = add nuw nsw i32 %a, 3
+  %sum8 = trunc i32 %sum to i8
+  %idx = zext nneg i8 %sum8 to i64
+  %q = getelementptr inbounds i32, ptr %p, i64 %idx
+  ret ptr %q
+}
+
+; sext additionally needs the sign bit of the narrow value to be clear.  Here
+; the sum reaches 130, so sext i8 of the truncation is negative and the
+; constant must stay inside the truncation.
+define ptr @sext_of_lossy_trunc(ptr %p, i32 %x) {
+; CHECK-LABEL: define ptr @sext_of_lossy_trunc(
+; CHECK-SAME: ptr [[P:%.*]], i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = and i32 [[X]], 127
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[SUM8]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 12
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %a = and i32 %x, 127
+  %sum = add nuw nsw i32 %a, 3
+  %sum8 = trunc i32 %sum to i8
+  %idx = sext i8 %sum8 to i64
+  %q = getelementptr i32, ptr %p, i64 %idx
+  ret ptr %q
+}
+
+; With the sign bits of the sum and of the remaining operand both clear, sext
+; and zext agree on every involved value, so the constant can be hoisted out
+; of a sext'ed truncation as well.
+define ptr @sext_of_lossless_trunc(ptr %p, i32 %x) {
+; CHECK-LABEL: define ptr @sext_of_lossless_trunc(
+; CHECK-SAME: ptr [[P:%.*]], i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = and i32 [[X]], 63
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[SUM8]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 12
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %a = and i32 %x, 63
+  %sum = add nuw nsw i32 %a, 3
+  %sum8 = trunc i32 %sum to i8
+  %idx = sext i8 %sum8 to i64
+  %q = getelementptr i32, ptr %p, i64 %idx
+  ret ptr %q
+}
+
+; The trunc of the sum being lossless is not enough: here %sum is 124 or 127,
+; so trunc+sext reproduces it, but the remaining operand %a is 224 or 227 and
+; wraps at i8.  Distributing the sext would compute sext(trunc(%a)) - 100 =
+; -132 instead of 124.  Counter-example from review of llvm#221381.
+define ptr @sext_of_lossy_operand(ptr %p, i1 %c) {
+; CHECK-LABEL: define ptr @sext_of_lossy_operand(
+; CHECK-SAME: ptr [[P:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = select i1 [[C]], i32 224, i32 227
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[SUM8]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 -400
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %a = select i1 %c, i32 224, i32 227
+  %sum = add nsw i32 %a, -100
+  %sum8 = trunc i32 %sum to i8
+  %idx = sext i8 %sum8 to i64
+  %q = getelementptr i32, ptr %p, i64 %idx
+  ret ptr %q
+}
+
+; A truncation with no extension above it still distributes over the add,
+; because truncation is exact in modular arithmetic.
+define ptr @bare_trunc(ptr %p, i128 %i) {
+; CHECK-LABEL: define ptr @bare_trunc(
+; CHECK-SAME: ptr [[P:%.*]], i128 [[I:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i128 [[I]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i64 [[TMP1]], 2
+; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP2]]
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 4
+; CHECK-NEXT:    ret ptr [[Q]]
+;
+  %idx = add i128 %i, 1
+  %idx.conv = trunc i128 %idx to i64
+  %q = getelementptr i32, ptr %p, i64 %idx.conv
+  ret ptr %q
+}

>From d28138f9d914902930c0887cb0b5ed3cbfd2893d Mon Sep 17 00:00:00 2001
From: Fujun Han <fujun.han at iluvatar.com>
Date: Wed, 9 Sep 2026 19:08:02 +0800
Subject: [PATCH 2/2] [SeparateConstOffsetFromGEP] Stop distributing sext/zext
 over lossy trunc

find() traces through a trunc while keeping a pending sext/zext, and
canTraceInto() then justifies distributing the extension into the
operands of an add/sub by checking its nuw/nsw flags. Those flags
hold at the width of the add but say nothing about wrapping at the
truncation width, so

  zext i64 (trunc i8 (add nuw i32 (zext i8 251), (zext i8 5)))

which is 0 was rebuilt as 251 + 5 = 256.

With a pending sext/zext, only trace through a trunc when the
truncated value is entirely constant (a plain constant fold), or when
it is an add with an immediate constant operand and known bits prove
that both the sum and the remaining operand survive the trunc+ext
round trip. The constant then equals sum - operand exactly and needs
no proof of its own beyond being non-negative when the value is
zero-extended.

Proving the round trip for the sum alone is not enough: in the review
counter-example, %sum = add nsw (select i1 %c, i32 224, i32 227), -100
is 124 or 127 and truncates to i8 losslessly, but the remaining
operand 224/227 wraps at i8, so hoisting -100 would rebuild the index
as sext(trunc(%a)) - 100 = -132 instead of 124.

Assisted-by: Cursor (Claude Fable 5)
Signed-off-by: Fujun Han <fujun.han at iluvatar.com>
Co-authored-by: Cursor <cursoragent at cursor.com>
---
 .../Scalar/SeparateConstOffsetFromGEP.cpp     | 71 ++++++++++++++++++-
 .../ext-of-trunc-add-wrap.ll                  | 27 +++----
 2 files changed, 79 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 4870b8c888279..7d245d8e2df3c 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -298,6 +298,13 @@ class ConstantOffsetExtractor {
   bool canTraceInto(bool SignExtended, bool ZeroExtended, BinaryOperator *BO,
                     GetElementPtrInst *GEP, Value *Idx);
 
+  /// A helper function that returns whether truncating \p V to \p BitWidth and
+  /// extending it back reproduces \p V, so that a sext/zext applied after the
+  /// truncation can still be reasoned about at the width of \p V.
+  ///
+  /// \p Signed Whether the value is extended back with sext rather than zext.
+  bool truncPreservesValue(Value *V, unsigned BitWidth, bool Signed);
+
   /// Analyze a xor expression, and identify the bits in the constant operand
   /// that are disjoint from the base operand's known set bits. For these
   /// disjoint bits, a xor is equivalent to an addition, which allows us to
@@ -684,6 +691,18 @@ bool ConstantOffsetExtractor::canTraceInto(bool SignExtended, bool ZeroExtended,
   return true;
 }
 
+bool ConstantOffsetExtractor::truncPreservesValue(Value *V, unsigned BitWidth,
+                                                  bool Signed) {
+  unsigned SrcBits = V->getType()->getScalarSizeInBits();
+  if (SrcBits <= BitWidth)
+    return true;
+
+  // Extending back with sext additionally requires the sign bit of the
+  // truncated value to be clear, so that sext and zext agree.
+  unsigned KeepBits = Signed ? BitWidth - 1 : BitWidth;
+  return computeKnownBits(V, DL).countMinLeadingZeros() >= SrcBits - KeepBits;
+}
+
 APInt ConstantOffsetExtractor::findInEitherOperand(BinaryOperator *BO,
                                                    bool SignExtended,
                                                    bool ZeroExtended) {
@@ -741,9 +760,55 @@ APInt ConstantOffsetExtractor::find(Value *V, GetElementPtrInst *GEP,
     else if (BO->getOpcode() == Instruction::Xor)
       ConstantOffset = extractDisjointBitsFromXor(BO);
   } else if (isa<TruncInst>(V)) {
-    ConstantOffset =
-        find(U->getOperand(0), GEP, Idx, SignExtended, ZeroExtended)
-            .trunc(BitWidth);
+    Value *TruncOp = U->getOperand(0);
+    if (!SignExtended && !ZeroExtended) {
+      // With no pending extension, truncation distributes over add/sub in
+      // modular arithmetic, so any constant found in the wider operand stays
+      // valid after truncating it.
+      ConstantOffset =
+          find(TruncOp, GEP, Idx, SignExtended, ZeroExtended).trunc(BitWidth);
+    } else if (isa<ConstantInt>(TruncOp)) {
+      // The truncated value is entirely constant, so there is no remainder
+      // and the pending extensions apply to the constant itself; extracting
+      // it is a plain constant fold.
+      ConstantOffset =
+          find(TruncOp, GEP, Idx, SignExtended, ZeroExtended).trunc(BitWidth);
+    } else if (auto *BO = dyn_cast<BinaryOperator>(TruncOp)) {
+      // A pending sext/zext applies to the truncated value, so splitting a
+      // constant C out of the truncation rewrites ext(trunc(a + C)) into
+      // ext(trunc(a)) + C.  That is only sound when a, a + C, and C each
+      // survive the trunc+ext round trip.  nuw/nsw on the add cannot prove
+      // this: the flags hold at the width of the add and say nothing about
+      // wrapping at the truncation width.  Instead, restrict the search to a
+      // constant that is an immediate operand of the truncated add, so that
+      // known bits can prove the round trip for both the sum and the
+      // remaining operand; C = sum - operand then needs no proof of its own
+      // beyond being non-negative when the value is zero-extended.
+      //
+      // Checking only the sum is not enough.  In
+      //   %a   = select i1 %c, i32 224, i32 227
+      //   %sum = add nsw i32 %a, -100        ; 124 or 127, trunc is lossless
+      //   %idx = sext i8 (trunc i8 %sum) to i64
+      // the sum survives the round trip but %a does not: hoisting -100 would
+      // rebuild the index as sext(trunc(%a)) - 100 = -132 instead of 124.
+      ConstantInt *CI = nullptr;
+      Value *Remainder = nullptr;
+      if (BO->getOpcode() == Instruction::Add &&
+          canTraceInto(SignExtended, ZeroExtended, BO, GEP, Idx)) {
+        if ((CI = dyn_cast<ConstantInt>(BO->getOperand(1))))
+          Remainder = BO->getOperand(0);
+        else if ((CI = dyn_cast<ConstantInt>(BO->getOperand(0))))
+          Remainder = BO->getOperand(1);
+      }
+      if (CI && !CI->isZero() &&
+          (!ZeroExtended || CI->getValue().isNonNegative()) &&
+          truncPreservesValue(TruncOp, BitWidth, SignExtended) &&
+          truncPreservesValue(Remainder, BitWidth, SignExtended)) {
+        UserChain.push_back(CI);
+        UserChain.push_back(BO);
+        ConstantOffset = CI->getValue().trunc(BitWidth);
+      }
+    }
   } else if (isa<SExtInst>(V)) {
     ConstantOffset =
         find(U->getOperand(0), GEP, Idx, /* SignExtended */ true, ZeroExtended)
diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll
index 94f5b6c1208c6..092ac833e62f7 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/ext-of-trunc-add-wrap.ll
@@ -12,17 +12,14 @@
 define ptr @zext_of_lossy_trunc(ptr %p, i8 %a, i64 %iv) {
 ; CHECK-LABEL: define ptr @zext_of_lossy_trunc(
 ; CHECK-SAME: ptr [[P:%.*]], i8 [[A:%.*]], i64 [[IV:%.*]]) {
+; CHECK-NEXT:    [[IV_NEXT:%.*]] = add nuw nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[T:%.*]] = trunc i64 [[IV_NEXT]] to i8
 ; CHECK-NEXT:    [[AZ:%.*]] = zext i8 [[A]] to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[AZ]] to i8
-; CHECK-NEXT:    [[TMP2:%.*]] = zext nneg i8 [[TMP1]] to i64
-; CHECK-NEXT:    [[T:%.*]] = trunc i64 [[IV]] to i8
 ; CHECK-NEXT:    [[TZ:%.*]] = zext i8 [[T]] to i32
-; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[TZ]] to i8
+; CHECK-NEXT:    [[SUM:%.*]] = add nuw nsw i32 [[AZ]], [[TZ]]
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[SUM]] to i8
 ; CHECK-NEXT:    [[IDX:%.*]] = zext nneg i8 [[SUM8]] to i64
-; CHECK-NEXT:    [[SUM2:%.*]] = add i64 [[TMP2]], [[IDX]]
-; CHECK-NEXT:    [[TMP7:%.*]] = shl i64 [[SUM2]], 2
-; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP7]]
-; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 4
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr inbounds i32, ptr [[P]], i64 [[IDX]]
 ; CHECK-NEXT:    ret ptr [[Q]]
 ;
   %iv.next = add nuw nsw i64 %iv, 1
@@ -64,11 +61,10 @@ define ptr @sext_of_lossy_trunc(ptr %p, i32 %x) {
 ; CHECK-LABEL: define ptr @sext_of_lossy_trunc(
 ; CHECK-SAME: ptr [[P:%.*]], i32 [[X:%.*]]) {
 ; CHECK-NEXT:    [[A:%.*]] = and i32 [[X]], 127
-; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[SUM:%.*]] = add nuw nsw i32 [[A]], 3
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[SUM]] to i8
 ; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[SUM8]] to i64
-; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
-; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
-; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 12
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i32, ptr [[P]], i64 [[IDX]]
 ; CHECK-NEXT:    ret ptr [[Q]]
 ;
   %a = and i32 %x, 127
@@ -109,11 +105,10 @@ define ptr @sext_of_lossy_operand(ptr %p, i1 %c) {
 ; CHECK-LABEL: define ptr @sext_of_lossy_operand(
 ; CHECK-SAME: ptr [[P:%.*]], i1 [[C:%.*]]) {
 ; CHECK-NEXT:    [[A:%.*]] = select i1 [[C]], i32 224, i32 227
-; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[A]] to i8
+; CHECK-NEXT:    [[SUM:%.*]] = add nsw i32 [[A]], -100
+; CHECK-NEXT:    [[SUM8:%.*]] = trunc i32 [[SUM]] to i8
 ; CHECK-NEXT:    [[IDX:%.*]] = sext i8 [[SUM8]] to i64
-; CHECK-NEXT:    [[TMP3:%.*]] = shl i64 [[IDX]], 2
-; CHECK-NEXT:    [[UGLYGEP:%.*]] = getelementptr i8, ptr [[P]], i64 [[TMP3]]
-; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[UGLYGEP]], i64 -400
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr i32, ptr [[P]], i64 [[IDX]]
 ; CHECK-NEXT:    ret ptr [[Q]]
 ;
   %a = select i1 %c, i32 224, i32 227



More information about the llvm-commits mailing list