[SCEV] Transfer gep nusw and nuw flags

nusw implies nsw offset and nuw base+offset arithmetic if offset
is non-negative. nuw implies nuw offset and base+offset arithmetic.
As usual, we can only transfer is poison implies UB.

Author: nikic

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -3746,21 +3746,23 @@ ScalarEvolution::getGEPExpr(GEPOperator *GEP,
   // getSCEV(Base)->getType() has the same address space as Base->getType()
   // because SCEV::getType() preserves the address space.
   Type *IntIdxTy = getEffectiveSCEVType(BaseExpr->getType());
-  const bool AssumeInBoundsFlags = [&]() {
-    if (!GEP->isInBounds())
-      return false;
-
+  GEPNoWrapFlags NW = GEP->getNoWrapFlags();
+  if (NW != GEPNoWrapFlags::none()) {
     // We'd like to propagate flags from the IR to the corresponding SCEV nodes,
     // but to do that, we have to ensure that said flag is valid in the entire
     // defined scope of the SCEV.
-    auto *GEPI = dyn_cast<Instruction>(GEP);
     // TODO: non-instructions have global scope.  We might be able to prove
     // some global scope cases
-    return GEPI && isSCEVExprNeverPoison(GEPI);
-  }();
+    auto *GEPI = dyn_cast<Instruction>(GEP);
+    if (!GEPI || !isSCEVExprNeverPoison(GEPI))
+      NW = GEPNoWrapFlags::none();
+  }
 
-  SCEV::NoWrapFlags OffsetWrap =
-    AssumeInBoundsFlags ? SCEV::FlagNSW : SCEV::FlagAnyWrap;
+  SCEV::NoWrapFlags OffsetWrap = SCEV::FlagAnyWrap;
+  if (NW.hasNoUnsignedSignedWrap())
+    OffsetWrap = setFlags(OffsetWrap, SCEV::FlagNSW);
+  if (NW.hasNoUnsignedWrap())
+    OffsetWrap = setFlags(OffsetWrap, SCEV::FlagNUW);
 
   Type *CurTy = GEP->getType();
   bool FirstIter = true;
@@ -3806,8 +3808,9 @@ ScalarEvolution::getGEPExpr(GEPOperator *GEP,
   // Add the base address and the offset. We cannot use the nsw flag, as the
   // base address is unsigned. However, if we know that the offset is
   // non-negative, we can use nuw.
-  SCEV::NoWrapFlags BaseWrap = AssumeInBoundsFlags && isKnownNonNegative(Offset)
-                                   ? SCEV::FlagNUW : SCEV::FlagAnyWrap;
+  bool NUW = NW.hasNoUnsignedWrap() ||
+             (NW.hasNoUnsignedSignedWrap() && isKnownNonNegative(Offset));
+  SCEV::NoWrapFlags BaseWrap = NUW ? SCEV::FlagNUW : SCEV::FlagAnyWrap;
   auto *GEPExpr = getAddExpr(BaseExpr, Offset, BaseWrap);
   assert(BaseExpr->getType() == GEPExpr->getType() &&
          "GEP should not change type mid-flight.");

llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll

@@ -1956,3 +1956,86 @@ define noundef i32 @add-recurse-inline() {
   %res = add nuw i32 %x, %y
   ret i32 %res
 }
+
+define noundef ptr @gep_inbounds(ptr %p, i64 %index) {
+; CHECK-LABEL: 'gep_inbounds'
+; CHECK-NEXT:  Classifying expressions for: @gep_inbounds
+; CHECK-NEXT:    %gep = getelementptr inbounds i32, ptr %p, i64 %index
+; CHECK-NEXT:    --> ((4 * %index)<nsw> + %p) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_inbounds
+;
+  %gep = getelementptr inbounds i32, ptr %p, i64 %index
+  ret ptr %gep
+}
+
+define noundef ptr @gep_inbounds_nneg(ptr %p, i32 %index) {
+; CHECK-LABEL: 'gep_inbounds_nneg'
+; CHECK-NEXT:  Classifying expressions for: @gep_inbounds_nneg
+; CHECK-NEXT:    %index.ext = zext i32 %index to i64
+; CHECK-NEXT:    --> (zext i32 %index to i64) U: [0,4294967296) S: [0,4294967296)
+; CHECK-NEXT:    %gep = getelementptr inbounds i32, ptr %p, i64 %index.ext
+; CHECK-NEXT:    --> ((4 * (zext i32 %index to i64))<nuw><nsw> + %p)<nuw> U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_inbounds_nneg
+;
+  %index.ext = zext i32 %index to i64
+  %gep = getelementptr inbounds i32, ptr %p, i64 %index.ext
+  ret ptr %gep
+}
+
+define noundef ptr @gep_nusw(ptr %p, i64 %index) {
+; CHECK-LABEL: 'gep_nusw'
+; CHECK-NEXT:  Classifying expressions for: @gep_nusw
+; CHECK-NEXT:    %gep = getelementptr nusw i32, ptr %p, i64 %index
+; CHECK-NEXT:    --> ((4 * %index)<nsw> + %p) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_nusw
+;
+  %gep = getelementptr nusw i32, ptr %p, i64 %index
+  ret ptr %gep
+}
+
+define noundef ptr @gep_nusw_nneg(ptr %p, i32 %index) {
+; CHECK-LABEL: 'gep_nusw_nneg'
+; CHECK-NEXT:  Classifying expressions for: @gep_nusw_nneg
+; CHECK-NEXT:    %index.ext = zext i32 %index to i64
+; CHECK-NEXT:    --> (zext i32 %index to i64) U: [0,4294967296) S: [0,4294967296)
+; CHECK-NEXT:    %gep = getelementptr nusw i32, ptr %p, i64 %index.ext
+; CHECK-NEXT:    --> ((4 * (zext i32 %index to i64))<nuw><nsw> + %p)<nuw> U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_nusw_nneg
+;
+  %index.ext = zext i32 %index to i64
+  %gep = getelementptr nusw i32, ptr %p, i64 %index.ext
+  ret ptr %gep
+}
+
+define noundef ptr @gep_nuw(ptr %p, i64 %index) {
+; CHECK-LABEL: 'gep_nuw'
+; CHECK-NEXT:  Classifying expressions for: @gep_nuw
+; CHECK-NEXT:    %gep = getelementptr nuw i32, ptr %p, i64 %index
+; CHECK-NEXT:    --> ((4 * %index)<nuw> + %p)<nuw> U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_nuw
+;
+  %gep = getelementptr nuw i32, ptr %p, i64 %index
+  ret ptr %gep
+}
+
+define noundef ptr @gep_nusw_nuw(ptr %p, i64 %index) {
+; CHECK-LABEL: 'gep_nusw_nuw'
+; CHECK-NEXT:  Classifying expressions for: @gep_nusw_nuw
+; CHECK-NEXT:    %gep = getelementptr nusw nuw i32, ptr %p, i64 %index
+; CHECK-NEXT:    --> ((4 * %index)<nuw><nsw> + %p)<nuw> U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_nusw_nuw
+;
+  %gep = getelementptr nusw nuw i32, ptr %p, i64 %index
+  ret ptr %gep
+}
+
+define ptr @gep_nusw_nuw_missing_noundef(ptr %p, i64 %index) {
+; CHECK-LABEL: 'gep_nusw_nuw_missing_noundef'
+; CHECK-NEXT:  Classifying expressions for: @gep_nusw_nuw_missing_noundef
+; CHECK-NEXT:    %gep = getelementptr nusw nuw i32, ptr %p, i64 %index
+; CHECK-NEXT:    --> ((4 * %index) + %p) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @gep_nusw_nuw_missing_noundef
+;
+  %gep = getelementptr nusw nuw i32, ptr %p, i64 %index
+  ret ptr %gep
+}