[DSE,MSSA] Remove dead stores to stack-allocated locations

[Whitney]

Summary

This patch adds the support to remove dead stores to stack-allocated locations when -enable-dse-memoryssa=true.
For example:

%A = alloca i32
...
store i32 1, i32* %A
ret void

Note that dead stores to stack-allocated locations in end blocks can already be removed when -enable-dse-memoryssa=false.

[fhahn]

Thanks for working on this. I think a patch I've been working on for a while solves the same problem, but in a slightly more general way I think: D72631.

There's one more slightly larger outstanding patch to support 'memory terminators' like lifetime.end/free as well: D72410. I think those patches should get us to almost cover all cases legacy DSE does, but there is plenty of potential for further improvements. I would be more than happy to sync up to avoid duplicating work, if you are also interested in working on MemorySSA backed DSE :)

[Whitney]

In D82407#2110262, @fhahn wrote:

Thanks for working on this. I think a patch I've been working on for a while solves the same problem, but in a slightly more general way I think: D72631.

There's one more slightly larger outstanding patch to support 'memory terminators' like lifetime.end/free as well: D72410. I think those patches should get us to almost cover all cases legacy DSE does, but there is plenty of potential for further improvements. I would be more than happy to sync up to avoid duplicating work, if you are also interested in working on MemorySSA backed DSE :)

@fhahn Cool, cancelled this revision. Positive side is I understand better of how MemorySSA backed DSE works after working on this patch :)
Do you know if there are plans on eliminating dead stores to array (in a loop)?

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

static void removeAccessedObjects(const MemoryLocation &LoadedLoc,
   // Remove objects that could alias LoadedLoc.
   DeadStackObjects.remove_if([&](const Value *I) {
     // See if the loaded location could alias the stack location.
     MemoryLocation StackLoc(I, getPointerSize(I, DL, *TLI, F));
     return !AA->isNoAlias(StackLoc, LoadedLoc);
   });
 }
 
+// Collect stack objects that are dead at the end of the function.
+static void
+collectNonEscapedObjects(const Function &F,
+                         SmallSetVector<const Value *, 16> &NonEscapedObjects,
+                         const TargetLibraryInfo *TLI) {
+  // Find all of the alloca'd pointers in the entry block.
+  const BasicBlock &Entry = F.front();
+  for (const Instruction &I : Entry) {
+    if (isa<AllocaInst>(&I))
+      NonEscapedObjects.insert(&I);
+
+    // Okay, so these are dead heap objects, but if the pointer never escapes
+    // then it's leaked by this function anyways.
+    else if (isAllocLikeFn(&I, TLI) && !PointerMayBeCaptured(&I, true, true))
+      NonEscapedObjects.insert(&I);
+  }
+
+  // Treat byval or inalloca arguments the same, stores to them are dead at the
+  // end of the function.
+  for (const Argument &AI : F.args())
+    if (AI.hasPassPointeeByValueAttr())
+      NonEscapedObjects.insert(&AI);
+}
+
 /// Remove dead stores to stack-allocated locations in the function end block.
 /// Ex:
 /// %A = alloca i32
 /// ...
 /// store i32 1, i32* %A
 /// ret void
 static bool handleEndBlock(BasicBlock &BB, AliasAnalysis *AA,
                            MemoryDependenceResults *MD,
                            const TargetLibraryInfo *TLI,
                            InstOverlapIntervalsTy &IOL,
                            MapVector<Instruction *, bool> &ThrowableInst) {
   bool MadeChange = false;
 
   // Keep track of all of the stack objects that are dead at the end of the
   // function.
   SmallSetVector<const Value*, 16> DeadStackObjects;
-
-  // Find all of the alloca'd pointers in the entry block.
-  BasicBlock &Entry = BB.getParent()->front();
-  for (Instruction &I : Entry) {
-    if (isa<AllocaInst>(&I))
-      DeadStackObjects.insert(&I);
-
-    // Okay, so these are dead heap objects, but if the pointer never escapes
-    // then it's leaked by this function anyways.
-    else if (isAllocLikeFn(&I, TLI) && !PointerMayBeCaptured(&I, true, true))
-      DeadStackObjects.insert(&I);
-  }
-
-  // Treat byval or inalloca arguments the same, stores to them are dead at the
-  // end of the function.
-  for (Argument &AI : BB.getParent()->args())
-    if (AI.hasPassPointeeByValueAttr())
-      DeadStackObjects.insert(&AI);
+  collectNonEscapedObjects(*BB.getParent(), DeadStackObjects, TLI);
 
   const DataLayout &DL = BB.getModule()->getDataLayout();
 
   // Scan the basic block backwards
   for (BasicBlock::iterator BBI = BB.end(); BBI != BB.begin(); ){
     --BBI;
 
     // If we find a store, check to see if it points into a dead stack value.
for (unsigned I = 0; I < ToCheck.size(); I++) {
       }
     }
   }
 
   if (EnablePartialOverwriteTracking)
     for (auto &KV : State.IOLs)
       MadeChange |= removePartiallyOverlappedStores(&AA, DL, KV.second);
 
