First cut of co-routine differentiation

Author: asl

include/swift/AST/IndexSubset.h

@@ -108,8 +108,10 @@ class IndexSubset : public llvm::FoldingSetNode {
   static IndexSubset *get(ASTContext &ctx, unsigned capacity,
                           ArrayRef<unsigned> indices) {
     SmallBitVector indicesBitVec(capacity, false);
-    for (auto index : indices)
+    for (auto index : indices) {
+      assert(index < capacity);
       indicesBitVec.set(index);
+    }
     return IndexSubset::get(ctx, indicesBitVec);
   }
 

include/swift/AST/Types.h

@@ -5130,8 +5130,11 @@ class SILFunctionType final
   /// Returns the number of function potential semantic results:
   ///  * Usual results
   ///  * Inout parameters
+  ///  * yields
   unsigned getNumAutoDiffSemanticResults() const {
-    return getNumResults() + getNumAutoDiffSemanticResultsParameters();
+    return getNumResults() +
+           getNumAutoDiffSemanticResultsParameters() +
+           getNumIndirectFormalYields();
   }
 
   /// Get the generic signature that the component types are specified

include/swift/SIL/SILFunctionConventions.h

@@ -248,6 +248,14 @@ class SILFunctionConventions {
            idx < indirectResults + getNumIndirectSILErrorResults();
   }
 
+  unsigned getNumAutoDiffSemanticResults() const {
+    return funcTy->getNumAutoDiffSemanticResults();
+  }
+
+  unsigned getNumAutoDiffSemanticResultParameters() const {
+    return funcTy->getNumAutoDiffSemanticResultsParameters();
+  }
+
   /// Are any SIL results passed as address-typed arguments?
   bool hasIndirectSILResults() const { return getNumIndirectSILResults() != 0; }
   bool hasIndirectSILErrorResults() const { return getNumIndirectSILErrorResults() != 0; }

include/swift/SILOptimizer/Differentiation/ADContext.h

@@ -17,6 +17,7 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_ADCONTEXT_H
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_ADCONTEXT_H
 
+#include "swift/SIL/ApplySite.h"
 #include "swift/SILOptimizer/Differentiation/Common.h"
 #include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
 
@@ -51,6 +52,12 @@ struct NestedApplyInfo {
   /// The original pullback type before reabstraction. `None` if the pullback
   /// type is not reabstracted.
   llvm::Optional<CanSILFunctionType> originalPullbackType;
+  /// Index of `apply` pullback in nested pullback call
+  unsigned pullbackIdx = -1U;
+  /// Pullback value itself that is memoized in some cases (e.g. pullback is
+  /// called by `begin_apply`, but should be destroyed after `end_apply`).
+  SILValue pullback = SILValue();
+  SILValue beginApplyToken = SILValue();
 };
 
 /// Per-module contextual information for the Differentiation pass.
@@ -97,7 +104,7 @@ class ADContext {
 
   /// Mapping from original `apply` instructions to their corresponding
   /// `NestedApplyInfo`s.
-  llvm::DenseMap<ApplyInst *, NestedApplyInfo> nestedApplyInfo;
+  llvm::DenseMap<FullApplySite, NestedApplyInfo> nestedApplyInfo;
 
   /// List of generated functions (JVPs, VJPs, pullbacks, and thunks).
   /// Saved for deletion during cleanup.
@@ -185,7 +192,7 @@ class ADContext {
     invokers.insert({witness, DifferentiationInvoker(witness)});
   }
 
-  llvm::DenseMap<ApplyInst *, NestedApplyInfo> &getNestedApplyInfo() {
+  llvm::DenseMap<FullApplySite, NestedApplyInfo> &getNestedApplyInfo() {
     return nestedApplyInfo;
   }
 

include/swift/SILOptimizer/Differentiation/Common.h

@@ -20,6 +20,7 @@
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/SemanticAttrs.h"
+#include "swift/SIL/ApplySite.h"
 #include "swift/SIL/SILDifferentiabilityWitness.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/Projection.h"
@@ -112,15 +113,15 @@ void collectAllDirectResultsInTypeOrder(SILFunction &function,
 /// Given a function call site, gathers all of its actual results (both direct
 /// and indirect) in an order defined by its result type.
 void collectAllActualResultsInTypeOrder(
-    ApplyInst *ai, ArrayRef<SILValue> extractedDirectResults,
+    FullApplySite fai, ArrayRef<SILValue> extractedDirectResults,
     SmallVectorImpl<SILValue> &results);
 
 /// For an `apply` instruction with active results, compute:
 /// - The results of the `apply` instruction, in type order.
 /// - The set of minimal parameter and result indices for differentiating the
 ///   `apply` instruction.
 void collectMinimalIndicesForFunctionCall(
-    ApplyInst *ai, const AutoDiffConfig &parentConfig,
+    FullApplySite fai, const AutoDiffConfig &parentConfig,
     const DifferentiableActivityInfo &activityInfo,
     SmallVectorImpl<SILValue> &results, SmallVectorImpl<unsigned> &paramIndices,
     SmallVectorImpl<unsigned> &resultIndices);

include/swift/SILOptimizer/Differentiation/LinearMapInfo.h

@@ -77,9 +77,9 @@ class LinearMapInfo {
   /// For differentials: these are successor enums.
   llvm::DenseMap<SILBasicBlock *, EnumDecl *> branchingTraceDecls;
 
-  /// Mapping from `apply` instructions in the original function to the
+  /// Mapping from `apply` / `begin_apply` instructions in the original function to the
   /// corresponding linear map tuple type index.
-  llvm::DenseMap<ApplyInst *, unsigned> linearMapIndexMap;
+  llvm::DenseMap<FullApplySite, unsigned> linearMapIndexMap;
 
   /// Mapping from predecessor-successor basic block pairs in the original
   /// function to the corresponding branching trace enum case.
@@ -112,9 +112,9 @@ class LinearMapInfo {
   void populateBranchingTraceDecl(SILBasicBlock *originalBB,
                                   SILLoopInfo *loopInfo);
 
-  /// Given an `apply` instruction, conditionally gets a linear map tuple field
-  /// AST type for its linear map function if it is active.
-  Type getLinearMapType(ADContext &context, ApplyInst *ai);
+  /// Given an `apply` / `begin_apply` instruction, conditionally gets a linear
+  /// map tuple field AST type for its linear map function if it is active.
+  Type getLinearMapType(ADContext &context, FullApplySite fai);
 
   /// Generates linear map struct and branching enum declarations for the given
   /// function. Linear map structs are populated with linear map fields and a
@@ -180,18 +180,18 @@ class LinearMapInfo {
   }
 
   /// Finds the linear map index in the pullback tuple for the given
-  /// `apply` instruction in the original function.
-  unsigned lookUpLinearMapIndex(ApplyInst *ai) const {
-    assert(ai->getFunction() == original);
-    auto lookup = linearMapIndexMap.find(ai);
+  /// `apply` / `begin_apply` instruction in the original function.
+  unsigned lookUpLinearMapIndex(FullApplySite fas) const {
+    assert(fas->getFunction() == original);
+    auto lookup = linearMapIndexMap.find(fas);
     assert(lookup != linearMapIndexMap.end() &&
            "No linear map field corresponding to the given `apply`");
     return lookup->getSecond();
   }
 
-  Type lookUpLinearMapType(ApplyInst *ai) const {
-    unsigned idx = lookUpLinearMapIndex(ai);
-    return getLinearMapTupleType(ai->getParentBlock())->getElement(idx).getType();
+  Type lookUpLinearMapType(FullApplySite fas) const {
+    unsigned idx = lookUpLinearMapIndex(fas);
+    return getLinearMapTupleType(fas->getParent())->getElement(idx).getType();
   }
 
   bool hasHeapAllocatedContext() const {

include/swift/SILOptimizer/Differentiation/Thunk.h

@@ -56,6 +56,11 @@ SILFunction *getOrCreateReabstractionThunk(SILOptFunctionBuilder &fb,
                                            CanSILFunctionType fromType,
                                            CanSILFunctionType toType);
 
+SILValue reabstractCoroutine(
+    SILBuilder &builder, SILOptFunctionBuilder &fb, SILLocation loc,
+    SILValue fn, CanSILFunctionType toType,
+    std::function<SubstitutionMap(SubstitutionMap)> remapSubstitutions);
+
 /// Reabstracts the given function-typed value `fn` to the target type `toType`.
 /// Remaps substitutions using `remapSubstitutions`.
 SILValue reabstractFunction(

lib/AST/AutoDiff.cpp

@@ -101,6 +101,10 @@ DifferentiabilityWitnessFunctionKind::getAsDerivativeFunctionKind() const {
   llvm_unreachable("invalid derivative kind");
 }
 
+void AutoDiffConfig::dump() const {
+  print(llvm::errs());
+}
+
 void AutoDiffConfig::print(llvm::raw_ostream &s) const {
   s << "(parameters=";
   parameterIndices->print(s);

lib/SIL/IR/SILFunctionType.cpp

@@ -258,6 +258,13 @@ IndexSubset *SILFunctionType::getDifferentiabilityResultIndices() {
 
   numSemanticResults += getNumAutoDiffSemanticResultsParameters();
 
+  // Check yields.
+  for (auto yieldAndIndex : enumerate(getYields()))
+    if (!yieldAndIndex.value().hasOption(
+          SILParameterInfo::NotDifferentiable))
+      resultIndices.push_back(numSemanticResults + yieldAndIndex.index());
+
+  numSemanticResults += getNumIndirectFormalYields();
   return IndexSubset::get(getASTContext(), numSemanticResults, resultIndices);
 }
 
@@ -555,10 +562,14 @@ static CanSILFunctionType getAutoDiffDifferentialType(
         param.getConvention());
     differentialParams.push_back({paramTanType, paramConv});
   }
+
   SmallVector<SILResultInfo, 1> differentialResults;
+  unsigned firstSemanticParamResultIdx = originalFnTy->getNumResults();
+  unsigned firstYieldResultIndex = originalFnTy->getNumResults() +
+      originalFnTy->getNumAutoDiffSemanticResultsParameters();
   for (auto resultIndex : resultIndices->getIndices()) {
     // Handle formal original result.
-    if (resultIndex < originalFnTy->getNumResults()) {
+    if (resultIndex < firstSemanticParamResultIdx) {
       auto &result = originalResults[resultIndex];
       auto resultTanType = getAutoDiffTangentTypeForLinearMap(
           result.getInterfaceType(), lookupConformance,
@@ -571,26 +582,38 @@ static CanSILFunctionType getAutoDiffDifferentialType(
           result.getConvention());
       differentialResults.push_back({resultTanType, resultConv});
       continue;
-    }
-    // Handle original semantic result parameters.
-    auto resultParamIndex = resultIndex - originalFnTy->getNumResults();
-    auto resultParamIt = std::next(
+    } else if (resultIndex < firstYieldResultIndex) {
+      // Handle original semantic result parameters.
+      auto resultParamIndex = resultIndex - originalFnTy->getNumResults();
+      auto resultParamIt = std::next(
         originalFnTy->getAutoDiffSemanticResultsParameters().begin(),
         resultParamIndex);
-    auto paramIndex =
-      std::distance(originalFnTy->getParameters().begin(), &*resultParamIt);
-    // If the original semantic result parameter is a differentiability
-    // parameter, then it already has a corresponding differential
-    // parameter. Skip adding a corresponding differential result.
-    if (parameterIndices->contains(paramIndex))
-      continue;
+      auto paramIndex =
+        std::distance(originalFnTy->getParameters().begin(), &*resultParamIt);
+      // If the original semantic result parameter is a differentiability
+      // parameter, then it already has a corresponding differential
+      // parameter. Skip adding a corresponding differential result.
+      if (parameterIndices->contains(paramIndex))
+        continue;
 
-    auto resultParam = originalFnTy->getParameters()[paramIndex];
-    auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
-      resultParam.getInterfaceType(), lookupConformance,
+      auto resultParam = originalFnTy->getParameters()[paramIndex];
+      auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
+        resultParam.getInterfaceType(), lookupConformance,
         substGenericParams, substReplacements, ctx);
-    differentialResults.emplace_back(resultParamTanType,
-                                     ResultConvention::Indirect);
+      differentialResults.emplace_back(resultParamTanType,
+                                       ResultConvention::Indirect);
+    } else {
+      assert(originalFnTy->isCoroutine());
+      assert(originalFnTy->getCoroutineKind() == SILCoroutineKind::YieldOnce);
+      auto yieldResultIndex = resultIndex - firstYieldResultIndex;
+      auto yieldResult = originalFnTy->getYields()[yieldResultIndex];
+      auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
+        yieldResult.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
+      ParameterConvention paramTanConvention = yieldResult.getConvention();
+      assert(yieldResult.getConvention() == ParameterConvention::Indirect_Inout);
+      differentialParams.emplace_back(resultParamTanType, paramTanConvention);
+    }
   }
 
   SubstitutionMap substitutions;
@@ -604,9 +627,9 @@ static CanSILFunctionType getAutoDiffDifferentialType(
   }
   return SILFunctionType::get(
       GenericSignature(), SILFunctionType::ExtInfo(), SILCoroutineKind::None,
-      ParameterConvention::Direct_Guaranteed, differentialParams, {},
-      differentialResults, llvm::None, substitutions,
-      /*invocationSubstitutions*/ SubstitutionMap(), ctx);
+      ParameterConvention::Direct_Guaranteed,
+      differentialParams, {}, differentialResults, llvm::None,
+      substitutions, /*invocationSubstitutions*/ SubstitutionMap(), ctx);
 }
 
 /// Returns the pullback type for the given original function type, parameter
@@ -696,11 +719,15 @@ static CanSILFunctionType getAutoDiffPullbackType(
     return conv;
   };
 
-  // Collect pullback parameters.
+  // Collect pullback parameters & yields
   SmallVector<SILParameterInfo, 1> pullbackParams;
+  SmallVector<SILYieldInfo, 1> pullbackYields;
+  unsigned firstSemanticParamResultIdx = originalFnTy->getNumResults();
+  unsigned firstYieldResultIndex = originalFnTy->getNumResults() +
+      originalFnTy->getNumAutoDiffSemanticResultsParameters();
   for (auto resultIndex : resultIndices->getIndices()) {
     // Handle formal original result.
-    if (resultIndex < originalFnTy->getNumResults()) {
+    if (resultIndex < firstSemanticParamResultIdx) {
       auto &origRes = originalResults[resultIndex];
       auto resultTanType = getAutoDiffTangentTypeForLinearMap(
           origRes.getInterfaceType(), lookupConformance,
@@ -712,28 +739,38 @@ static CanSILFunctionType getAutoDiffPullbackType(
               ->getCanonicalType(),
           origRes.getConvention());
       pullbackParams.emplace_back(resultTanType, paramConv);
-      continue;
+    } else if (resultIndex < firstYieldResultIndex) {
+      // Handle original semantic result parameters.
+      auto resultParamIndex = resultIndex - firstSemanticParamResultIdx;
+      auto resultParamIt = std::next(
+        originalFnTy->getAutoDiffSemanticResultsParameters().begin(),
+        resultParamIndex);
+      auto paramIndex =
+        std::distance(originalFnTy->getParameters().begin(), &*resultParamIt);
+      auto resultParam = originalFnTy->getParameters()[paramIndex];
+      // The pullback parameter convention depends on whether the original `inout`
+      // parameter is a differentiability parameter.
+      // - If yes, the pullback parameter convention is `@inout`.
+      // - If no, the pullback parameter convention is `@in_guaranteed`.
+      auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
+        resultParam.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
+      ParameterConvention paramTanConvention = resultParam.getConvention();
+      if (!parameterIndices->contains(paramIndex))
+        paramTanConvention = ParameterConvention::Indirect_In_Guaranteed;
+      pullbackParams.emplace_back(resultParamTanType, paramTanConvention);
+    } else {
+      assert(originalFnTy->isCoroutine());
+      assert(originalFnTy->getCoroutineKind() == SILCoroutineKind::YieldOnce);
+      auto yieldResultIndex = resultIndex - firstYieldResultIndex;
+      auto yieldResult = originalFnTy->getYields()[yieldResultIndex];
+      auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
+        yieldResult.getInterfaceType(), lookupConformance,
+        substGenericParams, substReplacements, ctx);
+      ParameterConvention paramTanConvention = yieldResult.getConvention();
+      assert(yieldResult.getConvention() == ParameterConvention::Indirect_Inout);
+      pullbackYields.emplace_back(resultParamTanType, paramTanConvention);
     }
-    // Handle original semantic result parameters.
-    auto resultParamIndex = resultIndex - originalFnTy->getNumResults();
-    auto resultParamIt = std::next(
-      originalFnTy->getAutoDiffSemanticResultsParameters().begin(),
-      resultParamIndex);
-    auto paramIndex =
-      std::distance(originalFnTy->getParameters().begin(), &*resultParamIt);
-    auto resultParam = originalFnTy->getParameters()[paramIndex];
-    // The pullback parameter convention depends on whether the original `inout`
-    // parameter is a differentiability parameter.
-    // - If yes, the pullback parameter convention is `@inout`.
-    // - If no, the pullback parameter convention is `@in_guaranteed`.
-    auto resultParamTanType = getAutoDiffTangentTypeForLinearMap(
-      resultParam.getInterfaceType(), lookupConformance,
-      substGenericParams, substReplacements, ctx);
-    ParameterConvention paramTanConvention = resultParam.getConvention();
-    if (!parameterIndices->contains(paramIndex))
-      paramTanConvention = ParameterConvention::Indirect_In_Guaranteed;
-
-    pullbackParams.emplace_back(resultParamTanType, paramTanConvention);
   }
 
   // Collect pullback results.
@@ -756,6 +793,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
         param.getConvention());
     pullbackResults.push_back({paramTanType, resultTanConvention});
   }
+
   SubstitutionMap substitutions;
   if (!substGenericParams.empty()) {
     auto genericSig =
@@ -768,7 +806,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
   return SILFunctionType::get(
       GenericSignature(), SILFunctionType::ExtInfo(), originalFnTy->getCoroutineKind(),
       ParameterConvention::Direct_Guaranteed,
-      pullbackParams, {}, pullbackResults, llvm::None, substitutions,
+      pullbackParams, pullbackYields, pullbackResults, llvm::None, substitutions,
       /*invocationSubstitutions*/ SubstitutionMap(), ctx);
 }
 
@@ -920,10 +958,9 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
   // Compute the derivative function results.
   SmallVector<SILResultInfo, 4> newResults;
   newResults.reserve(getNumResults() + 1);
-  for (auto &result : constrainedOriginalFnTy->getResults()) {
+  for (auto &result : constrainedOriginalFnTy->getResults())
     newResults.push_back(result);
-  }
-  newResults.push_back({closureType, ResultConvention::Owned});
+  newResults.emplace_back(closureType, ResultConvention::Owned);
 
   // Compute the derivative function ExtInfo.
   // If original function is `@convention(c)`, the derivative function should