[llvm-branch-commits] [DSE] Make DSE eliminate stores to objects with a sized dead_on_return (PR #173694)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Aiden Grossman (boomanaiden154)

<details>
<summary>Changes</summary>

dead_on_return is made optionally sized in #<!-- -->171712. This patch adds
handling in DSE so that we can actually eliminate stores to pointer
parameters marked with a sized dead_on_return attribute. We do not
eliminate stores where the store may overlap with bytes that are not
known to be dead after return.


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


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp (+22-5) 
- (modified) llvm/test/Transforms/DeadStoreElimination/simple.ll (+43) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 87393cd595e97..25d473b5beb72 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -965,6 +965,7 @@ struct DSEState {
   // Keep track of all of the objects that are invisible to the caller after
   // the function returns.
   DenseMap<const Value *, bool> InvisibleToCallerAfterRet;
+  DenseMap<const Value *, uint64_t> InvisibleToCallerAfterRetBounded;
   // Keep track of blocks with throwing instructions not modeled in MemorySSA.
   SmallPtrSet<BasicBlock *, 16> ThrowingBlocks;
   // Post-order numbers for each basic block. Used to figure out if memory
@@ -1016,11 +1017,18 @@ struct DSEState {
 
     // Treat byval, inalloca or dead on return arguments the same as Allocas,
     // stores to them are dead at the end of the function.
-    for (Argument &AI : F.args())
+    for (Argument &AI : F.args()) {
       if (AI.hasPassPointeeByValueCopyAttr() ||
           (AI.getType()->isPointerTy() &&
            AI.getDeadOnReturnInfo().coversAllReachableMemory()))
         InvisibleToCallerAfterRet.insert({&AI, true});
+      if (AI.getType()->isPointerTy() &&
+          !AI.getDeadOnReturnInfo().coversAllReachableMemory()) {
+        if (uint64_t DeadOnReturnBytes =
+                AI.getDeadOnReturnInfo().getNumberOfDeadBytes())
+          InvisibleToCallerAfterRetBounded.insert({&AI, DeadOnReturnBytes});
+      }
+    }
 
     // Collect whether there is any irreducible control flow in the function.
     ContainsIrreducibleLoops = mayContainIrreducibleControl(F, &LI);
@@ -1206,11 +1214,20 @@ struct DSEState {
     return OW_None;
   }
 
-  bool isInvisibleToCallerAfterRet(const Value *V) {
+  bool isInvisibleToCallerAfterRet(const Value *V, const Value *Ptr, const LocationSize StoreSize) {
     if (isa<AllocaInst>(V))
       return true;
 
     auto I = InvisibleToCallerAfterRet.insert({V, false});
+    if (I.second && InvisibleToCallerAfterRetBounded.contains(V)) {
+      int64_t ValueOffset;
+      const Value *BaseValue =
+          GetPointerBaseWithConstantOffset(Ptr, ValueOffset, DL);
+      assert(BaseValue == V);
+      if (ValueOffset + StoreSize.toRaw() <
+          InvisibleToCallerAfterRetBounded[BaseValue])
+        return true;
+    }
     if (I.second && isInvisibleToCallerOnUnwind(V) && isNoAliasCall(V))
       I.first->second = capturesNothing(PointerMayBeCaptured(
           V, /*ReturnCaptures=*/true, CaptureComponents::Provenance));
@@ -1757,7 +1774,7 @@ struct DSEState {
           BasicBlock *MaybeKillingBlock = UseInst->getParent();
           if (PostOrderNumbers.find(MaybeKillingBlock)->second <
               PostOrderNumbers.find(MaybeDeadAccess->getBlock())->second) {
-            if (!isInvisibleToCallerAfterRet(KillingUndObj)) {
+            if (!isInvisibleToCallerAfterRet(KillingUndObj, KillingLoc.Ptr, KillingLoc.Size)) {
               LLVM_DEBUG(dbgs()
                          << "    ... found killing def " << *UseInst << "\n");
               KillingDefs.insert(UseInst);
@@ -1775,7 +1792,7 @@ struct DSEState {
     // For accesses to locations visible after the function returns, make sure
     // that the location is dead (=overwritten) along all paths from
     // MaybeDeadAccess to the exit.
-    if (!isInvisibleToCallerAfterRet(KillingUndObj)) {
+    if (!isInvisibleToCallerAfterRet(KillingUndObj, KillingLoc.Ptr, KillingLoc.Size)) {
       SmallPtrSet<BasicBlock *, 16> KillingBlocks;
       for (Instruction *KD : KillingDefs)
         KillingBlocks.insert(KD->getParent());
@@ -1982,7 +1999,7 @@ struct DSEState {
         // underlying objects is very uncommon. If it turns out to be important,
         // we can use getUnderlyingObjects here instead.
         const Value *UO = getUnderlyingObject(DefLoc->Ptr);
-        if (!isInvisibleToCallerAfterRet(UO))
+        if (!isInvisibleToCallerAfterRet(UO, DefLoc->Ptr, DefLoc->Size))
           continue;
 
         if (isWriteAtEndOfFunction(Def, *DefLoc)) {
diff --git a/llvm/test/Transforms/DeadStoreElimination/simple.ll b/llvm/test/Transforms/DeadStoreElimination/simple.ll
index 9d28395a4ccd0..7619842ea18cf 100644
--- a/llvm/test/Transforms/DeadStoreElimination/simple.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/simple.ll
@@ -886,5 +886,48 @@ define ptr @test_dead_on_return_ptr_returned(ptr dead_on_return %p) {
   ret ptr %p
 }
 
+define void @test_dead_on_return_oob(ptr dead_on_return(4) %p) {
+; CHECK-LABEL: @test_dead_on_return_oob(
+; CHECK-NEXT:    [[LOCAL_VAR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    call void @opaque(ptr [[LOCAL_VAR]])
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr i8, ptr [[P:%.*]], i64 8
+; CHECK-NEXT:    store ptr [[LOCAL_VAR]], ptr [[P1]], align 8
+; CHECK-NEXT:    ret void
+;
+  %local.var = alloca ptr
+  call void @opaque(ptr %local.var)
+  %p1 = getelementptr i8, ptr %p, i64 8
+  store ptr %local.var, ptr %p1
+  ret void
+}
+
+define void @test_dead_on_return_inbounds(ptr dead_on_return(16) %p) {
+; CHECK-LABEL: @test_dead_on_return_inbounds(
+; CHECK-NEXT:    [[LOCAL_VAR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    call void @opaque(ptr [[LOCAL_VAR]])
+; CHECK-NEXT:    ret void
+;
+  %local.var = alloca ptr
+  call void @opaque(ptr %local.var)
+  %p1 = getelementptr inbounds i8, ptr %p, i64 2
+  store ptr %local.var, ptr %p1
+  ret void
+}
+
+define void @test_on_return_overlapping_oob(ptr dead_on_return(8) %p) {
+; CHECK-LABEL: @test_on_return_overlapping_oob(
+; CHECK-NEXT:    [[LOCAL_VAR:%.*]] = alloca ptr, align 8
+; CHECK-NEXT:    call void @opaque(ptr [[LOCAL_VAR]])
+; CHECK-NEXT:    [[P1:%.*]] = getelementptr inbounds i8, ptr [[P:%.*]], i64 4
+; CHECK-NEXT:    store ptr [[LOCAL_VAR]], ptr [[P1]], align 8
+; CHECK-NEXT:    ret void
+;
+  %local.var = alloca ptr
+  call void @opaque(ptr %local.var)
+  %p1 = getelementptr inbounds i8, ptr %p, i64 4
+  store ptr %local.var, ptr %p1
+  ret void
+}
+
 declare void @opaque(ptr)
 declare void @maythrow() memory(none)

``````````

</details>


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