[IR] Add dead_on_unwind attribute (#74289)

Add the `dead_on_unwind` attribute, which states that the caller will
not read from this argument if the call unwinds. This allows eliding
stores that could otherwise be visible on the unwind path, for example:

```
declare void @may_unwind()

define void @src(ptr noalias dead_on_unwind %out) {
    store i32 0, ptr %out
    call void @may_unwind()
    store i32 1, ptr %out
    ret void
}

define void @tgt(ptr noalias dead_on_unwind %out) {
    call void @may_unwind()
    store i32 1, ptr %out
    ret void
}
```

The optimization is not valid without `dead_on_unwind`, because the `i32
0` value might be read if `@may_unwind` unwinds.

This attribute is primarily intended to be used on sret arguments. In
fact, I previously wanted to change the semantics of sret to include
this "no read after unwind" property (see D116998), but based on the
feedback there it is better to keep these attributes orthogonal (sret is
an ABI attribute, dead_on_unwind is an optimization attribute). This is
a reboot of that change with a separate attribute.

Author: nikic

llvm/docs/LangRef.rst

@@ -1576,6 +1576,20 @@ Currently, only the following parameter attributes are defined:
     ``readonly`` or a ``memory`` attribute that does not contain
     ``argmem: write``.
 
+``dead_on_unwind``
+    At a high level, this attribute indicates that the pointer argument is dead
+    if the call unwinds, in the sense that the caller will not depend on the
+    contents of the memory. Stores that would only be visible on the unwind
+    path can be elided.
+
+    More precisely, the behavior is as-if any memory written through the
+    pointer during the execution of the function is overwritten with a poison
+    value on unwind. This includes memory written by the implicit write implied
+    by the ``writable`` attribute. The caller is allowed to access the affected
+    memory, but all loads that are not preceded by a store will return poison.
+
+    This attribute cannot be applied to return values.
+
 .. _gc:
 
 Garbage Collector Strategy Names

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -723,6 +723,7 @@ enum AttributeKindCodes {
   ATTR_KIND_OPTIMIZE_FOR_DEBUGGING = 88,
   ATTR_KIND_WRITABLE = 89,
   ATTR_KIND_CORO_ONLY_DESTROY_WHEN_COMPLETE = 90,
+  ATTR_KIND_DEAD_ON_UNWIND = 91,
 };
 
 enum ComdatSelectionKindCodes {

llvm/include/llvm/IR/Attributes.td

@@ -151,6 +151,9 @@ def NoDuplicate : EnumAttr<"noduplicate", [FnAttr]>;
 /// Function does not deallocate memory.
 def NoFree : EnumAttr<"nofree", [FnAttr, ParamAttr]>;
 
+/// Argument is dead if the call unwinds.
+def DeadOnUnwind : EnumAttr<"dead_on_unwind", [ParamAttr]>;
+
 /// Disable implicit floating point insts.
 def NoImplicitFloat : EnumAttr<"noimplicitfloat", [FnAttr]>;
 

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -896,7 +896,7 @@ bool llvm::isNotVisibleOnUnwind(const Value *Object,
 
   // Byval goes out of scope on unwind.
   if (auto *A = dyn_cast<Argument>(Object))
-    return A->hasByValAttr();
+    return A->hasByValAttr() || A->hasAttribute(Attribute::DeadOnUnwind);
 
   // A noalias return is not accessible from any other code. If the pointer
   // does not escape prior to the unwind, then the caller cannot access the

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -2098,6 +2098,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::Writable;
   case bitc::ATTR_KIND_CORO_ONLY_DESTROY_WHEN_COMPLETE:
     return Attribute::CoroDestroyOnlyWhenComplete;
+  case bitc::ATTR_KIND_DEAD_ON_UNWIND:
+    return Attribute::DeadOnUnwind;
   }
 }
 

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -827,6 +827,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_WRITABLE;
   case Attribute::CoroDestroyOnlyWhenComplete:
     return bitc::ATTR_KIND_CORO_ONLY_DESTROY_WHEN_COMPLETE;
+  case Attribute::DeadOnUnwind:
+    return bitc::ATTR_KIND_DEAD_ON_UNWIND;
   case Attribute::EndAttrKinds:
     llvm_unreachable("Can not encode end-attribute kinds marker.");
   case Attribute::None:

llvm/lib/IR/Attributes.cpp

@@ -1962,7 +1962,8 @@ AttributeMask AttributeFuncs::typeIncompatible(Type *Ty,
           .addAttribute(Attribute::ReadOnly)
           .addAttribute(Attribute::Dereferenceable)
           .addAttribute(Attribute::DereferenceableOrNull)
-          .addAttribute(Attribute::Writable);
+          .addAttribute(Attribute::Writable)
+          .addAttribute(Attribute::DeadOnUnwind);
     if (ASK & ASK_UNSAFE_TO_DROP)
       Incompatible.addAttribute(Attribute::Nest)
           .addAttribute(Attribute::SwiftError)

llvm/lib/Transforms/Utils/CodeExtractor.cpp

@@ -998,6 +998,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
       case Attribute::ByRef:
       case Attribute::WriteOnly:
       case Attribute::Writable:
+      case Attribute::DeadOnUnwind:
       //  These are not really attributes.
       case Attribute::None:
       case Attribute::EndAttrKinds:

llvm/test/Bitcode/attributes.ll

@@ -521,6 +521,11 @@ define void @f90(ptr writable %p) {
   ret void
 }
 
+; CHECK: define void @f91(ptr dead_on_unwind %p)
+define void @f91(ptr dead_on_unwind %p) {
+  ret void
+}
+
 ; CHECK: attributes #0 = { noreturn }
 ; CHECK: attributes #1 = { nounwind }
 ; CHECK: attributes #2 = { memory(none) }

llvm/test/Transforms/DeadStoreElimination/simple.ll

@@ -514,13 +514,12 @@ define void @test34(ptr noalias %p) {
   store i32 0, ptr %p
   ret void
 }
-; Same as previous case, but with an sret argument.
-; TODO: The first store could be eliminated if sret is not visible on unwind.
-define void @test34_sret(ptr noalias sret(i32) %p) {
-; CHECK-LABEL: @test34_sret(
-; CHECK-NEXT:    store i32 1, ptr [[P:%.*]], align 4
+
+; Same as previous case, but with a dead_on_unwind argument.
+define void @test34_dead_on_unwind(ptr noalias dead_on_unwind %p) {
+; CHECK-LABEL: @test34_dead_on_unwind(
 ; CHECK-NEXT:    call void @unknown_func()
-; CHECK-NEXT:    store i32 0, ptr [[P]], align 4
+; CHECK-NEXT:    store i32 0, ptr [[P:%.*]], align 4
 ; CHECK-NEXT:    ret void
 ;
   store i32 1, ptr %p

llvm/test/Transforms/LICM/scalar-promote-unwind.ll

@@ -147,17 +147,15 @@ for.cond.cleanup:
   ret void
 }
 
-; TODO: sret could be specified to not be accessed on unwind either.
-define void @test_sret(ptr noalias sret(i32) %a, i1 zeroext %y) uwtable {
-; CHECK-LABEL: @test_sret(
+define void @test_dead_on_unwind(ptr noalias dead_on_unwind %a, i1 zeroext %y) uwtable {
+; CHECK-LABEL: @test_dead_on_unwind(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[A_PROMOTED:%.*]] = load i32, ptr [[A:%.*]], align 4
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[ADD1:%.*]] = phi i32 [ [[A_PROMOTED]], [[ENTRY:%.*]] ], [ [[ADD:%.*]], [[FOR_INC:%.*]] ]
 ; CHECK-NEXT:    [[I_03:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[INC:%.*]], [[FOR_INC]] ]
 ; CHECK-NEXT:    [[ADD]] = add nsw i32 [[ADD1]], 1
-; CHECK-NEXT:    store i32 [[ADD]], ptr [[A]], align 4
 ; CHECK-NEXT:    br i1 [[Y:%.*]], label [[IF_THEN:%.*]], label [[FOR_INC]]
 ; CHECK:       if.then:
 ; CHECK-NEXT:    tail call void @f()
@@ -167,6 +165,8 @@ define void @test_sret(ptr noalias sret(i32) %a, i1 zeroext %y) uwtable {
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i32 [[INC]], 10000
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
 ; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[ADD_LCSSA:%.*]] = phi i32 [ [[ADD]], [[FOR_INC]] ]
+; CHECK-NEXT:    store i32 [[ADD_LCSSA]], ptr [[A]], align 4
 ; CHECK-NEXT:    ret void
 ;
 entry:

llvm/test/Transforms/MemCpyOpt/callslot_throw.ll

@@ -56,9 +56,39 @@ entry:
   ret void
 }
 
-; TODO: With updated semantics, sret could also be invisible on unwind.
-define void @test_sret(ptr nocapture noalias dereferenceable(4) sret(i32) %x) {
-; CHECK-LABEL: @test_sret(
+define void @test_dead_on_unwind(ptr nocapture noalias writable dead_on_unwind dereferenceable(4) %x) {
+; CHECK-LABEL: @test_dead_on_unwind(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    call void @may_throw(ptr nonnull [[X:%.*]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %t = alloca i32, align 4
+  call void @may_throw(ptr nonnull %t)
+  %load = load i32, ptr %t, align 4
+  store i32 %load, ptr %x, align 4
+  ret void
+}
+
+; Same as previous test, but dereferenceability information is provided by sret.
+define void @test_dead_on_unwind_sret(ptr nocapture noalias writable dead_on_unwind sret(i32) %x) {
+; CHECK-LABEL: @test_dead_on_unwind_sret(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[T:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    call void @may_throw(ptr nonnull [[X:%.*]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %t = alloca i32, align 4
+  call void @may_throw(ptr nonnull %t)
+  %load = load i32, ptr %t, align 4
+  store i32 %load, ptr %x, align 4
+  ret void
+}
+
+define void @test_dead_on_unwind_missing_writable(ptr nocapture noalias dead_on_unwind dereferenceable(4) %x) {
+; CHECK-LABEL: @test_dead_on_unwind_missing_writable(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[T:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    call void @may_throw(ptr nonnull [[T]])

llvm/test/Verifier/dead-on-unwind.ll

@@ -0,0 +1,7 @@
+; RUN: not llvm-as -disable-output %s 2>&1 | FileCheck %s
+
+; CHECK: Attribute 'dead_on_unwind' applied to incompatible type!
+; CHECK-NEXT: ptr @not_pointer
+define void @not_pointer(i32 dead_on_unwind %arg) {
+  ret void
+}