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

In most cases, GEPs should be canonicalized by InstCombine. Bail out on
non-canonical forms for simplicity.
Fixes
#155253 (comment).

Author: dtcxzyw

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)

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
+}