[llvm] [DSE] Update dereferenceable attributes when adjusting memintrinsic ptr (PR #125073)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-ir

Author: Björn Pettersson (bjope)

<details>
<summary>Changes</summary>

Consider IR like this
  call void @<!-- -->llvm.memset.p0.i64(ptr dereferenceable(28) %p, i8 0, i64 28, i1 false)
  store i32 1, ptr %p

It has been optimized like this:
  %p2 = getelementptr inbounds i8, ptr %p, i64 4
  call void @<!-- -->llvm.memset.p0.i64(ptr dereferenceable(28) %p2, i8 0, i64 24, i1 false)
  store i32 1, ptr %p

As the input IR doesn't guarantee that it is OK to deref 28 bytes starting at the adjusted pointer %p2 the transformation has been a bit flawed.

With this patch we make sure to also adjust the size of any dereferenceable/dereferenceable_or_null attributes when doing the transform in tryToShorten (when adjusting the start pointer). So now we will get dereferenceable(24) in the example above.

---
Full diff: https://github.com/llvm/llvm-project/pull/125073.diff


3 Files Affected:

- (modified) llvm/include/llvm/IR/InstrTypes.h (+5) 
- (modified) llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp (+18) 
- (modified) llvm/test/Transforms/DeadStoreElimination/OverwriteStoreBegin.ll (+30) 


``````````diff
diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
index 47ddc7555594c5..f26aafbe3f0ef8 100644
--- a/llvm/include/llvm/IR/InstrTypes.h
+++ b/llvm/include/llvm/IR/InstrTypes.h
@@ -1559,6 +1559,11 @@ class CallBase : public Instruction {
     Attrs = Attrs.addDereferenceableParamAttr(getContext(), i, Bytes);
   }
 
+  /// adds the dereferenceable attribute to the list of attributes.
+  void addDereferenceableOrNullParamAttr(unsigned i, uint64_t Bytes) {
+    Attrs = Attrs.addDereferenceableOrNullParamAttr(getContext(), i, Bytes);
+  }
+
   /// adds the dereferenceable attribute to the list of attributes.
   void addDereferenceableRetAttr(uint64_t Bytes) {
     Attrs = Attrs.addDereferenceableRetAttr(getContext(), Bytes);
diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 13f3de07c3c44d..c16134e3f01598 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -563,6 +563,23 @@ static void shortenAssignment(Instruction *Inst, Value *OriginalDest,
   for_each(LinkedDVRAssigns, InsertAssignForOverlap);
 }
 
+// Helper to trim or drop any dereferencable/dereferencable_or_null attributes
+// for a given argument, based on the new access being restricted to derefence
+// bytes in the range [Offset, Offset+Size).
+static void trimDereferencableAttrs(AnyMemIntrinsic *Intrinsic, unsigned Arg,
+                                    uint64_t Offset, uint64_t Size) {
+  uint64_t End = Offset + Size;
+  if (Intrinsic->getParamDereferenceableBytes(Arg) >= End)
+    Intrinsic->addDereferenceableParamAttr(Arg, Size);
+  else
+    Intrinsic->removeParamAttr(Arg, Attribute::Dereferenceable);
+  if (Intrinsic->getParamDereferenceableOrNullBytes(Arg) >= End)
+    Intrinsic->addDereferenceableOrNullParamAttr(Arg, Size);
+  else
+    Intrinsic->removeParamAttr(Arg, Attribute::DereferenceableOrNull);
+}
+
+
 static bool tryToShorten(Instruction *DeadI, int64_t &DeadStart,
                          uint64_t &DeadSize, int64_t KillingStart,
                          uint64_t KillingSize, bool IsOverwriteEnd) {
@@ -644,6 +661,7 @@ static bool tryToShorten(Instruction *DeadI, int64_t &DeadStart,
         DeadI->getIterator());
     NewDestGEP->setDebugLoc(DeadIntrinsic->getDebugLoc());
     DeadIntrinsic->setDest(NewDestGEP);
+    trimDereferencableAttrs(DeadIntrinsic, 0, ToRemoveSize, NewSize);
   }
 
   // Update attached dbg.assign intrinsics. Assume 8-bit byte.
diff --git a/llvm/test/Transforms/DeadStoreElimination/OverwriteStoreBegin.ll b/llvm/test/Transforms/DeadStoreElimination/OverwriteStoreBegin.ll
index bc1756f6ca9d1b..6c42d8e2c98338 100644
--- a/llvm/test/Transforms/DeadStoreElimination/OverwriteStoreBegin.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/OverwriteStoreBegin.ll
@@ -402,3 +402,33 @@ entry:
   store i64 1, ptr %p, align 1
   ret void
 }
+
+; Verify that we adjust the dereferenceable attribute.
+define void @dereferenceable(ptr nocapture %p) {
+; CHECK-LABEL: @dereferenceable(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 4
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 4 dereferenceable(24) [[TMP0]], i8 0, i64 24, i1 false)
+; CHECK-NEXT:    store i32 1, ptr [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @llvm.memset.p0.i64(ptr dereferenceable(28) align 4 %p, i8 0, i64 28, i1 false)
+  store i32 1, ptr %p, align 4
+  ret void
+}
+
+; Verify that we adjust the dereferenceable_or_null attribute.
+define void @dereferenceable_or_null(ptr nocapture %p) {
+; CHECK-LABEL: @dereferenceable_or_null(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 8
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 4 dereferenceable_or_null(20) [[TMP0]], i8 0, i64 20, i1 false)
+; CHECK-NEXT:    store i64 1, ptr [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @llvm.memset.p0.i64(ptr dereferenceable_or_null(28) align 4 %p, i8 0, i64 28, i1 false)
+  store i64 1, ptr %p, align 4
+  ret void
+}

``````````

</details>


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