Merge pull request #66779 from jckarter/noncopyable-address-only-borrowed-capture-5.9

jckarter · web-flow · commit 2d1cbb675b20 · 2023-06-20T19:44:41.000-07:00
[5.9] ClosureLifetimeFixup: Remove copy of borrowed move-only nonescaping captures when possible.
diff --git a/include/swift/SILOptimizer/Utils/InstOptUtils.h b/include/swift/SILOptimizer/Utils/InstOptUtils.h
@@ -289,7 +289,9 @@ void insertDestroyOfCapturedArguments(
     SILLocation loc = RegularLocation::getAutoGeneratedLocation());
 
 void insertDeallocOfCapturedArguments(PartialApplyInst *pai,
-                                      DominanceInfo *domInfo);
+    DominanceInfo *domInfo,
+    llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
+      [](SILValue arg) -> bool { return true; });
 
 /// This iterator 'looks through' one level of builtin expect users exposing all
 /// users of the looked through builtin expect instruction i.e it presents a
diff --git a/lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp b/lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp
@@ -15,9 +15,11 @@
 #include "swift/Basic/Defer.h"
 #include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/InstructionUtils.h"
+#include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SIL/SILValue.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
@@ -712,7 +714,171 @@ static SILValue tryRewriteToPartialApplyStack(
   // lifetime ends.
   SmallVector<SILInstruction *, 4> lifetimeEnds;
   collectStackClosureLifetimeEnds(lifetimeEnds, closureOp);
+  
+  // For noncopyable address-only captures, see if we can eliminate the copy
+  // that SILGen emitted to allow the original partial_apply to take ownership.
+  // We do this here because otherwise the move checker will see the copy as an
+  // attempt to consume the value, which we don't want.
+  SmallVector<SILBasicBlock *, 8> discoveredBlocks;
+  SSAPrunedLiveness closureLiveness(cvt->getFunction(), &discoveredBlocks);
+  closureLiveness.initializeDef(closureOp);
+  
+  SmallSetVector<SILValue, 4> borrowedOriginals;
+  
+  for (unsigned i : indices(newPA->getArgumentOperands())) {
+    auto &arg = newPA->getArgumentOperands()[i];
+    SILValue copy = arg.get();
+    // The temporary should be a local stack allocation.
+    LLVM_DEBUG(llvm::dbgs() << "considering whether to eliminate copy of capture\n";
+               copy->printInContext(llvm::dbgs());
+               llvm::dbgs() << "\n");
+
+    auto stack = dyn_cast<AllocStackInst>(copy);
+    if (!stack) {
+      LLVM_DEBUG(llvm::dbgs() << "-- not an alloc_stack\n");
+      continue;
+    }
+    
+    // This would be a nice optimization to attempt for all types, but for now,
+    // limit the effect to move-only types.
+    if (!copy->getType().isMoveOnly()) {
+      LLVM_DEBUG(llvm::dbgs() << "-- not move-only\n");
+      continue;
+    }
+    
+    // Is the capture a borrow?
+    auto paramIndex = newPA
+      ->getArgumentIndexForOperandIndex(i + newPA->getArgumentOperandNumber())
+      .getValue();
+    if (!newPA->getOrigCalleeType()->getParameters()[paramIndex]
+          .isIndirectInGuaranteed()) {
+      LLVM_DEBUG(llvm::dbgs() << "-- not an in_guaranteed parameter\n";
+                 newPA->getOrigCalleeType()->getParameters()[paramIndex]
+                   .print(llvm::dbgs());
+                 llvm::dbgs() << "\n");
+      continue;
+    }
     
+    // It needs to have been initialized by copying from somewhere else.
+    CopyAddrInst *initialization = nullptr;
+    MarkDependenceInst *markDep = nullptr;
+    for (auto *use : stack->getUses()) {
+      // Since we removed the `dealloc_stack`s from the capture arguments,
+      // the only uses of this stack slot should be the initialization, the
+      // partial application, and possibly a mark_dependence from the
+      // buffer to the partial application.
+      if (use->getUser() == newPA) {
+        continue;
+      }
+      if (auto mark = dyn_cast<MarkDependenceInst>(use->getUser())) {
+        // If we're marking dependence of the current partial_apply on this
+        // stack slot, that's fine.
+        if (mark->getValue() != newPA
+            || mark->getBase() != stack) {
+          LLVM_DEBUG(llvm::dbgs() << "-- had unexpected mark_dependence use\n";
+                     use->getUser()->print(llvm::dbgs());
+                     llvm::dbgs() << "\n");
+          
+          break;
+        }
+        markDep = mark;
+        continue;
+      }
+      
+      // If we saw more than just the initialization, this isn't a pattern we
+      // recognize.
+      if (initialization) {
+        LLVM_DEBUG(llvm::dbgs() << "-- had non-initialization, non-partial-apply use\n";
+                   use->getUser()->print(llvm::dbgs());
+                   llvm::dbgs() << "\n");
+                   
+        initialization = nullptr;
+        break;
+      }
+      if (auto possibleInit = dyn_cast<CopyAddrInst>(use->getUser())) {
+        // Should copy the source and initialize the destination.
+        if (possibleInit->isTakeOfSrc()
+            || !possibleInit->isInitializationOfDest()) {
+          LLVM_DEBUG(llvm::dbgs() << "-- had non-initialization, non-partial-apply use\n";
+                     use->getUser()->print(llvm::dbgs());
+                     llvm::dbgs() << "\n");
+
+          break;
+        }
+        // This is the initialization if there are no other uses.
+        initialization = possibleInit;
+        continue;
+      }
+      LLVM_DEBUG(llvm::dbgs() << "-- unrecognized use\n");
+      break;
+    }
+    if (!initialization) {
+      LLVM_DEBUG(llvm::dbgs() << "-- failed to find single initializing use\n");
+      continue;
+    }
+    
+    // The source should have no writes in the duration of the partial_apply's
+    // liveness.
+    auto orig = initialization->getSrc();
+    LLVM_DEBUG(llvm::dbgs() << "++ found original:\n";
+               orig->print(llvm::dbgs());
+               llvm::dbgs() << "\n");
+               
+    bool origIsUnusedDuringClosureLifetime = true;
+
+    class OrigUnusedDuringClosureLifetimeWalker final
+      : public TransitiveAddressWalker
+    {
+      SSAPrunedLiveness &closureLiveness;
+      bool &origIsUnusedDuringClosureLifetime;
+    public:
+      OrigUnusedDuringClosureLifetimeWalker(SSAPrunedLiveness &closureLiveness,
+                                        bool &origIsUnusedDuringClosureLifetime)
+        : closureLiveness(closureLiveness),
+          origIsUnusedDuringClosureLifetime(origIsUnusedDuringClosureLifetime)
+      {}
+    
+      virtual bool visitUse(Operand *origUse) override {
+        LLVM_DEBUG(llvm::dbgs() << "looking at use\n";
+                   origUse->getUser()->printInContext(llvm::dbgs());
+                   llvm::dbgs() << "\n");
+        
+        // If the user doesn't write to memory, then it's harmless.
+        if (!origUse->getUser()->mayWriteToMemory()) {
+          return true;
+        }
+        if (closureLiveness.isWithinBoundary(origUse->getUser())) {
+          origIsUnusedDuringClosureLifetime = false;
+          LLVM_DEBUG(llvm::dbgs() << "-- original has other possibly writing use during closure lifetime\n";
+                     origUse->getUser()->print(llvm::dbgs());
+                     llvm::dbgs() << "\n");
+          return false;
+        }
+        return true;
+      }
+    };
+    
+    OrigUnusedDuringClosureLifetimeWalker origUseWalker(closureLiveness,
+                                             origIsUnusedDuringClosureLifetime);
+    auto walkResult = std::move(origUseWalker).walk(orig);
+    
+    if (walkResult == AddressUseKind::Unknown
+        || !origIsUnusedDuringClosureLifetime) {
+      continue;
+    }
+
+    // OK, we can use the original. Eliminate the copy and replace it with the
+    // original.
+    LLVM_DEBUG(llvm::dbgs() << "++ replacing with original!\n");
+    arg.set(orig);
+    if (markDep) {
+      markDep->setBase(orig);
+    }
+    initialization->eraseFromParent();
+    stack->eraseFromParent();
+    borrowedOriginals.insert(orig);
+  }
+  
   /* DEBUG
   llvm::errs() << "=== found lifetime ends for\n";
   closureOp->dump();
@@ -724,7 +890,11 @@ static SILValue tryRewriteToPartialApplyStack(
     */
     SILBuilderWithScope builder(std::next(destroy->getIterator()));
     insertDestroyOfCapturedArguments(newPA, builder,
-                                     [&](SILValue arg) -> bool { return true; },
+                                     [&](SILValue arg) -> bool {
+                                       // Don't need to destroy if we borrowed
+                                       // in place .
+                                       return !borrowedOriginals.count(arg);
+                                     },
                                      newPA->getLoc());
   }
   /* DEBUG
@@ -750,7 +920,12 @@ static SILValue tryRewriteToPartialApplyStack(
     return closureOp;
 
   insertDeallocOfCapturedArguments(
-      newPA, dominanceAnalysis->get(closureUser->getFunction()));
+      newPA, dominanceAnalysis->get(closureUser->getFunction()),
+      [&](SILValue arg) -> bool {
+        // Don't need to destroy if we borrowed
+        // in place.
+        return !borrowedOriginals.count(arg);
+      });
 
   return closureOp;
 }
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp
@@ -3140,7 +3140,7 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
   }
 
   // Then gather all uses of our address by walking from def->uses. We use this
-  // to categorize the uses of this address into their ownership behavior (e.x.:
+  // to categorize the uses of this address into their ownership behavior (e.g.:
   // init, reinit, take, destroy, etc.).
   GatherUsesVisitor visitor(*this, addressUseState, markedAddress,
                             diagnosticEmitter);
diff --git a/lib/SILOptimizer/Utils/InstOptUtils.cpp b/lib/SILOptimizer/Utils/InstOptUtils.cpp
@@ -1495,7 +1495,9 @@ void swift::insertDestroyOfCapturedArguments(
 }
 
 void swift::insertDeallocOfCapturedArguments(PartialApplyInst *pai,
-                                             DominanceInfo *domInfo) {
+                         DominanceInfo *domInfo,
+                         llvm::function_ref<bool(SILValue)> shouldInsertDestroy)
+{
   assert(pai->isOnStack());
 
   ApplySite site(pai);
@@ -1509,6 +1511,9 @@ void swift::insertDeallocOfCapturedArguments(PartialApplyInst *pai,
       continue;
 
     SILValue argValue = arg.get();
+    if (!shouldInsertDestroy(argValue)) {
+      continue;
+    }
     if (auto moveWrapper =
             dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(argValue))
       argValue = moveWrapper->getOperand();
diff --git a/test/Interpreter/moveonly_deinit_autoclosure.swift b/test/Interpreter/moveonly_deinit_autoclosure.swift
@@ -0,0 +1,34 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+internal func _myPrecondition(
+  _ body: @autoclosure () -> Bool
+) {
+  guard body() else {
+    fatalError("condition failed")
+  }
+}
+
+var deinitCounter = 0
+
+struct MyCounter<T>: ~Copyable {
+  let expectedCount = 1
+  let box: T
+  init(_ t: T) {
+    self.box = t
+  }
+  deinit {
+    print("hello")
+    deinitCounter += 1
+    _myPrecondition(deinitCounter == self.expectedCount)
+  }
+}
+
+func test() {
+  _ = MyCounter(4343)
+}
+
+// CHECK: hello
+test()
+// CHECK-NEXT: great success
+print("great success")
diff --git a/test/SILOptimizer/moveonly_closure_lifetime_fixup_address_only_borrowed_captures.swift b/test/SILOptimizer/moveonly_closure_lifetime_fixup_address_only_borrowed_captures.swift
@@ -0,0 +1,64 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+// When ClosureLifetimeFixup promotes a closure to the stack, we should
+// eliminate the temporary copy of any borrowed move-only captures, so that
+// they get checked as borrows and not consumes.
+
+struct E<T>: ~Copyable {
+    var t: T
+}
+
+var escaper: () -> () = {}
+
+func nonescaping(_ x: () -> ()) { }
+func escaping(_ x: @escaping () -> ()) {  escaper = x }
+
+func borrow<T>(_: borrowing E<T>) {}
+
+// CHECK-LABEL: sil {{.*}}16testNeverEscaped
+func testNeverEscaped<T>(_ e: borrowing E<T>) {
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}(%0) :
+    nonescaping {
+        borrow(e)
+    }
+}
+
+func nonescaping(_ x: () -> (), and y: () -> ()) {}
+
+// CHECK-LABEL: sil {{.*}}21testNeverEscapedTwice
+func testNeverEscapedTwice<T>(_ e: borrowing E<T>) {
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}(%0) :
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}(%0) :
+    nonescaping({ borrow(e) }, and: { borrow(e) })
+}
+
+// CHECK-LABEL: sil {{.*}}16testEscapedLater
+func testEscapedLater<T>(_ e: consuming E<T>) {
+    // CHECK: [[BOX:%.*]] = alloc_box
+    // CHECK: [[CONTENTS:%.*]] = project_box [[BOX]]
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[CONTENTS]])
+    nonescaping {
+        borrow(e)
+    }
+
+    escaping {
+        borrow(e)
+    }
+}
+
+// CHECK-LABEL: sil {{.*}}17testEscapedLater2
+func testEscapedLater2<T>(_ e: consuming E<T>) {
+    // CHECK: [[BOX:%.*]] = alloc_box
+    // CHECK: [[CONTENTS:%.*]] = project_box [[BOX]]
+    // CHECK: partial_apply {{.*}}[on_stack] {{.*}}([[CONTENTS]])
+    let f = e
+
+    nonescaping {
+        borrow(f)
+    }
+
+    escaping {
+        borrow(f)
+    }
+}
+

Original file line number	Diff line number	Diff line change
`@@ -3140,7 +3140,7 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(`
`3140`	`3140`	`}`
`3141`	`3141`
`3142`	`3142`	`// Then gather all uses of our address by walking from def->uses. We use this`
`3143`		`- // to categorize the uses of this address into their ownership behavior (e.x.:`
	`3143`	`+ // to categorize the uses of this address into their ownership behavior (e.g.:`
`3144`	`3144`	`// init, reinit, take, destroy, etc.).`
`3145`	`3145`	`GatherUsesVisitor visitor(*this, addressUseState, markedAddress,`
`3146`	`3146`	`diagnosticEmitter);`