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

[Björn Pettersson 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();
----------------
bjope wrote:

> If you don't match the index size, you are introducing an implicit sext() on the indices. The transform is not valid in that case.
Oh, right, the transform isn't even valid if dropping the flags in that case.

It puzzeled me for awhile how implicit sext/trunc of indices would impact these rewrites (also for the legacy code with NSW considering that there are no checks on index size). Now I found out that we make sure that the gep index operand is using the correct index type size earlier in this function.  So we can expect that we match with the index size when doing the rewrites into GEP+GEP here.

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