[llvm-branch-commits] [llvm] [InstCombine] Improve inbounds preservation for ADD+GEP -> GEP+GEP (PR #135155)

[Nikita Popov via llvm-branch-commits]

================
@@ -3086,13 +3086,51 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
   if (GEPType->isVectorTy())
     return nullptr;
 
+  if (!GEP.isInBounds()) {
+    unsigned IdxWidth =
+        DL.getIndexSizeInBits(PtrOp->getType()->getPointerAddressSpace());
+    APInt BasePtrOffset(IdxWidth, 0);
+    Value *UnderlyingPtrOp =
+        PtrOp->stripAndAccumulateInBoundsConstantOffsets(DL, BasePtrOffset);
+    bool CanBeNull, CanBeFreed;
+    uint64_t DerefBytes = UnderlyingPtrOp->getPointerDereferenceableBytes(
+        DL, CanBeNull, CanBeFreed);
+    if (!CanBeNull && !CanBeFreed && DerefBytes != 0) {
+      if (GEP.accumulateConstantOffset(DL, BasePtrOffset) &&
+          BasePtrOffset.isNonNegative()) {
+        APInt AllocSize(IdxWidth, DerefBytes);
+        if (BasePtrOffset.ule(AllocSize)) {
+          return GetElementPtrInst::CreateInBounds(
+              GEP.getSourceElementType(), PtrOp, Indices, GEP.getName());
+        }
+      }
+    }
+  }
+
+  // nusw + nneg -> nuw
+  if (GEP.hasNoUnsignedSignedWrap() && !GEP.hasNoUnsignedWrap() &&
+      all_of(GEP.indices(), [&](Value *Idx) {
+        return isKnownNonNegative(Idx, SQ.getWithInstruction(&GEP));
+      })) {
+    GEP.setNoWrapFlags(GEP.getNoWrapFlags() | GEPNoWrapFlags::noUnsignedWrap());
+    return &GEP;
+  }
+
+  // These rewrites is trying to preserve inbounds/nuw attributes. So we want to
+  // do this after having tried to derive "nuw" above.
   if (GEP.getNumIndices() == 1) {
-    // We can only preserve inbounds if the original gep is inbounds, the add
-    // is nsw, and the add operands are non-negative.
-    auto CanPreserveInBounds = [&](bool AddIsNSW, Value *Idx1, Value *Idx2) {
-      SimplifyQuery Q = SQ.getWithInstruction(&GEP);
-      return GEP.isInBounds() && AddIsNSW && isKnownNonNegative(Idx1, Q) &&
-             isKnownNonNegative(Idx2, Q);
+    auto GetPreservedNoWrapFlags = [&](bool AddIsNUW, Value *Idx1,
+                                       Value *Idx2) {
+      // Preserve "inbounds nuw" if the original gep is "inbounds nuw", and the
+      // add is "nuw". Preserve "nuw" if the original gep is "nuw", and the add
+      // is "nuw".
+      GEPNoWrapFlags Flags = GEPNoWrapFlags::none();
+      if (GEP.hasNoUnsignedWrap() && AddIsNUW) {
+        Flags |= GEPNoWrapFlags::noUnsignedWrap();
+        if (GEP.isInBounds())
+          Flags |= GEPNoWrapFlags::inBounds();
----------------
nikic wrote:

Same comment here, we should just be returning the original Flags instead of reconstructing them. Note that this works for all of `nuw`, `inbounds nuw` and `nusw nuw`: https://alive2.llvm.org/ce/z/QSweWW

https://github.com/llvm/llvm-project/pull/135155