[AutoDiff] First cut of coroutines differentiation (#71461)

This PR implements first set of changes required to support autodiff for coroutines. It mostly targeted to `_modify` accessors in standard library (and beyond), but overall implementation is quite generic.

There are some specifics of implementation and known limitations:
 - Only `@yield_once` coroutines are naturally supported
 - VJP is a coroutine itself: it yields the results *and* returns a pullback closure as a normal return. This allows us to capture values produced in resume part of a coroutine (this is required for defers and other cleanups / commits)
 - Pullback is a coroutine, we assume that coroutine cannot abort and therefore we execute the original coroutine in reverse from return via yield and then back to the entry
 - It seems there is no semantically sane way to support `_read` coroutines (as we will need to "accept" adjoints via yields), therefore only coroutines with inout yields are supported (`_modify` accessors). Pullbacks of such coroutines take adjoint buffer as input argument, yield this buffer (to accumulate adjoint values in the caller) and finally return the adjoints indirectly.
 - Coroutines (as opposed to normal functions) are not first-class values: there is no AST type for them, one cannot e.g. store them into tuples, etc. So, everywhere where AST type is required, we have to hack around.
 - As there is no AST type for coroutines, there is no way one could register custom derivative for coroutines. So far only compiler-produced derivatives are supported
 - There are lots of common things wrt normal function apply's, but still there are subtle but important differences. I tried to organize the code to enable code reuse, still it was not always possible, so some code duplication could be seen
 - The order of how pullback closures are produced in VJP is a bit different: for normal apply's VJP produces both value and pullback closure via a single nested VJP apply. This is not so anymore with coroutine VJP's: yielded values are produced at `begin_apply` site and pullback closure is available only from `end_apply`, so we need to track the order in which pullbacks are produced (and arrange consumption of the values accordingly – effectively delay them)
 - On the way some complementary changes were required in e.g. mangler / demangler

This patch covers the generation of derivatives up to SIL level, however, it is not enough as codegen of `partial_apply` of a coroutine is completely broken. The fix for this will be submitted separately as it is not directly autodiff-related.

---------

Co-authored-by: Andrew Savonichev <[email protected]>
Co-authored-by: Richard Wei <[email protected]>

Author: 3 people

include/swift/AST/ASTDemangler.h

@@ -150,7 +150,9 @@ class ASTBuilder {
 
   Type createImplFunctionType(
       Demangle::ImplParameterConvention calleeConvention,
+      Demangle::ImplCoroutineKind coroutineKind,
       ArrayRef<Demangle::ImplFunctionParam<Type>> params,
+      ArrayRef<Demangle::ImplFunctionYield<Type>> yields,
       ArrayRef<Demangle::ImplFunctionResult<Type>> results,
       std::optional<Demangle::ImplFunctionResult<Type>> errorResult,
       ImplFunctionTypeFlags flags);

include/swift/AST/DiagnosticsSIL.def

@@ -607,6 +607,8 @@ NOTE(autodiff_cannot_differentiate_through_multiple_results,none,
      "cannot differentiate through multiple results", ())
 NOTE(autodiff_cannot_differentiate_through_inout_arguments,none,
      "cannot differentiate through 'inout' arguments", ())
+NOTE(autodiff_cannot_differentiate_through_direct_yield,none,
+     "cannot differentiate through a '_read' accessor", ())
 NOTE(autodiff_enums_unsupported,none,
      "differentiating enum values is not yet supported", ())
 NOTE(autodiff_stored_property_parent_not_differentiable,none,

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

@@ -5174,8 +5174,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() +
+           getNumYields();
   }
 
   /// Get the generic signature that the component types are specified

include/swift/Demangling/Demangle.h

@@ -553,6 +553,7 @@ struct [[nodiscard]] ManglingError {
     UnknownEncoding,
     InvalidImplCalleeConvention,
     InvalidImplDifferentiability,
+    InvalidImplCoroutineKind,
     InvalidImplFunctionAttribute,
     InvalidImplParameterConvention,
     InvalidImplParameterTransferring,

include/swift/Demangling/DemangleNodes.def

@@ -139,6 +139,7 @@ NODE(ImplFunctionAttribute)
 NODE(ImplFunctionConvention)
 NODE(ImplFunctionConventionName)
 NODE(ImplFunctionType)
+NODE(ImplCoroutineKind)
 NODE(ImplInvocationSubstitutions)
 CONTEXT_NODE(ImplicitClosure)
 NODE(ImplParameter)

include/swift/Demangling/TypeDecoder.h

@@ -48,6 +48,12 @@ enum class ImplMetatypeRepresentation {
   ObjC,
 };
 
+enum class ImplCoroutineKind {
+  None,
+  YieldOnce,
+  YieldMany,
+};
+
 /// Describe a function parameter, parameterized on the type
 /// representation.
 template <typename BuiltType>
@@ -188,6 +194,9 @@ class ImplFunctionParam {
   BuiltType getType() const { return Type; }
 };
 
+template<typename Type>
+using ImplFunctionYield = ImplFunctionParam<Type>;
+
 enum class ImplResultConvention {
   Indirect,
   Owned,
@@ -1023,9 +1032,11 @@ class TypeDecoder {
     case NodeKind::ImplFunctionType: {
       auto calleeConvention = ImplParameterConvention::Direct_Unowned;
       llvm::SmallVector<ImplFunctionParam<BuiltType>, 8> parameters;
+      llvm::SmallVector<ImplFunctionYield<BuiltType>, 8> yields;
       llvm::SmallVector<ImplFunctionResult<BuiltType>, 8> results;
       llvm::SmallVector<ImplFunctionResult<BuiltType>, 8> errorResults;
       ImplFunctionTypeFlags flags;
+      ImplCoroutineKind coroutineKind = ImplCoroutineKind::None;
 
       for (unsigned i = 0; i < Node->getNumChildren(); i++) {
         auto child = Node->getChild(i);
@@ -1066,6 +1077,15 @@ class TypeDecoder {
           } else if (child->getText() == "@async") {
             flags = flags.withAsync();
           }
+        } else if (child->getKind() == NodeKind::ImplCoroutineKind) {
+          if (!child->hasText())
+            return MAKE_NODE_TYPE_ERROR0(child, "expected text");
+          if (child->getText() == "yield_once") {
+            coroutineKind = ImplCoroutineKind::YieldOnce;
+          } else if (child->getText() == "yield_many") {
+            coroutineKind = ImplCoroutineKind::YieldMany;
+          } else
+            return MAKE_NODE_TYPE_ERROR0(child, "failed to decode coroutine kind");
         } else if (child->getKind() == NodeKind::ImplDifferentiabilityKind) {
           ImplFunctionDifferentiabilityKind implDiffKind;
           switch ((MangledDifferentiabilityKind)child->getIndex()) {
@@ -1088,10 +1108,14 @@ class TypeDecoder {
           if (decodeImplFunctionParam(child, depth + 1, parameters))
             return MAKE_NODE_TYPE_ERROR0(child,
                                          "failed to decode function parameter");
+        } else if (child->getKind() == NodeKind::ImplYield) {
+          if (decodeImplFunctionParam(child, depth + 1, yields))
+            return MAKE_NODE_TYPE_ERROR0(child,
+                                         "failed to decode function yields");
         } else if (child->getKind() == NodeKind::ImplResult) {
           if (decodeImplFunctionParam(child, depth + 1, results))
             return MAKE_NODE_TYPE_ERROR0(child,
-                                         "failed to decode function parameter");
+                                         "failed to decode function results");
         } else if (child->getKind() == NodeKind::ImplErrorResult) {
           if (decodeImplFunctionPart(child, depth + 1, errorResults))
             return MAKE_NODE_TYPE_ERROR0(child,
@@ -1115,11 +1139,10 @@ class TypeDecoder {
 
       // TODO: Some cases not handled above, but *probably* they cannot
       // appear as the types of values in SIL (yet?):
-      // - functions with yield returns
       // - functions with generic signatures
       // - foreign error conventions
-      return Builder.createImplFunctionType(calleeConvention,
-                                            parameters, results,
+      return Builder.createImplFunctionType(calleeConvention, coroutineKind,
+                                            parameters, yields, results,
                                             errorResult, flags);
     }
 

include/swift/RemoteInspection/TypeRefBuilder.h

@@ -1134,7 +1134,9 @@ class TypeRefBuilder {
 
   const FunctionTypeRef *createImplFunctionType(
       Demangle::ImplParameterConvention calleeConvention,
+      Demangle::ImplCoroutineKind coroutineKind,
       llvm::ArrayRef<Demangle::ImplFunctionParam<const TypeRef *>> params,
+      llvm::ArrayRef<Demangle::ImplFunctionYield<const TypeRef *>> yields,
       llvm::ArrayRef<Demangle::ImplFunctionResult<const TypeRef *>> results,
       std::optional<Demangle::ImplFunctionResult<const TypeRef *>> errorResult,
       ImplFunctionTypeFlags flags) {

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.
   std::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;
   }