Merge pull request #21933 from aschwaighofer/partial_apply_stack

Use stack allocation for Swift closures

Author: aschwaighofer

docs/SIL.rst

@@ -3347,11 +3347,12 @@ partial_apply
 `````````````
 ::
 
-  sil-instruction ::= 'partial_apply' callee-ownership-attr? sil-value
+  sil-instruction ::= 'partial_apply' callee-ownership-attr? on-stack-attr? sil-value
                         sil-apply-substitution-list?
                         '(' (sil-value (',' sil-value)*)? ')'
                         ':' sil-type
   callee-ownership-attr ::= '[callee_guaranteed]'
+  on-stack-attr ::= '[on_stack]'
 
   %c = partial_apply %0(%1, %2, ...) : $(Z..., A, B, ...) -> R
   // Note that the type of the callee '%0' is specified *after* the arguments
@@ -3369,12 +3370,18 @@ partial_apply
 Creates a closure by partially applying the function ``%0`` to a partial
 sequence of its arguments. In the instruction syntax, the type of the callee is
 specified after the argument list; the types of the argument and of the defined
-value are derived from the function type of the callee. The closure context will
-be allocated with retain count 1 and initialized to contain the values ``%1``,
+value are derived from the function type of the callee. If the ``partial_apply``
+has an escaping function type (not ``[on_stack]``) the closure context will be
+allocated with retain count 1 and initialized to contain the values ``%1``,
 ``%2``, etc.  The closed-over values will not be retained; that must be done
-separately before the ``partial_apply``. The closure does however take
-ownership of the partially applied arguments; when the closure reference
-count reaches zero, the contained values will be destroyed.
+separately before the ``partial_apply``. The closure does however take ownership
+of the partially applied arguments; when the closure reference count reaches
+zero, the contained values will be destroyed. If the ``partial_apply`` has a
+``@noescape`` function type (``partial_apply [on_stack]``) the closure context
+is allocated on the stack and intialized to contain the closed-over values. The
+closed-over values are not retained, lifetime of the closed-over values must be
+managed separately. The lifetime of the stack context of a ``partial_apply
+[stack]`` must be terminated with a ``dealloc_stack``.
 
 If the callee is generic, all of its generic parameters must be bound by the
 given substitution list. The arguments are given with these generic

include/swift/SIL/SILBuilder.h

@@ -327,10 +327,12 @@ class SILBuilder {
   // Type remapping
   //===--------------------------------------------------------------------===//
 
-  static SILType getPartialApplyResultType(SILType Ty, unsigned ArgCount,
-                                         SILModule &M,
-                                         SubstitutionMap subs,
-                                         ParameterConvention calleeConvention);
+  static SILType
+  getPartialApplyResultType(SILType Ty, unsigned ArgCount, SILModule &M,
+                            SubstitutionMap subs,
+                            ParameterConvention calleeConvention,
+                            PartialApplyInst::OnStackKind onStack =
+                                PartialApplyInst::OnStackKind::NotOnStack);
 
   //===--------------------------------------------------------------------===//
   // CFG Manipulation
@@ -476,10 +478,12 @@ class SILBuilder {
   PartialApplyInst *createPartialApply(
       SILLocation Loc, SILValue Fn, SubstitutionMap Subs,
       ArrayRef<SILValue> Args, ParameterConvention CalleeConvention,
+      PartialApplyInst::OnStackKind OnStack =
+          PartialApplyInst::OnStackKind::NotOnStack,
       const GenericSpecializationInformation *SpecializationInfo = nullptr) {
     return insert(PartialApplyInst::create(
         getSILDebugLocation(Loc), Fn, Args, Subs, CalleeConvention, *F,
-        C.OpenedArchetypes, SpecializationInfo));
+        C.OpenedArchetypes, SpecializationInfo, OnStack));
   }
 
   BeginApplyInst *createBeginApply(

include/swift/SIL/SILCloner.h

@@ -878,6 +878,7 @@ SILCloner<ImplClass>::visitPartialApplyInst(PartialApplyInst *Inst) {
                 getOpLocation(Inst->getLoc()), getOpValue(Inst->getCallee()),
                 getOpSubstitutionMap(Inst->getSubstitutionMap()), Args,
                 Inst->getType().getAs<SILFunctionType>()->getCalleeConvention(),
+                Inst->isOnStack(),
                 GenericSpecializationInformation::create(Inst, getBuilder())));
 }
 

include/swift/SIL/SILInstruction.h

@@ -2072,6 +2072,12 @@ class PartialApplyInst final
       public llvm::TrailingObjects<PartialApplyInst, Operand> {
   friend SILBuilder;
 
+public:
+  enum OnStackKind {
+    NotOnStack, OnStack
+  };
+
+private:
   PartialApplyInst(SILDebugLocation DebugLoc, SILValue Callee,
                    SILType SubstCalleeType,
                    SubstitutionMap Substitutions,
@@ -2084,7 +2090,8 @@ class PartialApplyInst final
   create(SILDebugLocation DebugLoc, SILValue Callee, ArrayRef<SILValue> Args,
          SubstitutionMap Substitutions, ParameterConvention CalleeConvention,
          SILFunction &F, SILOpenedArchetypesState &OpenedArchetypes,
-         const GenericSpecializationInformation *SpecializationInfo);
+         const GenericSpecializationInformation *SpecializationInfo,
+         OnStackKind onStack);
 
 public:
   /// Return the result function type of this partial apply.
@@ -2094,6 +2101,10 @@ class PartialApplyInst final
   bool hasCalleeGuaranteedContext() const {
     return getType().castTo<SILFunctionType>()->isCalleeGuaranteed();
   }
+
+  OnStackKind isOnStack() const {
+    return getFunctionType()->isNoEscape() ? OnStack : NotOnStack;
+  }
 };
 
 class BeginApplyInst;

include/swift/SIL/TypeSubstCloner.h

@@ -215,8 +215,8 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     PartialApplyInst *N = getBuilder().createPartialApply(
         getOpLocation(Inst->getLoc()), Helper.getCallee(),
         Helper.getSubstitutions(), Helper.getArguments(), ParamConvention,
-        GenericSpecializationInformation::create(
-          Inst, getBuilder()));
+        Inst->isOnStack(),
+        GenericSpecializationInformation::create(Inst, getBuilder()));
     recordClonedInstruction(Inst, N);
   }
 

include/swift/SILOptimizer/Utils/Local.h

@@ -212,6 +212,12 @@ void releasePartialApplyCapturedArg(
     SILBuilder &Builder, SILLocation Loc, SILValue Arg, SILParameterInfo PInfo,
     InstModCallbacks Callbacks = InstModCallbacks());
 
+/// Insert destroys of captured arguments of partial_apply [stack].
+void insertDestroyOfCapturedArguments(
+    PartialApplyInst *PAI, SILBuilder &B,
+    llvm::function_ref<bool(SILValue)> shouldInsertDestroy =
+        [](SILValue arg) -> bool { return true; });
+
 /// This computes the lifetime of a single SILValue.
 ///
 /// This does not compute a set of jointly postdominating use points. Instead it

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -98,10 +98,10 @@ class StackNesting {
   ///
   /// Each stack location is allocated by a single allocation instruction.
   struct StackLoc {
-    StackLoc(AllocationInst *Alloc) : Alloc(Alloc) { }
+    StackLoc(SingleValueInstruction *Alloc) : Alloc(Alloc) { }
 
     /// Back-link to the allocation instruction.
-    AllocationInst *Alloc;
+    SingleValueInstruction *Alloc;
 
     /// Bit-set which represents all alive locations at this allocation.
     /// It obviously includes this location itself. And it includes all "outer"

include/swift/Serialization/ModuleFormat.h

@@ -52,7 +52,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 471; // Last change: Add @available(_PackageDescription..)
+const uint16_t SWIFTMODULE_VERSION_MINOR = 472; // Last change: partial_apply [stack]
 
 using DeclIDField = BCFixed<31>;
 

lib/IRGen/GenFunc.cpp

@@ -1254,7 +1254,7 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
 
 /// Emit a partial application thunk for a function pointer applied to a partial
 /// set of argument values.
-void irgen::emitFunctionPartialApplication(
+Optional<StackAddress> irgen::emitFunctionPartialApplication(
     IRGenFunction &IGF, SILFunction &SILFn, const FunctionPointer &fn,
     llvm::Value *fnContext, Explosion &args, ArrayRef<SILParameterInfo> params,
     SubstitutionMap subs, CanSILFunctionType origType,
@@ -1270,8 +1270,6 @@ void irgen::emitFunctionPartialApplication(
   SmallVector<SILType, 4> argValTypes;
   SmallVector<ParameterConvention, 4> argConventions;
 
-  assert(!outType->isNoEscape());
-
   // Reserve space for polymorphic bindings.
   auto bindings = NecessaryBindings::forFunctionInvocations(IGF.IGM,
                                                             origType, subs);
@@ -1398,7 +1396,7 @@ void irgen::emitFunctionPartialApplication(
     llvm::Value *ctx = args.claimNext();
     ctx = IGF.Builder.CreateBitCast(ctx, IGF.IGM.RefCountedPtrTy);
     out.add(ctx);
-    return;
+    return {};
   }
 
   Optional<FunctionPointer> staticFn;
@@ -1437,14 +1435,20 @@ void irgen::emitFunctionPartialApplication(
     llvm::Value *ctx = args.claimNext();
     if (isIndirectFormalParameter(*singleRefcountedConvention))
       ctx = IGF.Builder.CreateLoad(ctx, IGF.IGM.getPointerAlignment());
+
+    auto expectedClosureTy =
+        outType->isNoEscape() ? IGF.IGM.OpaquePtrTy : IGF.IGM.RefCountedPtrTy;
+
     // We might get a struct containing a pointer e.g type <{ %AClass* }>
-    if (ctx->getType() != IGF.IGM.RefCountedPtrTy)
-      ctx = IGF.coerceValue(ctx, IGF.IGM.RefCountedPtrTy, IGF.IGM.DataLayout);
+    if (ctx->getType() != expectedClosureTy)
+      ctx = IGF.coerceValue(ctx, expectedClosureTy, IGF.IGM.DataLayout);
     out.add(ctx);
-    return;
+    if (outType->isNoEscape())
+      return StackAddress();
+    return {};
   }
 
-  // Store the context arguments on the heap.
+  // Store the context arguments on the heap/stack.
   assert(argValTypes.size() == argTypeInfos.size()
          && argTypeInfos.size() == argConventions.size()
          && "argument info lists out of sync");
@@ -1457,13 +1461,26 @@ void irgen::emitFunctionPartialApplication(
                                                        layout);
 
   llvm::Value *data;
+
+  Optional<StackAddress> stackAddr;
+
   if (args.empty() && layout.isKnownEmpty()) {
-    data = IGF.IGM.RefCountedNull;
+    if (outType->isNoEscape())
+      data = llvm::ConstantPointerNull::get(IGF.IGM.OpaquePtrTy);
+    else
+      data = IGF.IGM.RefCountedNull;
   } else {
-    // Allocate a new object.
-    HeapNonFixedOffsets offsets(IGF, layout);
 
-    data = IGF.emitUnmanagedAlloc(layout, "closure", descriptor, &offsets);
+    // Allocate a new object on the heap or stack.
+    HeapNonFixedOffsets offsets(IGF, layout);
+    if (outType->isNoEscape()) {
+      stackAddr = IGF.emitDynamicAlloca(
+          IGF.IGM.Int8Ty, layout.emitSize(IGF.IGM), Alignment(16));
+      stackAddr = stackAddr->withAddress(IGF.Builder.CreateBitCast(
+          stackAddr->getAddress(), IGF.IGM.OpaquePtrTy));
+      data = stackAddr->getAddress().getAddress();
+    } else
+      data = IGF.emitUnmanagedAlloc(layout, "closure", descriptor, &offsets);
     Address dataAddr = layout.emitCastTo(IGF, data);
 
     unsigned i = 0;
@@ -1531,6 +1548,7 @@ void irgen::emitFunctionPartialApplication(
   forwarder = IGF.Builder.CreateBitCast(forwarder, IGF.IGM.Int8PtrTy);
   out.add(forwarder);
   out.add(data);
+  return stackAddr;
 }
 
 /// Emit the block copy helper for a block.

lib/IRGen/GenFunc.h

@@ -47,7 +47,7 @@ namespace irgen {
 
   /// Emit a partial application thunk for a function pointer applied to a
   /// partial set of argument values.
-  void emitFunctionPartialApplication(
+  Optional<StackAddress> emitFunctionPartialApplication(
       IRGenFunction &IGF, SILFunction &SILFn, const FunctionPointer &fnPtr,
       llvm::Value *fnContext, Explosion &args,
       ArrayRef<SILParameterInfo> argTypes, SubstitutionMap subs,

lib/IRGen/GenOpaque.cpp

@@ -555,15 +555,17 @@ StackAddress IRGenFunction::emitDynamicAlloca(llvm::Type *eltTy,
 /// location before the dynamic alloca's call.
 void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address) {
   // In coroutines, unconditionally call llvm.coro.alloca.free.
-  if (isCoroutine()) {
+  // Except if the address is invalid, this happens when this is a StackAddress
+  // for a partial_apply [stack] that did not need a context object on the
+  // stack.
+  if (isCoroutine() && address.getAddress().isValid()) {
     auto allocToken = address.getExtraInfo();
     assert(allocToken && "dynamic alloca in coroutine without alloc token?");
     auto freeFn = llvm::Intrinsic::getDeclaration(
         &IGM.Module, llvm::Intrinsic::ID::coro_alloca_free);
     Builder.CreateCall(freeFn, allocToken);
     return;
   }
-
   // Otherwise, call llvm.stackrestore if an address was saved.
   auto savedSP = address.getExtraInfo();
   if (savedSP == nullptr)
@@ -1444,4 +1446,4 @@ void TypeInfo::storeExtraInhabitantTagDynamic(IRGenFunction &IGF,
   } else {
     storeEnumTagSinglePayload(IGF, tag, tag, address, T, isOutlined);
   }
-}
+}

lib/IRGen/IRGenSIL.cpp

@@ -374,6 +374,7 @@ class IRGenSILFunction :
 {
 public:
   llvm::DenseMap<SILValue, LoweredValue> LoweredValues;
+  llvm::DenseMap<SILValue, StackAddress> LoweredPartialApplyAllocations;
   llvm::DenseMap<SILType, LoweredValue> LoweredUndefs;
 
   /// All alloc_ref instructions which allocate the object on the stack.
@@ -2527,11 +2528,16 @@ void IRGenSILFunction::visitPartialApplyInst(swift::PartialApplyInst *i) {
 
   // Create the thunk and function value.
   Explosion function;
-  emitFunctionPartialApplication(
+  auto closureStackAddr = emitFunctionPartialApplication(
       *this, *CurSILFn, calleeFn, innerContext, llArgs, params,
       i->getSubstitutionMap(), origCalleeTy, i->getSubstCalleeType(),
       i->getType().castTo<SILFunctionType>(), function, false);
   setLoweredExplosion(v, function);
+
+  if (closureStackAddr) {
+    assert(i->isOnStack());
+    LoweredPartialApplyAllocations[v] = *closureStackAddr;
+  }
 }
 
 void IRGenSILFunction::visitIntegerLiteralInst(swift::IntegerLiteralInst *i) {
@@ -4074,6 +4080,13 @@ void IRGenSILFunction::visitAllocRefDynamicInst(swift::AllocRefDynamicInst *i) {
 }
 
 void IRGenSILFunction::visitDeallocStackInst(swift::DeallocStackInst *i) {
+  if (auto *closure = dyn_cast<PartialApplyInst>(i->getOperand())) {
+    assert(closure->isOnStack());
+    auto stackAddr = LoweredPartialApplyAllocations[i->getOperand()];
+    emitDeallocateDynamicAlloca(stackAddr);
+    return;
+  }
+
   auto allocatedType = i->getOperand()->getType();
   const TypeInfo &allocatedTI = getTypeInfo(allocatedType);
   StackAddress stackAddr = getLoweredStackAddress(i->getOperand());