+  // Keep track of all of the stack objects that are dead at the end of the
+  // function.
+  SmallSetVector<const Value *, 16> DeadStackObjects;
+  collectNonEscapedObjects(F, DeadStackObjects, &TLI);
+  for (MemoryDef *Current : State.MemDefs) {
+    if (State.SkipStores.count(Current))
+      continue;
+    Instruction *SI = Current->getMemoryInst();
+    if (!Current->use_empty())
+      continue;
+
+    LLVM_DEBUG(dbgs() << "No MemoryUse uses " << *Current << " (" << *SI
+                      << ")\n");
+
+    // See through pointer-to-pointer bitcasts
+    SmallVector<const Value *, 4> Pointers;
+    GetUnderlyingObjects(getStoredPointerOperand(SI), Pointers, DL);
+
+    // Stores to stack values are valid candidates for removal.
+    bool AllDead = all_of(Pointers, [&DeadStackObjects](const Value *Pointer) {
+      return DeadStackObjects.count(Pointer);
+    });
+    if (AllDead) {
+      State.deleteDeadInstruction(SI);
+      MadeChange = true;
+      continue;
+    }
+  }
+
   return MadeChange;
 }
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // DSE Pass
 //===----------------------------------------------------------------------===//
 PreservedAnalyses DSEPass::run(Function &F, FunctionAnalysisManager &AM) {

llvm/test/Transforms/DeadStoreElimination/MSSA/mda-with-dbg-values.ll

 !15 = !DISubroutineType(types: !16)
 !16 = !{null}
 !17 = !DILocalVariable(name: "i", scope: !14, file: !3, line: 7, type: !7)
 !18 = !DILocation(line: 7, column: 10, scope: !14)
 !19 = !DILocation(line: 8, column: 7, scope: !14)
 !20 = !DILocation(line: 9, column: 5, scope: !14)
 !21 = !DILocation(line: 10, column: 1, scope: !14)
 
-; Check that the store is removed and that the memcpy is still there
+; Check that the store and memcpy are removed
 ; CHECK-LABEL: foo
 ; CHECK-NOT:   store i8
-; CHECK:       call void @llvm.memcpy
+; CHECK-NOT:       call void @llvm.memcpy
 ; CHECK:       ret void

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-captures.ll

;
   call void @may_throw()
   br label %exit
 
 exit:
   store i8 1, i8* %m
   call void @capture(i8* %m)
   ret i8* %m
 }
