diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index e318ec94db4c3..66168a92f460e 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1213,7 +1213,8 @@ struct DSEState {
 
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second && isInvisibleToCallerOnUnwind(V) && isNoAliasCall(V))
-      I.first->second = !PointerMayBeCaptured(V, /*ReturnCaptures=*/true);
+      I.first->second = capturesNothing(PointerMayBeCaptured(
+          V, /*ReturnCaptures=*/true, CaptureComponents::Provenance));
     return I.first->second;
   }
 
@@ -1230,7 +1231,8 @@ struct DSEState {
       // with the killing MemoryDef. But we refrain from doing so for now to
       // limit compile-time and this does not cause any changes to the number
       // of stores removed on a large test set in practice.
-      I.first->second = PointerMayBeCaptured(V, /*ReturnCaptures=*/false);
+      I.first->second = capturesAnything(PointerMayBeCaptured(
+          V, /*ReturnCaptures=*/false, CaptureComponents::Provenance));
     return !I.first->second;
   }
 
diff --git a/llvm/test/Transforms/DeadStoreElimination/assume.ll b/llvm/test/Transforms/DeadStoreElimination/assume.ll
index ddf61542b903b..9df6d33afe10f 100644
--- a/llvm/test/Transforms/DeadStoreElimination/assume.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/assume.ll
@@ -8,7 +8,6 @@ define void @f() {
 ; CHECK-NEXT:    [[TMP1:%.*]] = call noalias ptr @_Znwm(i64 32)
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt ptr [[TMP1]], @global
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP2]])
-; CHECK-NEXT:    store i8 0, ptr [[TMP1]], align 1
 ; CHECK-NEXT:    ret void
 ;
   %tmp1 = call noalias ptr @_Znwm(i64 32)
@@ -23,7 +22,6 @@ define void @f2() {
 ; CHECK-NEXT:    [[TMP1:%.*]] = call noalias ptr @_Znwm(i64 32)
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt ptr [[TMP1]], @global
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP2]])
-; CHECK-NEXT:    store i8 0, ptr [[TMP1]], align 1
 ; CHECK-NEXT:    call void @quux(ptr @global)
 ; CHECK-NEXT:    ret void
 ;
diff --git a/llvm/test/Transforms/DeadStoreElimination/simple.ll b/llvm/test/Transforms/DeadStoreElimination/simple.ll
index af5b77c79acdc..f8e594b5626a0 100644
--- a/llvm/test/Transforms/DeadStoreElimination/simple.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/simple.ll
@@ -304,6 +304,58 @@ define void @custom_malloc_no_escape() {
   ret void
 }
 
+declare void @use.ptr(ptr)
+
+define void @malloc_no_escape_via_attr() {
+; CHECK-LABEL: @malloc_no_escape_via_attr(
+; CHECK-NEXT:    [[M:%.*]] = call ptr @malloc(i64 24)
+; CHECK-NEXT:    call void @use.ptr(ptr captures(none) [[M]])
+; CHECK-NEXT:    ret void
+;
+  %m = call ptr @malloc(i64 24)
+  call void @use.ptr(ptr captures(none) %m)
+  store i8 0, ptr %m
+  ret void
+}
+
+define void @malloc_address_only_escape() {
+; CHECK-LABEL: @malloc_address_only_escape(
+; CHECK-NEXT:    [[M:%.*]] = call ptr @malloc(i64 24)
+; CHECK-NEXT:    call void @use.ptr(ptr captures(address) [[M]])
+; CHECK-NEXT:    ret void
+;
+  %m = call ptr @malloc(i64 24)
+  call void @use.ptr(ptr captures(address) %m)
+  store i8 0, ptr %m
+  ret void
+}
+
+define void @malloc_provenance_escape() {
+; CHECK-LABEL: @malloc_provenance_escape(
+; CHECK-NEXT:    [[M:%.*]] = call ptr @malloc(i64 24)
+; CHECK-NEXT:    call void @use.ptr(ptr captures(provenance) [[M]])
+; CHECK-NEXT:    store i8 0, ptr [[M]], align 1
+; CHECK-NEXT:    ret void
+;
+  %m = call ptr @malloc(i64 24)
+  call void @use.ptr(ptr captures(provenance) %m)
+  store i8 0, ptr %m
+  ret void
+}
+
+define void @malloc_read_provenance_escape() {
+; CHECK-LABEL: @malloc_read_provenance_escape(
+; CHECK-NEXT:    [[M:%.*]] = call ptr @malloc(i64 24)
+; CHECK-NEXT:    call void @use.ptr(ptr captures(read_provenance) [[M]])
+; CHECK-NEXT:    store i8 0, ptr [[M]], align 1
+; CHECK-NEXT:    ret void
+;
+  %m = call ptr @malloc(i64 24)
+  call void @use.ptr(ptr captures(read_provenance) %m)
+  store i8 0, ptr %m
+  ret void
+}
+
 define void @test21() {
 ; CHECK-LABEL: @test21(
 ; CHECK-NEXT:    ret void
@@ -484,7 +536,7 @@ define i32 @test32(i1 %c, ptr %p, i32 %i, i1 %arg) {
 ; CHECK:       bb1:
 ; CHECK-NEXT:    store i32 [[V]], ptr [[P]], align 4
 ; CHECK-NEXT:    call void @unknown_func()
-; CHECK-NEXT:    br i1 %arg, label [[BB1]], label [[BB2:%.*]]
+; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB1]], label [[BB2:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    ret i32 0
 ;