Merge pull request swiftlang#62937 from jckarter/nonescaping-closure-move-only

[WIP] Treat nonescaping closure types as owned types in SIL; use borrowing to model capture lifetimes

Author: gottesmm

include/swift/SIL/OwnershipUtils.h

@@ -285,6 +285,8 @@ class BorrowingOperandKind {
     Apply,
     TryApply,
     Yield,
+    PartialApplyStack,
+    BeginAsyncLet,
   };
 
 private:
@@ -295,8 +297,8 @@ class BorrowingOperandKind {
 
   operator Kind() const { return value; }
 
-  static BorrowingOperandKind get(SILInstructionKind kind) {
-    switch (kind) {
+  static BorrowingOperandKind get(SILInstruction *i) {
+    switch (i->getKind()) {
     default:
       return Kind::Invalid;
     case SILInstructionKind::BeginBorrowInst:
@@ -311,6 +313,15 @@ class BorrowingOperandKind {
       return Kind::TryApply;
     case SILInstructionKind::YieldInst:
       return Kind::Yield;
+    case SILInstructionKind::PartialApplyInst:
+      return Kind::PartialApplyStack;
+    case SILInstructionKind::BuiltinInst: {
+      auto bi = cast<BuiltinInst>(i);
+      if (bi->getBuiltinKind() == BuiltinValueKind::StartAsyncLetWithLocalBuffer) {
+        return Kind::BeginAsyncLet;
+      }
+      return Kind::Invalid;
+    }
     }
   }
 
@@ -338,7 +349,7 @@ struct BorrowingOperand {
   BorrowingOperandKind kind;
 
   BorrowingOperand(Operand *op)
-      : op(op), kind(BorrowingOperandKind::get(op->getUser()->getKind())) {
+      : op(op), kind(BorrowingOperandKind::get(op->getUser())) {
     auto ownership = op->getOperandOwnership();
     if (ownership != OperandOwnership::Borrow
         && ownership != OperandOwnership::Reborrow) {
@@ -405,6 +416,8 @@ struct BorrowingOperand {
     case BorrowingOperandKind::Apply:
     case BorrowingOperandKind::TryApply:
     case BorrowingOperandKind::Yield:
+    case BorrowingOperandKind::PartialApplyStack:
+    case BorrowingOperandKind::BeginAsyncLet:
       return false;
     case BorrowingOperandKind::Branch:
       return true;
@@ -435,6 +448,8 @@ struct BorrowingOperand {
     case BorrowingOperandKind::Apply:
     case BorrowingOperandKind::TryApply:
     case BorrowingOperandKind::Yield:
+    case BorrowingOperandKind::PartialApplyStack:
+    case BorrowingOperandKind::BeginAsyncLet:
       return false;
     }
     llvm_unreachable("Covered switch isn't covered?!");

include/swift/SIL/SILBuilder.h

@@ -2760,7 +2760,7 @@ class SILBuilder {
       SILLocation Loc, NormalDifferentiableFunctionTypeComponent Extractee,
       SILValue Function, Optional<SILType> ExtracteeType = None) {
     return createDifferentiableFunctionExtract(
-        Loc, Extractee, Function, Function->getOwnershipKind(), ExtracteeType);
+        Loc, Extractee, Function, OwnershipKind::Guaranteed, ExtracteeType);
   }
 
   DifferentiableFunctionExtractInst *createDifferentiableFunctionExtract(

include/swift/SIL/SILCloner.h

@@ -1835,6 +1835,13 @@ template <typename ImplClass>
 void SILCloner<ImplClass>::visitCopyValueInst(CopyValueInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   if (!getBuilder().hasOwnership()) {
+    // Noescape closures become trivial after OSSA.
+    if (auto fnTy = Inst->getType().getAs<SILFunctionType>()) {
+      if (fnTy->isTrivialNoEscape()) {
+        return recordFoldedValue(Inst, getOpValue(Inst->getOperand()));
+      }
+    }
+  
     SILValue newValue = getBuilder().emitCopyValueOperation(
         getOpLocation(Inst->getLoc()), getOpValue(Inst->getOperand()));
     return recordFoldedValue(Inst, newValue);
@@ -1950,6 +1957,34 @@ template <typename ImplClass>
 void SILCloner<ImplClass>::visitDestroyValueInst(DestroyValueInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   if (!getBuilder().hasOwnership()) {
+    // Noescape closures become trivial after OSSA.
+    if (auto fnTy = Inst->getOperand()->getType().getAs<SILFunctionType>()) {
+      if (fnTy->isTrivialNoEscape()) {
+        // Destroying the partial_apply [stack] becomes the stack deallocation
+        // of the context.
+        // Look through mark_dependence and other wrapper instructions.
+        SILValue deallocOperand = Inst->getOperand();
+        while (true) {
+          if (auto mdi = dyn_cast<MarkDependenceInst>(deallocOperand)) {
+            deallocOperand = mdi->getValue();
+          } else if (isa<ConvertEscapeToNoEscapeInst>(deallocOperand)) {
+            break;
+          } else if (auto conv = dyn_cast<ConversionInst>(deallocOperand)) {
+            deallocOperand = conv->getConverted();
+          } else {
+            break;
+          }
+        }
+        if (isa<PartialApplyInst>(deallocOperand)) {
+          recordClonedInstruction(Inst,
+            getBuilder().createDeallocStack(getOpLocation(Inst->getLoc()),
+                                            getOpValue(deallocOperand)));
+        }
+        
+        return;
+      }
+    }
+  
     return recordClonedInstruction(
         Inst, getBuilder().createReleaseValue(
                   getOpLocation(Inst->getLoc()), getOpValue(Inst->getOperand()),

include/swift/SIL/SILInstruction.h

@@ -3019,6 +3019,9 @@ class PartialApplyInst final
   OnStackKind isOnStack() const {
     return getFunctionType()->isNoEscape() ? OnStack : NotOnStack;
   }
+  
+  /// Visit the instructions that end the lifetime of an OSSA on-stack closure.
+  bool visitOnStackLifetimeEnds(llvm::function_ref<bool (Operand*)> func) const;
 };
 
 class EndApplyInst;

include/swift/SIL/SILType.h

@@ -631,7 +631,7 @@ class SILType {
 
   /// Returns true if and only if this type is a first class move only
   /// type. NOTE: Returns false if the type is a move only wrapped type.
-  bool isMoveOnlyType() const;
+  bool isMoveOnlyNominalType() const;
 
   /// Returns true if this SILType is a move only wrapper type.
   ///

include/swift/SILOptimizer/Analysis/AccessSummaryAnalysis.h

@@ -122,6 +122,7 @@ class AccessSummaryAnalysis : public BottomUpIPAnalysis {
     unsigned getArgumentCount() const { return ArgAccesses.size(); }
 
     void print(raw_ostream &os, SILFunction *fn) const;
+    void dump(SILFunction *fn) const;
   };
 
   class FunctionInfo;

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -275,7 +275,8 @@ void releasePartialApplyCapturedArg(
 void insertDestroyOfCapturedArguments(
     PartialApplyInst *pai, SILBuilder &builder,
     llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
-        [](SILValue arg) -> bool { return true; });
+        [](SILValue arg) -> bool { return true; },
+    SILLocation loc = RegularLocation::getAutoGeneratedLocation());
 
 void insertDeallocOfCapturedArguments(PartialApplyInst *pai,
                                       DominanceInfo *domInfo);

lib/SIL/IR/OperandOwnership.cpp

@@ -549,9 +549,21 @@ OperandOwnershipClassifier::visitTryApplyInst(TryApplyInst *i) {
 
 OperandOwnership
 OperandOwnershipClassifier::visitPartialApplyInst(PartialApplyInst *i) {
-  // partial_apply [stack] does not take ownership of its operands.
+  // partial_apply [stack] borrows its operands.
   if (i->isOnStack()) {
-    return OperandOwnership::InstantaneousUse;
+    auto operandTy = getValue()->getType();
+    // Trivial values we can treat as trivial uses.
+    if (operandTy.isTrivial(*i->getFunction())) {
+      return OperandOwnership::TrivialUse;
+    }
+    
+    // Borrowing of address operands is ultimately handled by the move-only
+    // address checker and/or exclusivity checker rather than by value ownership.
+    if (operandTy.isAddress()) {
+      return OperandOwnership::TrivialUse;
+    }
+  
+    return OperandOwnership::Borrow;
   }
   // All non-trivial types should be captured.
   return OperandOwnership::ForwardingConsume;
@@ -833,15 +845,26 @@ OperandOwnership OperandOwnershipBuiltinClassifier::visitCopy(BuiltinInst *bi,
   }
 }
 BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, StartAsyncLet)
-BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, EndAsyncLet)
-BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, StartAsyncLetWithLocalBuffer)
-BUILTIN_OPERAND_OWNERSHIP(DestroyingConsume, EndAsyncLetLifetime)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, EndAsyncLet)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, EndAsyncLetLifetime)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroup)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroupWithFlags)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, DestroyTaskGroup)
 
 BUILTIN_OPERAND_OWNERSHIP(ForwardingConsume, COWBufferForReading)
 
+OperandOwnership
+OperandOwnershipBuiltinClassifier
+::visitStartAsyncLetWithLocalBuffer(BuiltinInst *bi, StringRef attr) {
+  if (&op == &bi->getOperandRef(0)) {
+    // The result buffer pointer is a trivial use.
+    return OperandOwnership::TrivialUse;
+  }
+  
+  // The closure is borrowed while the async let task is executing.
+  return OperandOwnership::Borrow;
+}
+
 const int PARAMETER_INDEX_CREATE_ASYNC_TASK_FUTURE_FUNCTION = 2;
 const int PARAMETER_INDEX_CREATE_ASYNC_TASK_GROUP_FUTURE_FUNCTION = 3;
 
@@ -900,7 +923,7 @@ visitResumeThrowingContinuationThrowing(BuiltinInst *bi, StringRef attr) {
   return OperandOwnership::TrivialUse;
 }
 
-BUILTIN_OPERAND_OWNERSHIP(TrivialUse, TaskRunInline)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, TaskRunInline)
 
 BUILTIN_OPERAND_OWNERSHIP(InteriorPointer, CancelAsyncTask)
 BUILTIN_OPERAND_OWNERSHIP(InteriorPointer, InitializeDefaultActor)

lib/SIL/IR/SILFunctionType.cpp

@@ -1212,10 +1212,20 @@ class Conventions {
                                 unsigned index, const AbstractionPattern &type,
                                 const TypeLowering &substTL) const {
     switch (ownership) {
-    case ValueOwnership::Default:
+    case ValueOwnership::Default: {
       if (forSelf)
         return getDirectSelfParameter(type);
-      return getDirectParameter(index, type, substTL);
+      auto convention = getDirectParameter(index, type, substTL);
+      // Nonescaping closures can only be borrowed across calls currently.
+      if (convention == ParameterConvention::Direct_Owned) {
+        if (auto fnTy = substTL.getLoweredType().getAs<SILFunctionType>()) {
+          if (fnTy->isTrivialNoEscape()) {
+            return ParameterConvention::Direct_Guaranteed;
+          }
+        }
+      }
+      return convention;
+    }
     case ValueOwnership::InOut:
       return ParameterConvention::Indirect_Inout;
     case ValueOwnership::Shared:

lib/SIL/IR/SILInstruction.cpp

@@ -1239,8 +1239,15 @@ bool SILInstruction::isAllocatingStack() const {
       return true;
   }
 
-  if (auto *PA = dyn_cast<PartialApplyInst>(this))
-    return PA->isOnStack();
+  // In OSSA, PartialApply is modeled as a value which borrows its operands
+  // and whose lifetime is ended by a `destroy_value`.
+  //
+  // After OSSA, we make the memory allocation and dependencies explicit again,
+  // with a `dealloc_stack` ending the closure's lifetime.
+  if (auto *PA = dyn_cast<PartialApplyInst>(this)) {
+    return PA->isOnStack()
+      && !PA->getFunction()->hasOwnership();
+  }
 
   if (auto *BI = dyn_cast<BuiltinInst>(this)) {
     if (BI->getBuiltinKind() == BuiltinValueKind::StackAlloc ||
@@ -1658,6 +1665,46 @@ bool SILInstruction::maySuspend() const {
   return false;
 }
 
+static bool visitRecursivelyLifetimeEndingUses(
+  SILValue i,
+  bool &noUsers,
+  llvm::function_ref<bool(Operand *)> func)
+{
+  for (Operand *use : i->getConsumingUses()) {
+    noUsers = false;
+    if (isa<DestroyValueInst>(use->getUser())) {
+      if (!func(use)) {
+        return false;
+      }
+      continue;
+    }
+    
+    // There shouldn't be any dead-end consumptions of a nonescaping
+    // partial_apply that don't forward it along, aside from destroy_value.
+    assert(use->getUser()->hasResults());
+    for (auto result : use->getUser()->getResults()) {
+      if (!visitRecursivelyLifetimeEndingUses(result, noUsers, func)) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+bool
+PartialApplyInst::visitOnStackLifetimeEnds(
+                             llvm::function_ref<bool (Operand *)> func) const {
+  assert(getFunction()->hasOwnership()
+         && isOnStack()
+         && "only meaningful for OSSA stack closures");
+  bool noUsers = true;
+
+  if (!visitRecursivelyLifetimeEndingUses(this, noUsers, func)) {
+    return false;
+  }
+  return !noUsers;
+}
+
 #ifndef NDEBUG
 
 //---