-; TODO: Remove store in exit.
+
 ; Stores to stack objects can be eliminated if they are not captured inside the function.
 define void @test_alloca_nocapture_1() {
 ; CHECK-LABEL: @test_alloca_nocapture_1(
-; CHECK-NEXT:    [[M:%.*]] = alloca i8
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    store i8 1, i8* [[M]]
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8
   store i8 0, i8* %m
   call void @foo()
   br label %exit
 
 exit:
   store i8 1, i8* %m
   ret void
 }
 
-; TODO: Remove store in exit.
 ; Cannot remove first store i8 0, i8* %m, as the call to @capture captures the object.
 define void @test_alloca_capture_1() {
 ; CHECK-LABEL: @test_alloca_capture_1(
 ; CHECK-NEXT:    [[M:%.*]] = alloca i8
 ; CHECK-NEXT:    store i8 0, i8* [[M]]
 ; CHECK-NEXT:    call void @capture(i8* [[M]])
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    store i8 1, i8* [[M]]
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8
   store i8 0, i8* %m
   call void @capture(i8* %m)
   br label %exit
 
 exit:
   store i8 1, i8* %m
   ret void
 }
 
-; TODO: Remove store at exit.
 ; We can remove the last store to %m, even though it escapes because the alloca
 ; becomes invalid after the function returns.
 define void @test_alloca_capture_2(%S1* %E) {
 ; CHECK-LABEL: @test_alloca_capture_2(
 ; CHECK-NEXT:    [[M:%.*]] = alloca i8
 ; CHECK-NEXT:    br label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[F_PTR:%.*]] = getelementptr [[S1:%.*]], %S1* [[E:%.*]], i32 0, i32 0
 ; CHECK-NEXT:    store i8* [[M]], i8** [[F_PTR]]
-; CHECK-NEXT:    store i8 1, i8* [[M]]
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8
   br label %exit
 
 exit:
   %f.ptr = getelementptr %S1, %S1* %E, i32 0, i32 0
   store i8* %m, i8** %f.ptr

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-exceptions.ll

 ; CHECK:       catch:
 ; CHECK-NEXT:    [[C:%.*]] = catchpad within [[CS1]] []
 ; CHECK-NEXT:    [[LV:%.*]] = load i32, i32* [[SV]]
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       cleanup:
 ; CHECK-NEXT:    [[C1:%.*]] = cleanuppad within none []
 ; CHECK-NEXT:    br label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    store i32 40, i32* [[SV]]
 ; CHECK-NEXT:    ret void
 ;
 block1:
   %sv = alloca i32
   br label %block2
 
 block2:
   store i32 20, i32* %sv

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-memintrinsics.ll

 bb3:
   ret void
 }
 
 ; Stores along  both exit paths.
 define void @alloca_3(i1 %c) {
 ; CHECK-LABEL: @alloca_3(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[P_ALLOCA:%.*]] = alloca [32 x i32]
-; CHECK-NEXT:    [[P:%.*]] = bitcast [32 x i32]* [[P_ALLOCA]] to i32*
-; CHECK-NEXT:    [[ARRAYIDX0:%.*]] = getelementptr inbounds i32, i32* [[P]], i64 1
-; CHECK-NEXT:    [[P3:%.*]] = bitcast i32* [[ARRAYIDX0]] to i8*
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, i8* [[P3]], i64 4
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[TMP0]], i8 0, i64 24, i1 false)
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[BB1:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    br label [[BB3:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
-; CHECK-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i32, i32* [[P]], i64 1
-; CHECK-NEXT:    store i32 1, i32* [[ARRAYIDX1]], align 4
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %P.alloca = alloca [32 x i32]
   %P = bitcast [32 x i32]* %P.alloca to i32*
   %arrayidx0 = getelementptr inbounds i32, i32* %P, i64 1
   %p3 = bitcast i32* %arrayidx0 to i8*
   call void @llvm.memset.p0i8.i64(i8* %p3, i8 0, i64 28, i32 4, i1 false)

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-multipath.ll

 %struct.blam.4 = type { %struct.bar.5, [4 x i8] }
 %struct.bar.5 = type <{ i64, i64*, i32, i64 }>
 
 ; Make sure we do not eliminate the store in %bb.
 define void @alloca_5(i1 %c) {
 ; CHECK-LABEL: @alloca_5(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[TMP:%.*]] = alloca [[STRUCT_BLAM_4:%.*]], align 8
-; CHECK-NEXT:    [[TMP36:%.*]] = getelementptr inbounds [[STRUCT_BLAM_4]], %struct.blam.4* [[TMP]], i64 0, i32 0, i32 1
-; CHECK-NEXT:    [[TMP37:%.*]] = bitcast i64** [[TMP36]] to i8*
 ; CHECK-NEXT:    [[TMP38:%.*]] = getelementptr inbounds [[STRUCT_BLAM_4]], %struct.blam.4* [[TMP]], i64 0, i32 0, i32 3
 ; CHECK-NEXT:    [[TMP39:%.*]] = bitcast i64* [[TMP38]] to i64*
 ; CHECK-NEXT:    store i64 0, i64* [[TMP39]], align 4
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[BB46:%.*]], label [[BB47:%.*]]
 ; CHECK:       bb46:
-; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* nonnull align 8 dereferenceable(20) [[TMP37]], i8 0, i64 26, i1 false)
 ; CHECK-NEXT:    ret void
 ; CHECK:       bb47:
 ; CHECK-NEXT:    [[TMP48:%.*]] = getelementptr inbounds [[STRUCT_BLAM_4]], %struct.blam.4* [[TMP]], i64 0, i32 0, i32 2
 ; CHECK-NEXT:    store i32 20, i32* [[TMP48]], align 8
 ; CHECK-NEXT:    br label [[BB52:%.*]]
 ; CHECK:       bb52:
 ; CHECK-NEXT:    br i1 [[C]], label [[BB68:%.*]], label [[BB59:%.*]]
 ; CHECK:       bb59:

llvm/test/Transforms/DeadStoreElimination/MSSA/multiblock-simple.ll

bb2:
   ret void
 bb3:
   ret void
 }
 
 
 define void @test11() {
 ; CHECK-LABEL: @test11(
-; CHECK-NEXT:    [[P:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    br i1 true, label [[BB1:%.*]], label [[BB2:%.*]]
 ; CHECK:       bb1:
-; CHECK-NEXT:    store i32 0, i32* [[P]], align 4
 ; CHECK-NEXT:    br label [[BB3:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    ret void
 ; CHECK:       bb3:
 ; CHECK-NEXT:    ret void
 ;
   %P = alloca i32
   store i32 0, i32* %P

llvm/test/Transforms/DeadStoreElimination/MSSA/simple.ll

 ; CHECK-NEXT:    store volatile i32 2, i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    store volatile i32 3, i32* [[P]], align 4
 ; CHECK-NEXT:    ret void
 ;
   store volatile i32 2, i32* %P, align 4
   store volatile i32 3, i32* %P, align 4
   ret void
 }
+
+define void @test48_volatile() {
+; CHECK-LABEL: @test48_volatile(
+; CHECK-NEXT:    [[P:%.*]] = alloca i32
+; CHECK-NEXT:    store volatile i32 2, i32* [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+  %P = alloca i32
+  store volatile i32 2, i32* %P, align 4
+  ret void
+}