Handle the mark_dependence chain when eliminating copies in ClosureLifetimeFixup

Author: Azoy

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -734,6 +734,7 @@ static SILValue tryRewriteToPartialApplyStack(
   unsigned appliedArgStartIdx =
         newPA->getOrigCalleeType()->getNumParameters() - newPA->getNumArguments();
 
+  MarkDependenceInst *markDepChain = nullptr;
   for (unsigned i : indices(newPA->getArgumentOperands())) {
     auto &arg = newPA->getArgumentOperands()[i];
     SILValue copy = arg.get();
@@ -781,6 +782,20 @@ static SILValue tryRewriteToPartialApplyStack(
         continue;
       }
       if (auto mark = dyn_cast<MarkDependenceInst>(use->getUser())) {
+        // When we insert mark_dependence for non-trivial address operands, we
+        // emit a chain that looks like:
+        //    %md = mark_dependence %pai on %0
+        //    %md2 = mark_dependence %md on %1
+        // to tie all of those operands together on the same partial_apply.
+        //
+        // FIXME: Should we not be chaining like this and just emit independent
+        // mark_dependence?
+        if (markDepChain && mark->getValue() == markDepChain) {
+          markDep = mark;
+          markDepChain = mark;
+          continue;
+        }
+
         // If we're marking dependence of the current partial_apply on this
         // stack slot, that's fine.
         if (mark->getValue() != newPA
@@ -792,6 +807,11 @@ static SILValue tryRewriteToPartialApplyStack(
           break;
         }
         markDep = mark;
+
+        if (!markDepChain) {
+          markDepChain = mark;
+        }
+
         continue;
       }
 

test/SILOptimizer/closure_lifetime_fixup.swift

@@ -446,3 +446,29 @@ public class F<T> {
 public func testClosureMethodParam<T>(f: F<T>) throws {
   try f.test { return f.t! }
 }
+
+struct AddressOnlyNoncopyableStruct: ~Copyable {
+  let x: Any = 123
+
+  borrowing func hello() {}
+}
+
+func simpleNonescapingClosure(with body: () -> ()) {
+  body()
+}
+
+// CHECK-LABEL: s22closure_lifetime_fixup27trySimpleNonescapingClosure
+// CHECK: [[FIRST:%.*]] = alloc_stack [var_decl] $AddressOnlyNoncopyableStruct, let, name "foo"
+// CHECK: [[SECOND:%.*]] = alloc_stack [var_decl] $AddressOnlyNoncopyableStruct, let, name "bar"
+// CHECK: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] %{{.*}}([[FIRST]], [[SECOND]])
+// CHECK: [[MD_ONE:%.*]] = mark_dependence [nonescaping] [[PA]] : $@noescape @callee_guaranteed () -> () on [[FIRST]] : $*AddressOnlyNoncopyableStruct
+// CHECK: [[MD_TWO:%.*]] = mark_dependence [nonescaping] [[MD_ONE]] : $@noescape @callee_guaranteed () -> () on [[SECOND]] : $*AddressOnlyNoncopyableStruct
+func trySimpleNonescapingClosure() {
+  let foo = AddressOnlyNoncopyableStruct()
+  let bar = AddressOnlyNoncopyableStruct()
+
+  simpleNonescapingClosure {
+    foo.hello() // OK
+    bar.hello() // OK
+  }
+}

test/SILOptimizer/stdlib/Atomics.swift

@@ -4,6 +4,9 @@
 
 import Synchronization
 
+@_silgen_name("testInt")
+func testInt(_: Int)
+
 //===----------------------------------------------------------------------===//
 // Ensure that we don't destroy the atomic before operations
 //===----------------------------------------------------------------------===//
@@ -77,3 +80,38 @@ func deadAtomic() {
   let _ = Atomic(0)
   let _ = Atomic<UnsafeRawPointer?>(nil)
 }
+
+//===----------------------------------------------------------------------===//
+// Closure Lifetime Fixup
+//===----------------------------------------------------------------------===//
+
+func nonescapingClosure(with body: () -> ()) {
+  body()
+}
+
+// CHECK-LABEL: sil {{.*}} @testNonescapingClosure {{.*}} {
+// CHECK:         {{%.*}} = alloc_stack [lexical] [var_decl] $Atomic<Int>, let, name "foo"
+// CHECK:         {{%.*}} = alloc_stack [lexical] [var_decl] $Atomic<Int>, let, name "bar"
+// CHECK:         builtin "atomicrmw_add_monotonic_Int[[PTR_SIZE]]"
+// CHECK:         builtin "atomicrmw_add_monotonic_Int[[PTR_SIZE]]"
+
+// Make sure there are no moves
+// CHECK-NOT:     alloc_stack $Atomic<Int>
+
+// Make sure we don't emit a partial application
+// CHECK-NOT:     partial_apply
+
+// CHECK-LABEL: } // end sil function 'testNonescapingClosure'
+@_silgen_name("testNonescapingClosure")
+func testNonescapingClosure() {
+  let foo = Atomic(0)
+  let bar = Atomic(1)
+
+  nonescapingClosure {
+    foo.wrappingAdd(1, ordering: .relaxed)
+    bar.wrappingAdd(1, ordering: .relaxed)
+  }
+
+  testInt(foo.load(ordering: .relaxed)) // OK
+  testInt(bar.load(ordering: .relaxed)) // OK
+}