[llvm] [ConstraintElim] Bail out on non-canonical GEPs (PR #156688)

[Yingwei Zheng via llvm-commits]

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

In most cases, GEPs should be canonicalized by InstCombine. Bail out on non-canonical forms for simplicity.
Fixes https://github.com/llvm/llvm-project/pull/155253#issuecomment-3248457478.


>From 02036ece8eeea15b5f6c9dcbb6e01d89ca64ccc5 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Wed, 3 Sep 2025 23:07:20 +0800
Subject: [PATCH] [ConstraintElim] Bail out on non-canonical GEPs

---
 .../Scalar/ConstraintElimination.cpp           |  8 ++++++--
 .../implied-by-bounded-memory-access.ll        | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 1b4d8c786cbb2..4acc3f2d84690 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1125,6 +1125,12 @@ static bool getConstraintFromMemoryAccess(GetElementPtrInst &GEP,
   if (Offset.VariableOffsets.size() != 1)
     return false;
 
+  uint64_t BitWidth = Offset.ConstantOffset.getBitWidth();
+  auto &[Index, Scale] = Offset.VariableOffsets.front();
+  // Bail out on non-canonical GEPs.
+  if (Index->getType()->getScalarSizeInBits() != BitWidth)
+    return false;
+
   ObjectSizeOpts Opts;
   // Workaround for gep inbounds, ptr null, idx.
   Opts.NullIsUnknownSize = true;
@@ -1140,8 +1146,6 @@ static bool getConstraintFromMemoryAccess(GetElementPtrInst &GEP,
   // With nuw flag, we know that the index addition doesn't have unsigned wrap.
   // If (AllocSize - (ConstOffset + AccessSize)) wraps around, there is no valid
   // value for Index.
-  uint64_t BitWidth = Offset.ConstantOffset.getBitWidth();
-  auto &[Index, Scale] = Offset.VariableOffsets.front();
   APInt MaxIndex = (APInt(BitWidth, Size->getFixedValue() - AccessSize,
                           /*isSigned=*/false, /*implicitTrunc=*/true) -
                     Offset.ConstantOffset)
diff --git a/llvm/test/Transforms/ConstraintElimination/implied-by-bounded-memory-access.ll b/llvm/test/Transforms/ConstraintElimination/implied-by-bounded-memory-access.ll
index 8e3862b5714d0..90caf52d04164 100644
--- a/llvm/test/Transforms/ConstraintElimination/implied-by-bounded-memory-access.ll
+++ b/llvm/test/Transforms/ConstraintElimination/implied-by-bounded-memory-access.ll
@@ -371,3 +371,21 @@ define i8 @load_from_null(i64 %idx) {
   %add = add i8 %load, %zext
   ret i8 %add
 }
+
+define i8 @load_global_non_canonical_gep(i32 %idx) {
+; CHECK-LABEL: define i8 @load_global_non_canonical_gep(
+; CHECK-SAME: i32 [[IDX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr nuw i8, ptr @g, i32 [[IDX]]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[GEP]], align 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[IDX]], 5
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i1 [[CMP]] to i8
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[LOAD]], [[ZEXT]]
+; CHECK-NEXT:    ret i8 [[ADD]]
+;
+  %gep = getelementptr nuw i8, ptr @g, i32 %idx
+  %load = load i8, ptr %gep
+  %cmp = icmp ult i32 %idx, 5
+  %zext = zext i1 %cmp to i8
+  %add = add i8 %load, %zext
+  ret i8 %add
+}